diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 5e4583a8..07c29e0c 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -690,12 +690,6 @@ func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler message.S return } - // Check if context is cancelled - if so, skip handler call to avoid blocking - if ctx.Context != nil && ctx.Context.Err() != nil { - ctx.Logger.Debug("Context cancelled, skipping sendAgentStreamEnd handler call") - return - } - endData := &message.EventStreamEndData{ RequestID: ctx.RequestID(), ContextID: ctx.ID, @@ -727,25 +721,17 @@ func (ast *Assistant) sendStreamEndOnError(ctx *context.Context, handler message // handleInterrupt handles the interrupt signal // This is called by the interrupt listener when a signal is received func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.InterruptSignal) error { - // Handle based on interrupt type switch signal.Type { case context.InterruptForce: - // Force interrupt: context is already cancelled in handleSignal - // LLM streaming will detect ctx.Interrupt.Context().Done() and stop - ctx.Logger.Debug("Force interrupt: stopping current operations immediately") - + ctx.Logger.Debug("Force interrupt received") + if ctx.Buffer != nil { + ctx.Buffer.FailCurrentStep(context.ResumeStatusInterrupted, + fmt.Errorf("interrupted by user")) + } case context.InterruptGraceful: - ctx.Logger.Debug("Graceful interrupt: will process after current step completes") - // Graceful interrupt: let current operation complete - // The signal is stored in current/pending, can be checked at checkpoints + ctx.Logger.Debug("Graceful interrupt received: messages=%d", len(signal.Messages)) } - // TODO: Implement actual interrupt handling logic: - // 1. For graceful: wait for current step, then merge messages and restart - // 2. For force: immediately stop and restart with new messages - // 3. Call Interrupted Hook if configured - // 4. Decide whether to continue, restart, or abort based on Hook response - return nil } diff --git a/agent/assistant/handlers/stream.go b/agent/assistant/handlers/stream.go index 1d0d1b6b..e62c252e 100644 --- a/agent/assistant/handlers/stream.go +++ b/agent/assistant/handlers/stream.go @@ -46,6 +46,9 @@ func DefaultStreamHandler(ctx *context.Context) message.StreamFunc { case message.ChunkToolCall: return state.handleToolCall(data) + case message.ChunkExecute: + return state.handleExecute(data) + case message.ChunkMetadata: return state.handleMetadata(data) @@ -72,9 +75,11 @@ type streamState struct { currentGroupID string // Current group ID (shared by all chunks in the group) currentType string // Track the current message type (text, thinking, tool_call) buffer []byte - chunkCount int // Track number of chunks in current group - messageSeq int // Message sequence number (for generating readable IDs) - groupStartTime time.Time // Track when group started + chunkCount int // Track number of chunks in current group + messageSeq int // Message sequence number (for generating readable IDs) + groupStartTime time.Time // Track when group started + lastExecStatus string // Last observed execute status in current group ("running", "completed", "error") + lastExecProps map[string]interface{} // Accumulated execute props for the current group (merged across chunks) } // handleStreamStart handles stream start event @@ -290,11 +295,75 @@ func (s *streamState) handleToolCall(data []byte) int { return 0 // Continue } -// handleMetadata handles metadata chunks (usage, finish_reason, etc.) +// handleExecute handles execute observation chunks from sandbox CLI agents. +// These represent tool actions observed inside the agent runtime (e.g., Bash, Read, Write). +func (s *streamState) handleExecute(data []byte) int { + if len(data) == 0 { + return 0 + } + + s.currentType = message.TypeExecute + s.buffer = append(s.buffer, data...) + s.chunkCount++ + s.messageSeq++ + + var props map[string]interface{} + if err := jsoniter.Unmarshal(data, &props); err != nil { + return 0 + } + + if st, ok := props["status"].(string); ok { + s.lastExecStatus = st + } + + if s.lastExecProps == nil { + s.lastExecProps = make(map[string]interface{}) + } + for k, v := range props { + s.lastExecProps[k] = v + } + + deltaAction := "merge" + + msg := &message.Message{ + ChunkID: s.ctx.IDGenerator.GenerateChunkID(), + MessageID: s.currentGroupID, + Type: message.TypeExecute, + Delta: true, + DeltaAction: deltaAction, + Props: props, + } + + if err := s.ctx.Send(msg); err != nil { + return 0 + } + + return 0 +} + +// handleMetadata handles metadata chunks (usage, finish_reason, result_summary, etc.) +// For sandbox CLI agents, this carries token usage and result summaries. func (s *streamState) handleMetadata(data []byte) int { - // Metadata is usually not displayed to users - // Could be logged or stored for analytics - return 0 // Continue + if len(data) == 0 { + return 0 + } + + var meta map[string]interface{} + if err := jsoniter.Unmarshal(data, &meta); err != nil { + return 0 + } + + if usage, ok := meta["usage"]; ok { + msg := output.NewEventMessage("token/usage", "", usage) + s.ctx.Send(msg) + } + + if summary, ok := meta["result_summary"]; ok { + msg := output.NewEventMessage("result/summary", "", summary) + s.ctx.Send(msg) + } + + return 0 } // handleError handles error chunks @@ -341,16 +410,20 @@ func (s *streamState) handleMessageEnd(data []byte) int { shouldSkipHistory := s.ctx.Stack != nil && s.ctx.Stack.Options != nil && s.ctx.Stack.Options.Skip != nil && s.ctx.Stack.Options.Skip.History - if s.ctx.Buffer != nil && len(s.buffer) > 0 && !shouldSkipHistory { + // Execute messages have two phases sharing the same message_id: + // 1. running — streamed for UI display only, NOT persisted + // 2. completed / error — the final state, persisted to the buffer + isExecuteRunning := msgType == message.TypeExecute && s.lastExecStatus == "running" + + if s.ctx.Buffer != nil && len(s.buffer) > 0 && !shouldSkipHistory && !isExecuteRunning { assistantID := "" if s.ctx.Stack != nil { assistantID = s.ctx.Stack.AssistantID } - // Build props based on message type var props map[string]interface{} - if msgType == message.TypeToolCall { - // For tool calls, try to parse the accumulated buffer as JSON + switch msgType { + case message.TypeToolCall: var toolCallData interface{} if err := jsoniter.Unmarshal(s.buffer, &toolCallData); err == nil { props = map[string]interface{}{ @@ -361,15 +434,25 @@ func (s *streamState) handleMessageEnd(data []byte) int { "content": string(s.buffer), } } - } else { - // For text/thinking, content is the accumulated text + case message.TypeExecute: + if s.lastExecProps != nil { + props = make(map[string]interface{}, len(s.lastExecProps)) + for k, v := range s.lastExecProps { + props[k] = v + } + } else { + props = map[string]interface{}{ + "content": string(s.buffer), + } + } + default: props = map[string]interface{}{ "content": string(s.buffer), } } s.ctx.Buffer.AddAssistantMessage( - s.currentGroupID, // Use the message ID + s.currentGroupID, msgType, props, blockID, @@ -403,6 +486,8 @@ func (s *streamState) handleMessageEnd(data []byte) int { s.currentType = "" s.buffer = []byte{} s.chunkCount = 0 + s.lastExecStatus = "" + s.lastExecProps = nil return 0 // Continue } diff --git a/agent/assistant/sandbox_v2.go b/agent/assistant/sandbox_v2.go index 1cab3a79..2ad61b64 100644 --- a/agent/assistant/sandbox_v2.go +++ b/agent/assistant/sandbox_v2.go @@ -1,10 +1,12 @@ package assistant import ( + stdContext "context" "fmt" "log" "os" "path/filepath" + "time" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -140,8 +142,10 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options) ctx.SetComputer(computer) cleanup := func() { - // Defensive fallback — executeSandboxV2Stream defer handles the - // normal case; this covers paths that never reach execution. + cleanCtx, cancel := stdContext.WithTimeout(stdContext.Background(), 5*time.Second) + defer cancel() + runner.Cleanup(cleanCtx, computer) + sandboxv2.LifecycleAction(cleanCtx, cfg, computer, manager) } return runner, computer, cleanup, loadingMsgID, nil @@ -194,6 +198,7 @@ func (ast *Assistant) executeSandboxV2Stream( SystemPrompt: systemPrompt, ChatID: ctx.ChatID, Token: tok, + Logger: ctx.Logger, } execReq := &sandboxv2.ExecuteRequest{ diff --git a/agent/context/log.go b/agent/context/log.go index af3ed978..eee34ab2 100644 --- a/agent/context/log.go +++ b/agent/context/log.go @@ -100,6 +100,11 @@ func WithParentID(parentID string) LoggerOption { // noopLogger is a shared no-op logger instance var noopLogger = &RequestLogger{noop: true} +// NoopLogger returns a shared no-op RequestLogger that silently discards all +// log calls. Use when a non-nil logger is required but no actual logging is +// desired (e.g., fallback when StreamRequest.Logger is nil). +func NoopLogger() *RequestLogger { return noopLogger } + // NewRequestLogger creates a new request-scoped logger with async processing func NewRequestLogger(assistantID, chatID, requestID string, opts ...LoggerOption) *RequestLogger { l := &RequestLogger{ diff --git a/agent/output/builtin.go b/agent/output/builtin.go index de62a233..85d88750 100644 --- a/agent/output/builtin.go +++ b/agent/output/builtin.go @@ -66,6 +66,27 @@ func NewToolCallMessage(id, name, arguments string) *message.Message { } } +// NewExecuteMessage creates an execute observation message for sandbox CLI agent actions. +// +// tool: the tool name (e.g., "Bash", "Read", "Write") +// toolID: the agent-side tool call ID +// input: the tool input (structured, may be nil for status-only updates) +// status: "running" | "completed" | "error" +func NewExecuteMessage(tool, toolID string, input interface{}, status string) *message.Message { + props := map[string]interface{}{ + "tool": tool, + "tool_id": toolID, + "status": status, + } + if input != nil { + props["input"] = input + } + return &message.Message{ + Type: message.TypeExecute, + Props: props, + } +} + // NewErrorMessage creates an error message func NewErrorMessage(msg, code string) *message.Message { return &message.Message{ @@ -135,7 +156,7 @@ func NewVideoMessage(url string) *message.Message { // IsBuiltinType checks if a message type is a built-in type func IsBuiltinType(msgType string) bool { switch msgType { - case message.TypeUserInput, message.TypeText, message.TypeThinking, message.TypeLoading, message.TypeToolCall, message.TypeError, message.TypeImage, message.TypeAudio, message.TypeVideo, message.TypeAction, message.TypeEvent: + case message.TypeUserInput, message.TypeText, message.TypeThinking, message.TypeLoading, message.TypeToolCall, message.TypeExecute, message.TypeError, message.TypeImage, message.TypeAudio, message.TypeVideo, message.TypeAction, message.TypeEvent: return true default: return false diff --git a/agent/output/message/types.go b/agent/output/message/types.go index 8f92c273..d3a81834 100644 --- a/agent/output/message/types.go +++ b/agent/output/message/types.go @@ -75,6 +75,9 @@ const ( TypeAudio = "audio" // Audio content TypeVideo = "video" // Video content + // Agent execution observation types + TypeExecute = "execute" // Agent tool execution observation (sandbox CLI agent actions, not LLM tool_call requests) + // System types (not visible in standard chat clients) TypeAction = "action" // System action (open panel, navigate, etc.) - silent in OpenAI clients TypeEvent = "event" // Lifecycle event (stream_start, stream_end, etc.) - CUI only, silent in OpenAI clients @@ -205,6 +208,34 @@ type VideoProps struct { Loop bool `json:"loop,omitempty"` // Whether to loop } +// ExecuteProps defines the standard structure for execute messages. +// Type: "execute" +// +// Represents an autonomous action taken by an external Agent (e.g., Claude CLI +// in a sandbox container, Codex CLI, or any future CLI-based Agent). Unlike +// tool_call (which is a call request that Yao dispatches), execute is an +// observation of an action that already happened inside the Agent's runtime. +// +// Lifecycle (via delta merge on the same MessageID): +// +// 1. message_start {type: "execute"} +// 2. ChunkExecute {tool, tool_id, input, status:"running"} — merge +// 3. ChunkExecute {tool_id, output, status:"completed"} — merge +// 4. message_end +type ExecuteProps struct { + Tool string `json:"tool"` // Tool name (e.g., "Bash", "Read", "Write", "mcp__github__search") + ToolID string `json:"tool_id"` // Agent-side tool call ID (e.g., "toolu_abc123") + Input interface{} `json:"input,omitempty"` // Tool input (structured, e.g., {"command":"ls -la"}) + Output interface{} `json:"output,omitempty"` // Tool execution result (via delta merge) + Status string `json:"status"` // "running" | "completed" | "error" + + IsError bool `json:"is_error,omitempty"` // Whether the tool execution failed + ExitCode *int `json:"exit_code,omitempty"` // Process exit code (Bash-type tools) + + Runner string `json:"runner,omitempty"` // Runner identifier (e.g., "claude-cli", "codex-cli") + Metadata map[string]interface{} `json:"metadata,omitempty"` // Extensible metadata per runner +} + // Delta action constants for incremental updates const ( DeltaAppend = "append" // Append (for arrays, strings) @@ -224,6 +255,7 @@ const ( ChunkToolCall StreamChunkType = "tool_call" // Tool/function call ChunkRefusal StreamChunkType = "refusal" // Model refusal ChunkMetadata StreamChunkType = "metadata" // Metadata (usage, finish_reason, etc.) + ChunkExecute StreamChunkType = "execute" // Agent execution observation (sandbox CLI agent tool use) ChunkError StreamChunkType = "error" // Error chunk ChunkUnknown StreamChunkType = "unknown" // Unknown/unrecognized chunk type diff --git a/agent/sandbox/v2/claude/command.go b/agent/sandbox/v2/claude/command.go new file mode 100644 index 00000000..42639b28 --- /dev/null +++ b/agent/sandbox/v2/claude/command.go @@ -0,0 +1,320 @@ +package claude + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/yaoapp/gou/connector" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +const defaultProxyPort = 3456 + +type command struct { + shell []string + env map[string]string + stdin []byte + workDir string +} + +func (r *ClaudeRunner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command { + computer := req.Computer + workDir := computer.GetWorkDir() + assistantID := "" + if req.Config != nil { + assistantID = req.Config.ID + } + + isContinuation := hasExistingSession(ctx, computer, p, assistantID) + env := buildEnv(req, p) + args := buildArgs(req, r, p, isContinuation, assistantID) + inputJSONL := buildInput(req.Messages, isContinuation) + + var systemPrompt string + envPrompt := buildSandboxEnvPrompt(p, workDir) + if !isContinuation && req.SystemPrompt != "" { + systemPrompt = req.SystemPrompt + "\n\n" + envPrompt + } else if !isContinuation { + systemPrompt = envPrompt + } + + promptFile := p.PathJoin(workDir, ".yao", "assistants", assistantID, "system-prompt.txt") + if assistantID == "" { + promptFile = p.PathJoin(workDir, ".yao", ".system-prompt.txt") + } + + script, stdin := p.BuildScript(scriptInput{ + args: args, + systemPrompt: systemPrompt, + inputJSONL: inputJSONL, + workDir: workDir, + promptFile: promptFile, + }) + + return command{ + shell: p.ShellCmd(script), + env: env, + stdin: stdin, + workDir: workDir, + } +} + +func buildEnv(req *types.StreamRequest, p platform) map[string]string { + env := make(map[string]string) + workDir := req.Computer.GetWorkDir() + + for k, v := range p.HomeEnv(workDir) { + env[k] = v + } + + assistantID := "" + if req.Config != nil { + assistantID = req.Config.ID + } + if assistantID != "" { + configDir := p.PathJoin(workDir, ".yao", "assistants", assistantID) + env["CLAUDE_CONFIG_DIR"] = configDir + } + + if req.Connector != nil { + setting := req.Connector.Setting() + host, _ := setting["host"].(string) + key, _ := setting["key"].(string) + model, _ := setting["model"].(string) + + if req.Connector.Is(connector.ANTHROPIC) { + env["ANTHROPIC_BASE_URL"] = host + env["ANTHROPIC_API_KEY"] = key + } else { + env["ANTHROPIC_BASE_URL"] = fmt.Sprintf("http://127.0.0.1:%d", defaultProxyPort) + env["ANTHROPIC_API_KEY"] = "dummy" + } + + if model != "" { + env["ANTHROPIC_MODEL"] = model + env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model + env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model + env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model + env["CLAUDE_CODE_SUBAGENT_MODEL"] = model + } + + if thinking, ok := setting["thinking"].(map[string]interface{}); ok { + thinkType, _ := thinking["type"].(string) + switch thinkType { + case "disabled": + env["MAX_THINKING_TOKENS"] = "0" + case "enabled": + if budget, ok := thinking["budget_tokens"].(float64); ok && budget > 0 { + env["MAX_THINKING_TOKENS"] = fmt.Sprintf("%d", int(budget)) + } + } + } + } + + if req.Config != nil && len(req.Config.Secrets) > 0 { + for k, v := range req.Config.Secrets { + env[k] = v + } + } + + if req.Token != nil { + if req.Token.Token != "" { + env["YAO_TOKEN"] = req.Token.Token + } + if req.Token.RefreshToken != "" { + env["YAO_REFRESH_TOKEN"] = req.Token.RefreshToken + } + } + + return env +} + +func buildArgs(req *types.StreamRequest, r *ClaudeRunner, p platform, isContinuation bool, assistantID string) []string { + var args []string + + permMode := "" + if req.Config != nil && req.Config.Runner.Options != nil { + if v, ok := req.Config.Runner.Options["permission_mode"]; ok { + permMode = fmt.Sprintf("%v", v) + } + } + if permMode == "bypassPermissions" { + args = append(args, "--dangerously-skip-permissions") + args = append(args, "--permission-mode", permMode) + } + + args = append(args, "--input-format", "stream-json") + args = append(args, "--output-format", "stream-json") + args = append(args, "--include-partial-messages") + args = append(args, "--verbose") + + if isContinuation { + args = append(args, "--continue") + } + + if req.Config != nil && req.Config.Runner.Options != nil { + for key, val := range req.Config.Runner.Options { + if flag, ok := claudeArgWhitelist[key]; ok { + args = append(args, flag, fmt.Sprintf("%v", val)) + } + } + } + + if r.hasMCP { + workDir := req.Computer.GetWorkDir() + mcpPath := p.PathJoin(workDir, ".yao", "assistants", assistantID, "mcp.json") + if assistantID == "" { + mcpPath = p.PathJoin(workDir, ".claude", "mcp.json") + } + args = append(args, "--mcp-config", mcpPath) + if r.mcpToolPattern != "" { + args = append(args, "--allowedTools", r.mcpToolPattern) + } + } + + return args +} + +func buildInput(messages []agentContext.Message, isContinuation bool) string { + if isContinuation { + return buildLastUserMessageJSONL(messages) + } + return buildFirstRequestJSONL(messages) +} + +func buildSandboxEnvPrompt(p platform, workDir string) string { + osName := p.OS() + if osName == "" { + osName = "linux" + } + shell := p.Shell() + if shell == "" { + shell = "bash" + } + + shellNote := p.EnvPromptNote() + + return fmt.Sprintf(`## Sandbox Environment + +- **Operating System**: %[2]s +- **Shell**: %[3]s +- **Working Directory**: %[1]s +- **File Access**: You have full read/write access to %[1]s%[4]s + +## User Attachments + +User-uploaded files (images, documents, code files, etc.) are placed in %[1]s/.attachments/{chatID}/ +Each chat session has its own subdirectory to avoid conflicts. +When the user references an attached file, read it from this directory using the Read or Bash tool. +For image files, you can view them directly as Claude supports vision on local files. +`, workDir, osName, shell, shellNote) +} + +func hasExistingSession(ctx context.Context, computer infra.Computer, p platform, assistantID string) bool { + workDir := computer.GetWorkDir() + var sessionDir string + if assistantID != "" { + configDir := p.PathJoin(workDir, ".yao", "assistants", assistantID) + sessionDir = p.PathJoin(configDir, "projects") + } else { + sessionDir = p.PathJoin(workDir, ".claude", "projects") + } + result, err := computer.Exec(ctx, p.ListDirCmd(sessionDir)) + if err != nil || result.ExitCode != 0 { + return false + } + return strings.TrimSpace(result.Stdout) != "" +} + +func buildMCPConfig(servers []types.MCPServer) []byte { + mcpServers := make(map[string]any, len(servers)) + for _, s := range servers { + name := s.ServerID + if name == "" { + continue + } + mcpServers[name] = map[string]any{ + "command": "tai", + "args": []string{"mcp", name}, + } + } + if len(mcpServers) == 0 { + mcpServers["yao"] = map[string]any{ + "command": "tai", + "args": []string{"mcp"}, + } + } + config := map[string]any{"mcpServers": mcpServers} + data, _ := json.Marshal(config) + return data +} + +func buildMCPAllowedTools(servers []types.MCPServer) string { + patterns := make([]string, 0, len(servers)) + for _, s := range servers { + if s.ServerID != "" { + patterns = append(patterns, fmt.Sprintf("mcp__%s__*", s.ServerID)) + } + } + if len(patterns) == 0 { + return "mcp__yao__*" + } + return strings.Join(patterns, ",") +} + +// buildFirstRequestJSONL builds the input JSONL for a new (non-continuation) +// Claude CLI session. Per the stream-json input protocol, only user messages +// should be sent; Claude CLI manages its own assistant history internally. +func buildFirstRequestJSONL(messages []agentContext.Message) string { + var lines []string + for _, msg := range messages { + if msg.Role != "user" { + continue + } + content := msg.Content + if content == nil { + content = "" + } + streamMsg := map[string]any{ + "type": "user", + "message": map[string]any{ + "role": "user", + "content": content, + }, + } + data, _ := json.Marshal(streamMsg) + lines = append(lines, string(data)) + } + return strings.Join(lines, "\n") +} + +func buildLastUserMessageJSONL(messages []agentContext.Message) string { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == "user" { + content := messages[i].Content + if content == nil { + content = "" + } + msg := map[string]any{ + "type": "user", + "message": map[string]any{ + "role": "user", + "content": content, + }, + } + data, _ := json.Marshal(msg) + return string(data) + } + } + return "" +} + +var claudeArgWhitelist = map[string]string{ + "max_turns": "--max-turns", + "disallowed_tools": "--disallowed-tools", + "allowed_tools": "--allowedTools", +} diff --git a/agent/sandbox/v2/claude/command_test.go b/agent/sandbox/v2/claude/command_test.go new file mode 100644 index 00000000..67e554fa --- /dev/null +++ b/agent/sandbox/v2/claude/command_test.go @@ -0,0 +1,384 @@ +package claude + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai/workspace" +) + +func testPlatform() platform { + return &darwinPlatform{posixBase: posixBase{ + os: "darwin", workDir: "/workspace", shell: "bash", tempDir: "/tmp", + }} +} + +// --- buildEnv --- + +func TestBuildEnv_HomeEnv(t *testing.T) { + req := &types.StreamRequest{Config: &types.SandboxConfig{}} + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "/workspace", env["HOME"]) +} + +func TestBuildEnv_ConfigDirIsolation(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ID: "my-assistant"}, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + expected := "/workspace/.yao/assistants/my-assistant" + assert.Equal(t, expected, env["CLAUDE_CONFIG_DIR"]) +} + +func TestBuildEnv_NoAssistantID(t *testing.T) { + req := &types.StreamRequest{Config: &types.SandboxConfig{}} + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + _, hasConfigDir := env["CLAUDE_CONFIG_DIR"] + assert.False(t, hasConfigDir, "should not set CLAUDE_CONFIG_DIR without assistantID") +} + +func TestBuildEnv_Token(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{}, + Token: &types.SandboxToken{ + Token: "tok123", + RefreshToken: "ref456", + }, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "tok123", env["YAO_TOKEN"]) + assert.Equal(t, "ref456", env["YAO_REFRESH_TOKEN"]) +} + +func TestBuildEnv_Secrets(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Secrets: map[string]string{ + "MY_SECRET": "secret_val", + }, + }, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "secret_val", env["MY_SECRET"]) +} + +// --- buildArgs --- + +func TestBuildArgs_Default(t *testing.T) { + req := &types.StreamRequest{Config: &types.SandboxConfig{}} + req.Computer = newFakeComputer("/workspace") + r := &ClaudeRunner{} + p := testPlatform() + + args := buildArgs(req, r, p, false, "") + assert.Contains(t, args, "--input-format") + assert.Contains(t, args, "stream-json") + assert.Contains(t, args, "--output-format") + assert.Contains(t, args, "--verbose") + assert.Contains(t, args, "--include-partial-messages") +} + +func TestBuildArgs_Continuation(t *testing.T) { + req := &types.StreamRequest{Config: &types.SandboxConfig{}} + req.Computer = newFakeComputer("/workspace") + r := &ClaudeRunner{} + p := testPlatform() + + args := buildArgs(req, r, p, true, "") + assert.Contains(t, args, "--continue") +} + +func TestBuildArgs_PermissionMode(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Options: map[string]interface{}{ + "permission_mode": "bypassPermissions", + }, + }, + }, + } + req.Computer = newFakeComputer("/workspace") + r := &ClaudeRunner{} + p := testPlatform() + + args := buildArgs(req, r, p, false, "") + assert.Contains(t, args, "--dangerously-skip-permissions") + assert.Contains(t, args, "--permission-mode") +} + +func TestBuildArgs_MCP(t *testing.T) { + req := &types.StreamRequest{Config: &types.SandboxConfig{}} + req.Computer = newFakeComputer("/workspace") + r := &ClaudeRunner{hasMCP: true, mcpToolPattern: "mcp__yao__*"} + p := testPlatform() + + args := buildArgs(req, r, p, false, "test-assistant") + assert.Contains(t, args, "--mcp-config") + assert.Contains(t, args, "--allowedTools") + assert.Contains(t, args, "mcp__yao__*") + + mcpIdx := -1 + for i, a := range args { + if a == "--mcp-config" { + mcpIdx = i + break + } + } + require.Greater(t, mcpIdx, -1) + mcpPath := args[mcpIdx+1] + assert.Contains(t, mcpPath, "test-assistant") +} + +func TestBuildArgs_WhitelistOptions(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Options: map[string]interface{}{ + "max_turns": 10, + }, + }, + }, + } + req.Computer = newFakeComputer("/workspace") + r := &ClaudeRunner{} + p := testPlatform() + + args := buildArgs(req, r, p, false, "") + assert.Contains(t, args, "--max-turns") +} + +// --- buildInput --- + +func TestBuildInput_FirstRequest(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + result := buildInput(msgs, false) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Len(t, lines, 1, "first request only includes user messages") + assert.Contains(t, lines[0], "hello") + assert.NotContains(t, result, "hi") +} + +func TestBuildInput_Continuation(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "user", Content: "first question"}, + {Role: "assistant", Content: "first answer"}, + {Role: "user", Content: "follow up"}, + } + result := buildInput(msgs, true) + var parsed map[string]any + err := json.Unmarshal([]byte(result), &parsed) + require.NoError(t, err) + assert.Equal(t, "user", parsed["type"]) + msg, _ := parsed["message"].(map[string]any) + assert.Equal(t, "follow up", msg["content"]) +} + +// --- buildFirstRequestJSONL --- + +func TestBuildFirstRequestJSONL_SkipsSystem(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "system", Content: "system prompt"}, + {Role: "user", Content: "hello"}, + } + result := buildFirstRequestJSONL(msgs) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Len(t, lines, 1, "system messages should be skipped") + assert.Contains(t, lines[0], "hello") +} + +func TestBuildFirstRequestJSONL_OnlyUserMessages(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "user", Content: "q1"}, + {Role: "assistant", Content: "a1"}, + {Role: "user", Content: "q2"}, + } + result := buildFirstRequestJSONL(msgs) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Len(t, lines, 2, "only user messages should be included") + for _, line := range lines { + var parsed map[string]any + err := json.Unmarshal([]byte(line), &parsed) + require.NoError(t, err) + assert.Equal(t, "user", parsed["type"]) + msg, _ := parsed["message"].(map[string]any) + assert.Equal(t, "user", msg["role"]) + } +} + +func TestBuildFirstRequestJSONL_AssistantOnlyMessages(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "assistant", Content: "a1"}, + {Role: "assistant", Content: "a2"}, + {Role: "assistant", Content: "a3"}, + {Role: "user", Content: "q1"}, + } + result := buildFirstRequestJSONL(msgs) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Len(t, lines, 1, "only the user message should be included") + assert.Contains(t, lines[0], "q1") + assert.NotContains(t, result, "a1") +} + +func TestBuildFirstRequestJSONL_NilContent(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "user", Content: nil}, + } + result := buildFirstRequestJSONL(msgs) + var parsed map[string]any + err := json.Unmarshal([]byte(strings.TrimSpace(result)), &parsed) + require.NoError(t, err) + msg, _ := parsed["message"].(map[string]any) + assert.Equal(t, "", msg["content"]) +} + +// --- buildLastUserMessageJSONL --- + +func TestBuildLastUserMessageJSONL_Found(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "user", Content: "first"}, + {Role: "assistant", Content: "reply"}, + {Role: "user", Content: "second"}, + } + result := buildLastUserMessageJSONL(msgs) + var parsed map[string]any + err := json.Unmarshal([]byte(result), &parsed) + require.NoError(t, err) + msg, _ := parsed["message"].(map[string]any) + assert.Equal(t, "second", msg["content"]) +} + +func TestBuildLastUserMessageJSONL_NoUser(t *testing.T) { + msgs := []agentContext.Message{ + {Role: "assistant", Content: "only assistant"}, + } + assert.Empty(t, buildLastUserMessageJSONL(msgs)) +} + +func TestBuildLastUserMessageJSONL_Empty(t *testing.T) { + assert.Empty(t, buildLastUserMessageJSONL(nil)) +} + +// --- buildMCPConfig --- + +func TestBuildMCPConfig_WithServers(t *testing.T) { + servers := []types.MCPServer{ + {ServerID: "server1"}, + {ServerID: "server2"}, + } + data := buildMCPConfig(servers) + var cfg map[string]any + err := json.Unmarshal(data, &cfg) + require.NoError(t, err) + mcpServers, ok := cfg["mcpServers"].(map[string]any) + require.True(t, ok) + assert.Contains(t, mcpServers, "server1") + assert.Contains(t, mcpServers, "server2") +} + +func TestBuildMCPConfig_Empty(t *testing.T) { + data := buildMCPConfig(nil) + var cfg map[string]any + err := json.Unmarshal(data, &cfg) + require.NoError(t, err) + mcpServers, ok := cfg["mcpServers"].(map[string]any) + require.True(t, ok) + assert.Contains(t, mcpServers, "yao", "should default to yao server") +} + +func TestBuildMCPConfig_EmptyServerID(t *testing.T) { + servers := []types.MCPServer{{ServerID: ""}} + data := buildMCPConfig(servers) + var cfg map[string]any + json.Unmarshal(data, &cfg) + mcpServers, _ := cfg["mcpServers"].(map[string]any) + assert.Contains(t, mcpServers, "yao", "empty serverID should fall back to default") +} + +// --- buildMCPAllowedTools --- + +func TestBuildMCPAllowedTools_WithServers(t *testing.T) { + servers := []types.MCPServer{ + {ServerID: "s1"}, + {ServerID: "s2"}, + } + result := buildMCPAllowedTools(servers) + assert.Contains(t, result, "mcp__s1__*") + assert.Contains(t, result, "mcp__s2__*") + assert.Contains(t, result, ",") +} + +func TestBuildMCPAllowedTools_Empty(t *testing.T) { + assert.Equal(t, "mcp__yao__*", buildMCPAllowedTools(nil)) +} + +// --- buildSandboxEnvPrompt --- + +func TestBuildSandboxEnvPrompt(t *testing.T) { + p := testPlatform() + prompt := buildSandboxEnvPrompt(p, "/workspace") + assert.Contains(t, prompt, "/workspace") + assert.Contains(t, prompt, "darwin") + assert.Contains(t, prompt, "bash") + assert.Contains(t, prompt, "Sandbox Environment") + assert.Contains(t, prompt, ".attachments") +} + +func TestBuildSandboxEnvPrompt_WindowsPlatform(t *testing.T) { + p := newWindowsPlatform(`C:\ws`, "pwsh", "") + prompt := buildSandboxEnvPrompt(p, `C:\ws`) + assert.Contains(t, prompt, "windows") + assert.Contains(t, prompt, "pwsh") + assert.Contains(t, prompt, `C:\ws`) + assert.Contains(t, prompt, "Windows desktop") +} + +// --- fakeComputer implements infra.Computer for unit tests --- + +type fakeComputer struct { + workDir string +} + +func newFakeComputer(workDir string) *fakeComputer { + return &fakeComputer{workDir: workDir} +} + +func (f *fakeComputer) GetWorkDir() string { return f.workDir } +func (f *fakeComputer) BindWorkplace(string) {} +func (f *fakeComputer) Workplace() workspace.FS { return nil } +func (f *fakeComputer) ComputerInfo() infra.ComputerInfo { + return infra.ComputerInfo{System: infra.SystemInfo{OS: "linux", Shell: "bash"}} +} +func (f *fakeComputer) Exec(_ context.Context, _ []string, _ ...infra.ExecOption) (*infra.ExecResult, error) { + return &infra.ExecResult{}, nil +} +func (f *fakeComputer) Stream(_ context.Context, _ []string, _ ...infra.ExecOption) (*infra.ExecStream, error) { + return nil, nil +} +func (f *fakeComputer) VNC(_ context.Context) (string, error) { return "", nil } +func (f *fakeComputer) Proxy(_ context.Context, _ int, _ string) (string, error) { return "", nil } diff --git a/agent/sandbox/v2/claude/oscompat.go b/agent/sandbox/v2/claude/oscompat.go deleted file mode 100644 index 7c7ddf3f..00000000 --- a/agent/sandbox/v2/claude/oscompat.go +++ /dev/null @@ -1,187 +0,0 @@ -package claude - -import ( - "fmt" - "path" - "strings" - - "github.com/yaoapp/yao/agent/sandbox/v2/types" - infra "github.com/yaoapp/yao/sandbox/v2" -) - -// osEnv captures OS-dependent paths and shell settings derived from the -// Computer's SystemInfo. All runner code should use osEnv instead of -// hardcoded Linux constants. -type osEnv struct { - OS string // "windows", "linux", "darwin", ... - Shell string // preferred shell binary: "bash", "pwsh", "cmd.exe", ... - WorkDir string // working directory on the target machine - UserHome string // user home directory (empty if irrelevant) - TempDir string // system temp directory -} - -func (e *osEnv) isWindows() bool { - return strings.EqualFold(e.OS, "windows") -} - -// resolveOSEnv builds an osEnv from the Computer's reported SystemInfo, -// falling back to SandboxConfig values where available, then to per-OS defaults. -func resolveOSEnv(computer infra.Computer, _ *types.SandboxConfig) *osEnv { - sys := computer.ComputerInfo().System - - env := &osEnv{ - OS: strings.ToLower(sys.OS), - Shell: sys.Shell, - TempDir: sys.TempDir, - WorkDir: computer.GetWorkDir(), - } - - if env.TempDir == "" { - env.TempDir = env.pathJoin(env.WorkDir, ".tmp") - } - - return env -} - -// shellCmd returns the command slice to run a script through the appropriate shell. -func (e *osEnv) shellCmd(script string) []string { - shell := strings.ToLower(e.Shell) - switch shell { - case "pwsh": - return []string{"pwsh", "-NoProfile", "-Command", script} - case "powershell": - return []string{"powershell", "-NoProfile", "-Command", script} - case "cmd.exe", "cmd": - return []string{"cmd.exe", "/C", script} - default: - return []string{"bash", "-c", script} - } -} - -// mkdirCmd returns a shell command string to create a directory (with parents). -func (e *osEnv) mkdirCmd(dir string) string { - if e.isWindows() { - return fmt.Sprintf(`if (!(Test-Path '%s')) { New-Item -ItemType Directory -Path '%s' -Force | Out-Null }`, dir, dir) - } - return fmt.Sprintf("mkdir -p %s", dir) -} - -// listDirCmd returns a command slice to list directory contents. -func (e *osEnv) listDirCmd(dir string) []string { - if e.isWindows() { - return e.shellCmd(fmt.Sprintf("Get-ChildItem -Name '%s'", dir)) - } - return []string{"ls", dir} -} - -// killProcessCmd returns a command slice to kill processes matching a pattern. -// On Windows, uses taskkill /T to kill the entire process tree, which handles -// child processes (chrome.exe, python3, etc.) that Claude CLI may have spawned. -func (e *osEnv) killProcessCmd(pattern string) []string { - if e.isWindows() { - // taskkill /F /T kills the process tree; fall back to Stop-Process. - script := fmt.Sprintf( - "Get-Process -ErrorAction SilentlyContinue | Where-Object {$_.ProcessName -like '*%s*'} | ForEach-Object { taskkill /F /T /PID $_.Id 2>$null }; "+ - "Get-Process -ErrorAction SilentlyContinue | Where-Object {$_.ProcessName -like '*%s*'} | Stop-Process -Force -ErrorAction SilentlyContinue", - pattern, pattern) - return e.shellCmd(script) - } - return []string{"sh", "-c", fmt.Sprintf("pkill -f '%s' || true", pattern)} -} - -// rootDir returns the filesystem root for the target OS. -func (e *osEnv) rootDir() string { - if e.isWindows() { - return `C:\` - } - return "/" -} - -// pathJoin joins path segments using the appropriate separator. -func (e *osEnv) pathJoin(parts ...string) string { - if e.isWindows() { - return strings.Join(parts, `\`) - } - return path.Join(parts...) -} - -// buildCLIScript builds the complete CLI invocation script for the target OS. -// Returns (script, stdin) — on Linux stdin is nil (heredoc handles it), -// on Windows stdin contains inputJSONL bytes to pass via gRPC Stdin. -func (e *osEnv) buildCLIScript(args []string, systemPrompt, inputJSONL string) (string, []byte) { - workDir := e.WorkDir - promptFile := e.pathJoin(workDir, ".yao", ".system-prompt.txt") - - if e.isWindows() { - return e.buildPowerShellScript(args, systemPrompt, inputJSONL, workDir, promptFile) - } - return e.buildBashScript(args, systemPrompt, inputJSONL, workDir, promptFile), nil -} - -func (e *osEnv) buildBashScript(args []string, systemPrompt, inputJSONL, workDir, promptFile string) string { - var b strings.Builder - - if e.UserHome != "" { - b.WriteString(fmt.Sprintf("touch %s/.Xauthority 2>/dev/null; ", e.UserHome)) - } - b.WriteString("touch \"$HOME/.Xauthority\" 2>/dev/null\n") - - if systemPrompt != "" { - b.WriteString(fmt.Sprintf("mkdir -p %s/.yao\n", workDir)) - b.WriteString(fmt.Sprintf("cat << 'PROMPTEOF' > %s\n", promptFile)) - b.WriteString(systemPrompt) - b.WriteString("\nPROMPTEOF\n") - args = append(args, "--append-system-prompt-file", promptFile) - } - - b.WriteString("cat << 'INPUTEOF' | claude -p") - for _, arg := range args { - b.WriteString(fmt.Sprintf(" %q", arg)) - } - b.WriteString(" 2>&1\n") - b.WriteString(inputJSONL) - b.WriteString("\nINPUTEOF") - - return b.String() -} - -// buildPowerShellScript builds a script that writes the system prompt file, -// then launches claude -p. inputJSONL is returned as stdin bytes to be passed -// directly via gRPC, bypassing PowerShell's encoding entirely. -func (e *osEnv) buildPowerShellScript(args []string, systemPrompt, inputJSONL, workDir, promptFile string) (string, []byte) { - var b strings.Builder - noBOM := "(New-Object System.Text.UTF8Encoding $false)" - - // Force UTF-8 for both input and output streams. - // On CJK Windows the default codepage is often GBK/GB2312 (936) - // which corrupts Claude CLI's UTF-8 JSON output. - b.WriteString("[Console]::InputEncoding = [System.Text.Encoding]::UTF8\n") - b.WriteString("[Console]::OutputEncoding = [System.Text.Encoding]::UTF8\n") - b.WriteString("$OutputEncoding = [System.Text.Encoding]::UTF8\n") - - // Ensure claude.exe can be found even when Tai runs as a different user. - // Claude CLI is typically installed per-user (e.g. C:\Users\X\.local\bin) - // which isn't in the PATH when Tai runs as a service or another account. - // Scan all user profiles for common install locations. - b.WriteString("foreach ($d in (Get-ChildItem 'C:\\Users' -Directory -ErrorAction SilentlyContinue)) {\n") - b.WriteString(" $p = Join-Path $d.FullName '.local\\bin'\n") - b.WriteString(" if (Test-Path (Join-Path $p 'claude.exe')) { $env:PATH = \"$p;$env:PATH\"; break }\n") - b.WriteString("}\n") - b.WriteString("if ($env:APPDATA) { $env:PATH = \"$env:APPDATA\\npm;$env:PATH\" }\n") - - yaoDir := e.pathJoin(workDir, ".yao") - b.WriteString(fmt.Sprintf("if (!(Test-Path '%s')) { New-Item -ItemType Directory -Path '%s' -Force | Out-Null }\n", yaoDir, yaoDir)) - - if systemPrompt != "" { - escaped := strings.ReplaceAll(systemPrompt, "'", "''") - b.WriteString(fmt.Sprintf("[IO.File]::WriteAllText('%s', @'\n%s\n'@, %s)\n", promptFile, escaped, noBOM)) - args = append(args, "--append-system-prompt-file", promptFile) - } - - b.WriteString("claude -p") - for _, arg := range args { - b.WriteString(fmt.Sprintf(" '%s'", strings.ReplaceAll(arg, "'", "''"))) - } - - return b.String(), []byte(inputJSONL + "\n") -} diff --git a/agent/sandbox/v2/claude/parse.go b/agent/sandbox/v2/claude/parse.go index a7acc435..86cfb227 100644 --- a/agent/sandbox/v2/claude/parse.go +++ b/agent/sandbox/v2/claude/parse.go @@ -4,30 +4,56 @@ import ( "bufio" "context" "encoding/json" - "errors" "fmt" "io" - "log" "strings" "time" - goujson "github.com/yaoapp/gou/json" - agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/output/message" ) -// errStreamCompleted is a sentinel indicating the parser received the terminal -// "result" message. It is NOT a real error — callers should treat it as -// successful completion of the stream. -var errStreamCompleted = errors.New("claude stream completed") +// streamParser is an explicit state machine for Claude CLI stream-json output. +// +// Each tool call gets its own message lifecycle: +// +// content_block_start -> message_start(id=exec-N-xxx) + ChunkExecute{tool, status:running} +// input_json_delta -> ChunkExecute{input_delta:...} (same message group) +// content_block_stop -> message_end(exec-N-xxx) +// ...later... +// user/tool_result -> message_start(id=exec-N-xxx, reuse!) + ChunkExecute{status:completed, output:...} + message_end +type streamParser struct { + handler message.StreamFunc + completed bool -// parseStreamJSON reads stream-json lines from Claude CLI stdout and -// pushes them through handler as standard StreamChunkType events. -func parseStreamJSON(ctx context.Context, stdout io.ReadCloser, handler message.StreamFunc) error { - // When the context is cancelled (upstream timeout / interrupt), close - // stdout so that scanner.Scan() unblocks immediately. Without this, - // a failed TerminateProcess (Access is denied) would leave us stuck - // forever on the read. + textActive bool + toolIndex int + curTool *toolState + + toolNames map[string]string // tool_id -> tool_name + toolMsgIDs map[string]string // tool_id -> message_id (for result reuse) + toolInputs map[string]string // tool_id -> full input JSON (for result replay) + toolSummaries map[string]string // tool_id -> summary (for result replay) +} + +type toolState struct { + id string + name string + msgID string + index int + inputJSON strings.Builder +} + +func newStreamParser(handler message.StreamFunc) *streamParser { + return &streamParser{ + handler: handler, + toolNames: make(map[string]string), + toolMsgIDs: make(map[string]string), + toolInputs: make(map[string]string), + toolSummaries: make(map[string]string), + } +} + +func (p *streamParser) parse(ctx context.Context, stdout io.ReadCloser) error { doneParsing := make(chan struct{}) defer close(doneParsing) go func() { @@ -39,20 +65,7 @@ func parseStreamJSON(ctx context.Context, stdout io.ReadCloser, handler message. }() scanner := bufio.NewScanner(stdout) - buf := make([]byte, 0, 64*1024) - scanner.Buffer(buf, 1024*1024) - - messageStarted := false - toolBlockActive := false - toolIndex := 0 - - type toolState struct { - id string - name string - index int - inputJSON strings.Builder - } - var currentTool *toolState + scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) for scanner.Scan() { line := scanner.Text() @@ -62,280 +75,39 @@ func parseStreamJSON(ctx context.Context, stdout io.ReadCloser, handler message. var msg map[string]any if err := json.Unmarshal([]byte(line), &msg); err != nil { + if len(line) > 200 { + fmt.Printf("[claude-parse] JSON unmarshal error: %v (line len=%d, prefix=%q)\n", err, len(line), line[:200]) + } else { + fmt.Printf("[claude-parse] JSON unmarshal error: %v (line=%q)\n", err, line) + } continue } msgType, _ := msg["type"].(string) - stopped := false + var stopped bool switch msgType { case "system": - if handler != nil { - data, _ := json.Marshal(msg) - if handler(message.ChunkMetadata, data) != 0 { - stopped = true - } - } - + stopped = p.handleSystem(msg) case "stream_event": - event, _ := msg["event"].(map[string]any) - if event == nil { - continue - } - eventType, _ := event["type"].(string) - - switch eventType { - case "content_block_start": - if cb, ok := event["content_block"].(map[string]any); ok { - blockType, _ := cb["type"].(string) - if blockType == "tool_use" { - toolName, _ := cb["name"].(string) - toolID, _ := cb["id"].(string) - if toolID == "" { - toolID = fmt.Sprintf("tool_%d_%d", toolIndex, time.Now().UnixNano()) - } - currentTool = &toolState{id: toolID, name: toolName, index: toolIndex} - toolIndex++ - - if handler != nil { - if messageStarted { - handler(message.ChunkMessageEnd, nil) - messageStarted = false - } - if !toolBlockActive { - startData := message.EventMessageStartData{ - MessageID: fmt.Sprintf("sandbox-tool-%d", time.Now().UnixNano()), - Type: "tool_call", - Timestamp: time.Now().UnixMilli(), - } - sd, _ := json.Marshal(startData) - if handler(message.ChunkMessageStart, sd) != 0 { - stopped = true - break - } - toolBlockActive = true - } - tcData, _ := json.Marshal([]map[string]any{{ - "index": currentTool.index, - "id": currentTool.id, - "type": "function", - "function": map[string]any{ - "name": toolName, - "arguments": "", - }, - }}) - if handler(message.ChunkToolCall, tcData) != 0 { - stopped = true - } - } - } - } - - case "content_block_delta": - if delta, ok := event["delta"].(map[string]any); ok { - deltaType, _ := delta["type"].(string) - switch deltaType { - case "text_delta": - if text, ok := delta["text"].(string); ok && text != "" { - text = strings.ReplaceAll(text, "\r\n", "\n") - text = strings.ReplaceAll(text, "\r", "\n") - if handler != nil { - if toolBlockActive { - handler(message.ChunkMessageEnd, nil) - toolBlockActive = false - messageStarted = false - } - if !messageStarted { - startData := message.EventMessageStartData{ - MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()), - Type: "text", - Timestamp: time.Now().UnixMilli(), - } - sd, _ := json.Marshal(startData) - if handler(message.ChunkMessageStart, sd) != 0 { - stopped = true - break - } - messageStarted = true - } - if handler(message.ChunkText, []byte(text)) != 0 { - stopped = true - } - } - } - case "input_json_delta": - if currentTool != nil { - if partial, ok := delta["partial_json"].(string); ok { - currentTool.inputJSON.WriteString(partial) - if handler != nil { - tcData, _ := json.Marshal([]map[string]any{{ - "index": currentTool.index, - "function": map[string]any{ - "arguments": partial, - }, - }}) - if handler(message.ChunkToolCall, tcData) != 0 { - stopped = true - } - } - } - } - } - } - - case "content_block_stop": - currentTool = nil - } - + stopped = p.handleStreamEvent(msg) case "assistant": - if msgData, ok := msg["message"].(map[string]any); ok { - stopReason, _ := msgData["stop_reason"].(string) - if stopReason != "" { - if contentArr, ok := msgData["content"].([]any); ok { - for _, item := range contentArr { - ci, ok := item.(map[string]any) - if !ok { - continue - } - itemType, _ := ci["type"].(string) - - if itemType == "tool_use" && handler != nil { - toolName, _ := ci["name"].(string) - toolID, _ := ci["id"].(string) - if toolID == "" { - toolID = fmt.Sprintf("tool_%d_%d", toolIndex, time.Now().UnixNano()) - } - inputRaw, _ := json.Marshal(ci["input"]) - idx := toolIndex - toolIndex++ - - if !toolBlockActive { - startData := message.EventMessageStartData{ - MessageID: fmt.Sprintf("sandbox-tool-%d", time.Now().UnixNano()), - Type: "tool_call", - Timestamp: time.Now().UnixMilli(), - } - sd, _ := json.Marshal(startData) - if handler(message.ChunkMessageStart, sd) != 0 { - stopped = true - break - } - toolBlockActive = true - } - tcData, _ := json.Marshal([]map[string]any{{ - "index": idx, - "id": toolID, - "type": "function", - "function": map[string]any{ - "name": toolName, - "arguments": string(inputRaw), - }, - }}) - if handler(message.ChunkToolCall, tcData) != 0 { - stopped = true - break - } - } - - if itemType == "text" { - if text, ok := ci["text"].(string); ok && text != "" && handler != nil && !messageStarted { - text = strings.ReplaceAll(text, "\r\n", "\n") - text = strings.ReplaceAll(text, "\r", "\n") - if toolBlockActive { - handler(message.ChunkMessageEnd, nil) - toolBlockActive = false - } - startData := message.EventMessageStartData{ - MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()), - Type: "text", - Timestamp: time.Now().UnixMilli(), - } - sd, _ := json.Marshal(startData) - if handler(message.ChunkMessageStart, sd) != 0 { - stopped = true - break - } - if handler(message.ChunkText, []byte(text)) != 0 { - stopped = true - break - } - messageStarted = true - } - } - } - } - - // Close any open message from the streaming phase. - // stream_event text_deltas set messageStarted=true but - // nothing resets it when the turn ends — the assistant - // message marks the turn boundary, so we must close - // the message here to keep state in sync with the - // stream handler (which already sent message_end). - if handler != nil { - if toolBlockActive { - handler(message.ChunkMessageEnd, nil) - toolBlockActive = false - } - if messageStarted { - handler(message.ChunkMessageEnd, nil) - messageStarted = false - } - } - } - } - + stopped = p.handleAssistant(msg) + case "user": + stopped = p.handleUser(msg) case "result": - isError, _ := msg["is_error"].(bool) - if isError { - if result, ok := msg["result"].(string); ok { - if handler != nil { - handler(message.ChunkError, []byte(result)) - } - return fmt.Errorf("Claude CLI error: %s", result) - } - } - if handler != nil { - if toolBlockActive { - handler(message.ChunkMessageEnd, nil) - toolBlockActive = false - } - if messageStarted { - handler(message.ChunkMessageEnd, nil) - } - } - // "result" is the terminal message in Claude CLI's stream-json - // protocol. Return immediately instead of continuing to - // scanner.Scan(), which would block forever if the process - // stays alive (e.g. child processes like chrome.exe keep the - // stdout pipe open). - return errStreamCompleted - + return p.handleResult(msg) case "error": - var errMsg string - switch e := msg["error"].(type) { - case string: - errMsg = e - case map[string]any: - errMsg, _ = e["message"].(string) - } - if errMsg != "" { - if handler != nil { - handler(message.ChunkError, []byte(errMsg)) - } - return fmt.Errorf("Claude CLI error: %s", errMsg) - } + return p.handleError(msg) } if stopped { - break + return nil } } if err := scanner.Err(); err != nil { - // If the context was cancelled (upstream timeout / interrupt), the - // stdout pipe was closed by the goroutine above. The resulting - // read error is expected — surface it as context.Canceled so the - // caller can handle it uniformly. + fmt.Printf("[claude-parse] scanner error: %v (ctx.Err=%v)\n", err, ctx.Err()) if ctx.Err() != nil { return ctx.Err() } @@ -344,53 +116,422 @@ func parseStreamJSON(ctx context.Context, stdout io.ReadCloser, handler message. return nil } -// buildFirstRequestJSONL builds JSONL with all messages for the first request. -func buildFirstRequestJSONL(messages []agentContext.Message) string { - var lines []string - for _, msg := range messages { - if msg.Role == "system" { - continue - } - content := msg.Content - if content == nil { - content = "" - } - streamMsg := map[string]any{ - "type": string(msg.Role), - "message": map[string]any{ - "role": string(msg.Role), - "content": content, - }, - } - data, _ := json.Marshal(streamMsg) - lines = append(lines, string(data)) +// --- Message lifecycle helpers --- + +func (p *streamParser) beginMessageWithID(id, msgType string) (stopped bool) { + startData := message.EventMessageStartData{ + MessageID: id, + Type: msgType, + Timestamp: time.Now().UnixMilli(), } - return strings.Join(lines, "\n") + sd, _ := json.Marshal(startData) + return p.handler != nil && p.handler(message.ChunkMessageStart, sd) != 0 } -// buildLastUserMessageJSONL builds JSONL with only the last user message. -func buildLastUserMessageJSONL(messages []agentContext.Message) string { - for i := len(messages) - 1; i >= 0; i-- { - if messages[i].Role == "user" { - content := messages[i].Content - if content == nil { - content = "" - } - msg := map[string]any{ - "type": "user", - "message": map[string]any{ - "role": "user", - "content": content, - }, - } - data, _ := json.Marshal(msg) - return string(data) +func (p *streamParser) beginMessage(msgType string) (messageID string, stopped bool) { + id := fmt.Sprintf("sandbox-%s-%s", msgType, message.GenerateNanoID()) + return id, p.beginMessageWithID(id, msgType) +} + +func (p *streamParser) endMessage() { + if p.handler != nil { + p.handler(message.ChunkMessageEnd, nil) + } +} + +func (p *streamParser) closeTextMessage() { + if p.textActive { + p.endMessage() + p.textActive = false + } +} + +func (p *streamParser) ensureTextMessage() (stopped bool) { + if !p.textActive { + _, stopped = p.beginMessage("text") + if stopped { + return true + } + p.textActive = true + } + return false +} + +func (p *streamParser) emitText(text string) (stopped bool) { + text = strings.ReplaceAll(text, "\r\n", "\n") + text = strings.ReplaceAll(text, "\r", "\n") + return p.handler != nil && p.handler(message.ChunkText, []byte(text)) != 0 +} + +func (p *streamParser) emitExecute(props map[string]any) (stopped bool) { + data, _ := json.Marshal(props) + return p.handler != nil && p.handler(message.ChunkExecute, data) != 0 +} + +func (p *streamParser) emitMetadata(data map[string]any) { + if p.handler == nil { + return + } + encoded, _ := json.Marshal(data) + p.handler(message.ChunkMetadata, encoded) +} + +// extractSummary builds a short human-readable summary from the tool input JSON. +func extractSummary(toolName string, inputJSON string) string { + if inputJSON == "" { + return "" + } + var obj map[string]any + if err := json.Unmarshal([]byte(inputJSON), &obj); err != nil { + return "" + } + + switch strings.ToLower(toolName) { + case "bash", "execute": + if cmd, ok := obj["command"].(string); ok { + return truncate(cmd, 80) + } + case "write", "create": + if fp, ok := obj["file_path"].(string); ok { + return fp + } + case "read": + if fp, ok := obj["file_path"].(string); ok { + return fp + } + case "edit": + if fp, ok := obj["file_path"].(string); ok { + return fp + } + } + + // Fallback: try common field names + for _, key := range []string{"path", "file_path", "command", "url", "query"} { + if v, ok := obj[key].(string); ok { + return truncate(v, 80) } } return "" } -// Suppress unused import warnings — goujson.Parse is used for tool description -// parsing in V1 and will be used for detailed tool descriptions in future. -var _ = goujson.Parse -var _ = log.Printf +func truncate(s string, max int) string { + s = strings.TrimSpace(s) + s = strings.ReplaceAll(s, "\n", " ") + if len(s) > max { + return s[:max] + "..." + } + return s +} + +// --- Event handlers --- + +func (p *streamParser) handleSystem(msg map[string]any) (stopped bool) { + if p.handler != nil { + data, _ := json.Marshal(msg) + return p.handler(message.ChunkMetadata, data) != 0 + } + return false +} + +func (p *streamParser) handleStreamEvent(msg map[string]any) (stopped bool) { + event, _ := msg["event"].(map[string]any) + if event == nil { + return false + } + eventType, _ := event["type"].(string) + + switch eventType { + case "content_block_start": + return p.onContentBlockStart(event) + case "content_block_delta": + return p.onContentBlockDelta(event) + case "content_block_stop": + return p.onContentBlockStop() + } + return false +} + +func (p *streamParser) onContentBlockStart(event map[string]any) (stopped bool) { + cb, ok := event["content_block"].(map[string]any) + if !ok { + return false + } + blockType, _ := cb["type"].(string) + if blockType != "tool_use" { + return false + } + + p.closeTextMessage() + + toolName, _ := cb["name"].(string) + toolID, _ := cb["id"].(string) + if toolID == "" { + toolID = fmt.Sprintf("tool_%d_%d", p.toolIndex, time.Now().UnixNano()) + } + + msgID, stopped := p.beginMessage("execute") + if stopped { + return true + } + + p.curTool = &toolState{id: toolID, name: toolName, msgID: msgID, index: p.toolIndex} + p.toolIndex++ + p.toolNames[toolID] = toolName + p.toolMsgIDs[toolID] = msgID + + if p.handler == nil { + return false + } + + return p.emitExecute(map[string]any{ + "tool": toolName, + "tool_id": toolID, + "status": "running", + "runner": "claude-cli", + }) +} + +func (p *streamParser) onContentBlockStop() (stopped bool) { + if p.curTool != nil { + toolID := p.curTool.id + inputStr := p.curTool.inputJSON.String() + if inputStr != "" { + p.toolInputs[toolID] = inputStr + summary := extractSummary(p.curTool.name, inputStr) + if summary != "" { + p.toolSummaries[toolID] = summary + p.emitExecute(map[string]any{ + "summary": summary, + }) + } + } + p.endMessage() + p.curTool = nil + } + return false +} + +func (p *streamParser) onContentBlockDelta(event map[string]any) (stopped bool) { + delta, ok := event["delta"].(map[string]any) + if !ok { + return false + } + + deltaType, _ := delta["type"].(string) + + switch deltaType { + case "text_delta": + text, _ := delta["text"].(string) + if text == "" { + return false + } + if p.ensureTextMessage() { + return true + } + return p.emitText(text) + + case "input_json_delta": + if p.curTool == nil { + return false + } + partial, _ := delta["partial_json"].(string) + if partial == "" { + return false + } + p.curTool.inputJSON.WriteString(partial) + if p.handler != nil { + return p.emitExecute(map[string]any{ + "input_delta": p.curTool.inputJSON.String(), + }) + } + } + return false +} + +func (p *streamParser) handleAssistant(msg map[string]any) (stopped bool) { + msgData, _ := msg["message"].(map[string]any) + if msgData == nil { + return false + } + + if usage, ok := msgData["usage"].(map[string]any); ok { + p.emitMetadata(map[string]any{ + "usage": usage, + }) + } + + stopReason, _ := msgData["stop_reason"].(string) + if stopReason == "" { + return false + } + + contentArr, _ := msgData["content"].([]any) + for _, item := range contentArr { + ci, ok := item.(map[string]any) + if !ok { + continue + } + itemType, _ := ci["type"].(string) + + if itemType == "tool_use" && p.handler != nil { + p.closeTextMessage() + + toolName, _ := ci["name"].(string) + toolID, _ := ci["id"].(string) + if toolID == "" { + toolID = fmt.Sprintf("tool_%d_%d", p.toolIndex, time.Now().UnixNano()) + } + + msgID, stopped := p.beginMessage("execute") + if stopped { + return true + } + + p.toolIndex++ + p.toolNames[toolID] = toolName + p.toolMsgIDs[toolID] = msgID + + inputRaw, _ := json.Marshal(ci["input"]) + inputStr := string(inputRaw) + p.toolInputs[toolID] = inputStr + summary := extractSummary(toolName, inputStr) + p.toolSummaries[toolID] = summary + + if p.emitExecute(map[string]any{ + "tool": toolName, + "tool_id": toolID, + "input": json.RawMessage(inputRaw), + "summary": summary, + "status": "running", + "runner": "claude-cli", + }) { + p.endMessage() + return true + } + p.endMessage() + } + + if itemType == "text" { + text, _ := ci["text"].(string) + if text == "" || p.handler == nil { + continue + } + if p.ensureTextMessage() { + return true + } + if p.emitText(text) { + return true + } + } + } + + p.closeTextMessage() + return false +} + +func (p *streamParser) handleUser(msg map[string]any) (stopped bool) { + msgData, _ := msg["message"].(map[string]any) + if msgData == nil { + return false + } + + contentArr, _ := msgData["content"].([]interface{}) + for _, item := range contentArr { + ci, ok := item.(map[string]any) + if !ok { + continue + } + + ciType, _ := ci["type"].(string) + if ciType != "tool_result" { + continue + } + + toolUseID, _ := ci["tool_use_id"].(string) + content := ci["content"] + isError, _ := ci["is_error"].(bool) + + status := "completed" + if isError { + status = "error" + } + + execProps := map[string]any{ + "tool_id": toolUseID, + "output": content, + "status": status, + "is_error": isError, + } + if name, ok := p.toolNames[toolUseID]; ok { + execProps["tool"] = name + } + if input, ok := p.toolInputs[toolUseID]; ok { + execProps["input"] = json.RawMessage(input) + } + if summary, ok := p.toolSummaries[toolUseID]; ok { + execProps["summary"] = summary + } + + if reuseMsgID, ok := p.toolMsgIDs[toolUseID]; ok { + if p.beginMessageWithID(reuseMsgID, "execute") { + return true + } + } else { + if _, stopped := p.beginMessage("execute"); stopped { + return true + } + } + + if p.emitExecute(execProps) { + p.endMessage() + return true + } + p.endMessage() + } + return false +} + +func (p *streamParser) handleResult(msg map[string]any) error { + isError, _ := msg["is_error"].(bool) + if isError { + if result, ok := msg["result"].(string); ok { + if p.handler != nil { + p.handler(message.ChunkError, []byte(result)) + } + return fmt.Errorf("Claude CLI error: %s", result) + } + } + + p.closeTextMessage() + + if p.handler != nil { + p.emitMetadata(map[string]any{ + "result_summary": map[string]any{ + "total_cost_usd": msg["total_cost_usd"], + "duration_ms": msg["duration_ms"], + "num_turns": msg["num_turns"], + "usage": msg["usage"], + }, + }) + } + + p.completed = true + return nil +} + +func (p *streamParser) handleError(msg map[string]any) error { + var errMsg string + switch e := msg["error"].(type) { + case string: + errMsg = e + case map[string]any: + errMsg, _ = e["message"].(string) + } + if errMsg != "" { + if p.handler != nil { + p.handler(message.ChunkError, []byte(errMsg)) + } + return fmt.Errorf("Claude CLI error: %s", errMsg) + } + return nil +} diff --git a/agent/sandbox/v2/claude/parse_test.go b/agent/sandbox/v2/claude/parse_test.go new file mode 100644 index 00000000..a6c38c5d --- /dev/null +++ b/agent/sandbox/v2/claude/parse_test.go @@ -0,0 +1,630 @@ +package claude + +import ( + "context" + "encoding/json" + "io" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/output/message" +) + +type chunkRecord struct { + Type message.StreamChunkType + Data json.RawMessage +} + +func recordingHandler(out *[]chunkRecord) message.StreamFunc { + return func(chunkType message.StreamChunkType, data []byte) int { + cp := make([]byte, len(data)) + copy(cp, data) + *out = append(*out, chunkRecord{Type: chunkType, Data: cp}) + return 0 + } +} + +func stoppingHandler(stopAfter int) (message.StreamFunc, *[]chunkRecord) { + var out []chunkRecord + count := 0 + fn := func(chunkType message.StreamChunkType, data []byte) int { + cp := make([]byte, len(data)) + copy(cp, data) + out = append(out, chunkRecord{Type: chunkType, Data: cp}) + count++ + if count >= stopAfter { + return 1 + } + return 0 + } + return fn, &out +} + +func jsonLine(v interface{}) string { + b, _ := json.Marshal(v) + return string(b) +} + +func pipeWithLines(lines ...string) io.ReadCloser { + return io.NopCloser(strings.NewReader(strings.Join(lines, "\n") + "\n")) +} + +// --- helper: extract message_id from ChunkMessageStart data --- +func extractMessageID(data json.RawMessage) string { + var d map[string]any + json.Unmarshal(data, &d) + if id, ok := d["message_id"].(string); ok { + return id + } + return "" +} + +func TestParser_TextOnly(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{"type": "system", "session_id": "abc"}), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "text_delta", "text": "Hello "}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "text_delta", "text": "world"}, + }, + }), + jsonLine(map[string]any{ + "type": "assistant", + "message": map[string]any{ + "stop_reason": "end_turn", + "content": []any{map[string]any{"type": "text", "text": "Hello world"}}, + }, + }), + jsonLine(map[string]any{ + "type": "result", + "total_cost_usd": 0.001, + "duration_ms": 1234, + "num_turns": 1, + "usage": map[string]any{"input_tokens": 10, "output_tokens": 20}, + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + assert.True(t, p.completed) + + hasMessageStart := false + hasText := false + hasMessageEnd := false + hasResultMeta := false + for _, c := range chunks { + switch c.Type { + case message.ChunkMessageStart: + hasMessageStart = true + case message.ChunkText: + hasText = true + case message.ChunkMessageEnd: + hasMessageEnd = true + case message.ChunkMetadata: + var meta map[string]any + json.Unmarshal(c.Data, &meta) + if _, ok := meta["result_summary"]; ok { + hasResultMeta = true + } + } + } + assert.True(t, hasMessageStart, "should emit message_start") + assert.True(t, hasText, "should emit text chunks") + assert.True(t, hasMessageEnd, "should emit message_end") + assert.True(t, hasResultMeta, "should emit result_summary metadata") +} + +func TestParser_ToolUseAndResult(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{ + "type": "tool_use", + "name": "Bash", + "id": "tool_123", + }, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"command":"ls`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `"}`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + jsonLine(map[string]any{ + "type": "user", + "message": map[string]any{ + "content": []any{ + map[string]any{ + "type": "tool_result", + "tool_use_id": "tool_123", + "content": "file1.txt\nfile2.txt", + "is_error": false, + }, + }, + }, + }), + jsonLine(map[string]any{ + "type": "result", + "num_turns": 1, + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + assert.True(t, p.completed) + + var execChunks []map[string]any + for _, c := range chunks { + if c.Type == message.ChunkExecute { + var data map[string]any + json.Unmarshal(c.Data, &data) + execChunks = append(execChunks, data) + } + } + require.GreaterOrEqual(t, len(execChunks), 2, "should have at least 2 execute chunks (start + result)") + + assert.Equal(t, "Bash", execChunks[0]["tool"]) + assert.Equal(t, "tool_123", execChunks[0]["tool_id"]) + assert.Equal(t, "running", execChunks[0]["status"]) + + lastExec := execChunks[len(execChunks)-1] + assert.Equal(t, "tool_123", lastExec["tool_id"]) + assert.Equal(t, "completed", lastExec["status"]) + assert.Equal(t, "Bash", lastExec["tool"], "tool_result should carry tool name") +} + +// TestParser_ToolIndependentMessages verifies that each tool call gets its own +// message_start/message_end pair, and tool_result reuses the same message_id. +func TestParser_ToolIndependentMessages(t *testing.T) { + lines := []string{ + // Tool 1: Write + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{"type": "tool_use", "name": "Write", "id": "t_write"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"file_path":"server.js"}`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + // Tool 2: Bash + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{"type": "tool_use", "name": "Bash", "id": "t_bash"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + // Results + jsonLine(map[string]any{ + "type": "user", + "message": map[string]any{ + "content": []any{ + map[string]any{"type": "tool_result", "tool_use_id": "t_write", "content": "ok"}, + map[string]any{"type": "tool_result", "tool_use_id": "t_bash", "content": "done"}, + }, + }, + }), + jsonLine(map[string]any{"type": "result", "num_turns": 1}), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + + // Collect all message_start IDs and their order + var msgStarts []string + for _, c := range chunks { + if c.Type == message.ChunkMessageStart { + msgStarts = append(msgStarts, extractMessageID(c.Data)) + } + } + + // Should have 4 message_starts: Write(running), Bash(running), Write(result), Bash(result) + require.Equal(t, 4, len(msgStarts), "should have 4 message_start events") + + writeMsgID := msgStarts[0] + bashMsgID := msgStarts[1] + assert.NotEqual(t, writeMsgID, bashMsgID, "Write and Bash should have different message_ids") + + // tool_result should reuse the original message_id + assert.Equal(t, writeMsgID, msgStarts[2], "Write tool_result should reuse Write message_id") + assert.Equal(t, bashMsgID, msgStarts[3], "Bash tool_result should reuse Bash message_id") + + // Count message_end events (should match message_start) + endCount := 0 + for _, c := range chunks { + if c.Type == message.ChunkMessageEnd { + endCount++ + } + } + assert.Equal(t, 4, endCount, "each message_start should have matching message_end") +} + +// TestParser_ToolSummaryExtraction verifies that summary is extracted from tool input. +func TestParser_ToolSummaryExtraction(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{"type": "tool_use", "name": "Bash", "id": "t1"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"command":"ls -la /workspace"}`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + jsonLine(map[string]any{"type": "result", "num_turns": 1}), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + + // Look for summary in the last execute chunk before message_end + var summaryFound bool + for _, c := range chunks { + if c.Type == message.ChunkExecute { + var data map[string]any + json.Unmarshal(c.Data, &data) + if s, ok := data["summary"].(string); ok && s != "" { + summaryFound = true + assert.Equal(t, "ls -la /workspace", s) + } + } + } + assert.True(t, summaryFound, "should emit summary from tool input") +} + +func TestParser_UsageMetadata(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "assistant", + "message": map[string]any{ + "usage": map[string]any{ + "input_tokens": 100, + "output_tokens": 50, + }, + "stop_reason": "end_turn", + "content": []any{}, + }, + }), + jsonLine(map[string]any{"type": "result", "num_turns": 1}), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + + var usageMeta map[string]any + for _, c := range chunks { + if c.Type == message.ChunkMetadata { + var meta map[string]any + json.Unmarshal(c.Data, &meta) + if u, ok := meta["usage"]; ok { + usageMeta, _ = u.(map[string]any) + } + } + } + require.NotNil(t, usageMeta, "should emit usage metadata") + assert.Equal(t, float64(100), usageMeta["input_tokens"]) + assert.Equal(t, float64(50), usageMeta["output_tokens"]) +} + +func TestParser_ResultSummary(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "result", + "total_cost_usd": 0.05, + "duration_ms": 5000, + "num_turns": 3, + "usage": map[string]any{"input_tokens": 500, "output_tokens": 200}, + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + assert.True(t, p.completed) + + var summary map[string]any + for _, c := range chunks { + if c.Type == message.ChunkMetadata { + var meta map[string]any + json.Unmarshal(c.Data, &meta) + if s, ok := meta["result_summary"]; ok { + summary, _ = s.(map[string]any) + } + } + } + require.NotNil(t, summary, "should emit result_summary") + assert.Equal(t, float64(0.05), summary["total_cost_usd"]) + assert.Equal(t, float64(5000), summary["duration_ms"]) + assert.Equal(t, float64(3), summary["num_turns"]) +} + +func TestParser_ErrorMessage(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "error", + "error": map[string]any{"message": "rate limit exceeded"}, + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.Error(t, err) + assert.Contains(t, err.Error(), "rate limit exceeded") + assert.False(t, p.completed) + + hasError := false + for _, c := range chunks { + if c.Type == message.ChunkError { + hasError = true + assert.Contains(t, string(c.Data), "rate limit exceeded") + } + } + assert.True(t, hasError, "should emit error chunk") +} + +func TestParser_ResultIsError(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "result", + "is_error": true, + "result": "authentication failed", + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.Error(t, err) + assert.Contains(t, err.Error(), "authentication failed") +} + +func TestParser_ContextCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + r, w := io.Pipe() + + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + p := newStreamParser(nil) + err := p.parse(ctx, r) + _ = w.Close() + + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestParser_HandlerStopsStream(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "text_delta", "text": "first"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "text_delta", "text": "second"}, + }, + }), + } + + handler, chunks := stoppingHandler(2) + p := newStreamParser(handler) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + + assert.LessOrEqual(t, len(*chunks), 3, "should stop early") +} + +func TestParser_EmptyAndInvalidLines(t *testing.T) { + lines := []string{ + "", + "not json at all", + " ", + `{"type": "result", "num_turns": 1}`, + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + assert.True(t, p.completed, "should complete despite invalid lines") +} + +func TestParser_MultiTurnConversation(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{"type": "tool_use", "name": "Read", "id": "t1"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + jsonLine(map[string]any{ + "type": "assistant", + "message": map[string]any{ + "stop_reason": "tool_use", + "content": []any{ + map[string]any{"type": "tool_use", "name": "Read", "id": "t1", "input": map[string]any{"path": "/tmp"}}, + }, + }, + }), + jsonLine(map[string]any{ + "type": "user", + "message": map[string]any{ + "content": []any{ + map[string]any{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + }, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "text_delta", "text": "Done reading"}, + }, + }), + jsonLine(map[string]any{ + "type": "assistant", + "message": map[string]any{ + "stop_reason": "end_turn", + "content": []any{map[string]any{"type": "text", "text": "Done reading"}}, + }, + }), + jsonLine(map[string]any{"type": "result", "num_turns": 2}), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + assert.True(t, p.completed) + + startCount := 0 + endCount := 0 + for _, c := range chunks { + switch c.Type { + case message.ChunkMessageStart: + startCount++ + case message.ChunkMessageEnd: + endCount++ + } + } + // streaming tool_use(start) + streaming tool_use(stop) + assistant tool_use + tool_result + text = 5 starts + assert.GreaterOrEqual(t, startCount, 3, "should have multiple message starts for multi-turn") + assert.Equal(t, startCount, endCount, "each message_start should have matching message_end") +} + +func TestParser_ErrorStringFormat(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "error", + "error": "simple string error", + }), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.Error(t, err) + assert.Contains(t, err.Error(), "simple string error") +} + +func TestParser_InputDeltaAccumulation(t *testing.T) { + lines := []string{ + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_start", + "content_block": map[string]any{"type": "tool_use", "name": "Bash", "id": "t1"}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"com`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{ + "type": "content_block_delta", + "delta": map[string]any{"type": "input_json_delta", "partial_json": `mand":"ls"}`}, + }, + }), + jsonLine(map[string]any{ + "type": "stream_event", + "event": map[string]any{"type": "content_block_stop"}, + }), + jsonLine(map[string]any{"type": "result", "num_turns": 1}), + } + + var chunks []chunkRecord + p := newStreamParser(recordingHandler(&chunks)) + err := p.parse(context.Background(), pipeWithLines(lines...)) + require.NoError(t, err) + + // The last input_delta chunk should contain the full accumulated input + var lastDelta string + for _, c := range chunks { + if c.Type == message.ChunkExecute { + var data map[string]any + json.Unmarshal(c.Data, &data) + if d, ok := data["input_delta"].(string); ok { + lastDelta = d + } + } + } + assert.Equal(t, `{"command":"ls"}`, lastDelta, "input_delta should accumulate all fragments") +} diff --git a/agent/sandbox/v2/claude/plat_linuxos.go b/agent/sandbox/v2/claude/plat_linuxos.go new file mode 100644 index 00000000..e9c6c3b1 --- /dev/null +++ b/agent/sandbox/v2/claude/plat_linuxos.go @@ -0,0 +1,39 @@ +package claude + +import "fmt" + +// linuxPlatform overrides Linux container-specific behavior. +type linuxPlatform struct { + posixBase + hasDisplay bool // true when the container has a DISPLAY env var (Desktop/VNC) + sysHome string // original system HOME (e.g., /root) before redirection +} + +// XauthoritySetup handles X11 authentication cookie for Desktop containers. +// +// Desktop containers (VNC/noVNC): X Server creates .Xauthority under the +// system home (e.g., /root), but we redirect HOME to the workspace. GUI +// tools look for $HOME/.Xauthority, so we must copy the real cookie. +// +// Headless containers: no X Server, nothing to do. +func (p *linuxPlatform) XauthoritySetup(workDir string) string { + if !p.hasDisplay || p.sysHome == "" { + return "" + } + src := p.sysHome + "/.Xauthority" + dst := workDir + "/.Xauthority" + return fmt.Sprintf("[ -f %q ] && cp %q %q 2>/dev/null\n", src, src, dst) +} + +func (p *linuxPlatform) EnvPromptNote() string { + if p.hasDisplay { + return ` +- **Desktop Environment**: You have access to a Linux desktop via VNC (GUI applications, browsers, etc.) +- **Important**: When you launch GUI applications, do NOT close them unless explicitly asked` + } + return "" +} + +func (p *linuxPlatform) BuildScript(in scriptInput) (string, []byte) { + return p.buildBashScript(in, p.XauthoritySetup(in.workDir)), nil +} diff --git a/agent/sandbox/v2/claude/plat_macos.go b/agent/sandbox/v2/claude/plat_macos.go new file mode 100644 index 00000000..7407a773 --- /dev/null +++ b/agent/sandbox/v2/claude/plat_macos.go @@ -0,0 +1,17 @@ +package claude + +// darwinPlatform overrides macOS-specific behavior. +// Most methods are inherited from posixBase. +type darwinPlatform struct { + posixBase +} + +func (p *darwinPlatform) EnvPromptNote() string { + return ` +- **Desktop Environment**: You have access to the macOS desktop (GUI applications, browsers, Finder, etc.) +- **Important**: When you launch GUI applications, do NOT close them unless explicitly asked` +} + +func (p *darwinPlatform) BuildScript(in scriptInput) (string, []byte) { + return p.buildBashScript(in, ""), nil +} diff --git a/agent/sandbox/v2/claude/plat_win.go b/agent/sandbox/v2/claude/plat_win.go new file mode 100644 index 00000000..79cdc4f5 --- /dev/null +++ b/agent/sandbox/v2/claude/plat_win.go @@ -0,0 +1,114 @@ +package claude + +import ( + "fmt" + "strings" +) + +// windowsPlatform is a standalone implementation for Windows — no posixBase. +type windowsPlatform struct { + workDir string + shell string + tempDir string +} + +func newWindowsPlatform(workDir, shell, tempDir string) *windowsPlatform { + if shell == "" { + shell = "pwsh" + } + if tempDir == "" { + tempDir = workDir + `\.tmp` + } + return &windowsPlatform{workDir: workDir, shell: shell, tempDir: tempDir} +} + +func (w *windowsPlatform) OS() string { return "windows" } +func (w *windowsPlatform) Shell() string { return w.shell } +func (w *windowsPlatform) RootDir() string { return `C:\` } +func (w *windowsPlatform) ConfigDir() string { return `.claude` } +func (w *windowsPlatform) XauthoritySetup(_ string) string { return "" } + +func (w *windowsPlatform) PathJoin(parts ...string) string { + return strings.Join(parts, `\`) +} + +// HomeEnv sets all four HOME-related variables to prevent Windows path +// resolution issues. The HOME variable is critical: without it, tools like +// Git for Windows inherit the host HOME, causing ~ to expand to the wrong +// directory. See: anthropics/claude-code#13138 +func (w *windowsPlatform) HomeEnv(workDir string) map[string]string { + env := map[string]string{ + "HOME": workDir, + "USERPROFILE": workDir, + } + if len(workDir) >= 2 && workDir[1] == ':' { + env["HOMEDRIVE"] = workDir[:2] + env["HOMEPATH"] = workDir[2:] + } + return env +} + +func (w *windowsPlatform) EnvPromptNote() string { + return ` +- **Desktop Environment**: You have full access to the Windows desktop (GUI applications, browsers, etc.) +- **Important**: When you launch GUI applications (browsers, editors, etc.), do NOT close them unless explicitly asked — the user expects them to remain open` +} + +func (w *windowsPlatform) ShellCmd(script string) []string { + shell := strings.ToLower(w.shell) + switch shell { + case "pwsh": + return []string{"pwsh", "-NoProfile", "-Command", script} + case "powershell": + return []string{"powershell", "-NoProfile", "-Command", script} + case "cmd.exe", "cmd": + return []string{"cmd.exe", "/C", script} + default: + return []string{"pwsh", "-NoProfile", "-Command", script} + } +} + +func (w *windowsPlatform) KillCmd(pattern string) []string { + script := fmt.Sprintf( + "Get-Process -ErrorAction SilentlyContinue | Where-Object {$_.ProcessName -like '*%s*'} | ForEach-Object { taskkill /F /T /PID $_.Id 2>$null }; "+ + "Get-Process -ErrorAction SilentlyContinue | Where-Object {$_.ProcessName -like '*%s*'} | Stop-Process -Force -ErrorAction SilentlyContinue", + pattern, pattern) + return w.ShellCmd(script) +} + +func (w *windowsPlatform) ListDirCmd(dir string) []string { + return w.ShellCmd(fmt.Sprintf("Get-ChildItem -Name '%s'", dir)) +} + +func (w *windowsPlatform) BuildScript(in scriptInput) (string, []byte) { + var b strings.Builder + noBOM := "(New-Object System.Text.UTF8Encoding $false)" + + b.WriteString("[Console]::InputEncoding = [System.Text.Encoding]::UTF8\n") + b.WriteString("[Console]::OutputEncoding = [System.Text.Encoding]::UTF8\n") + b.WriteString("$OutputEncoding = [System.Text.Encoding]::UTF8\n") + + b.WriteString("foreach ($d in (Get-ChildItem 'C:\\Users' -Directory -ErrorAction SilentlyContinue)) {\n") + b.WriteString(" $p = Join-Path $d.FullName '.local\\bin'\n") + b.WriteString(" if (Test-Path (Join-Path $p 'claude.exe')) { $env:PATH = \"$p;$env:PATH\"; break }\n") + b.WriteString("}\n") + b.WriteString("if ($env:APPDATA) { $env:PATH = \"$env:APPDATA\\npm;$env:PATH\" }\n") + + yaoDir := w.PathJoin(in.workDir, ".yao") + b.WriteString(fmt.Sprintf("if (!(Test-Path '%s')) { New-Item -ItemType Directory -Path '%s' -Force | Out-Null }\n", yaoDir, yaoDir)) + + if in.systemPrompt != "" { + promptDir := in.promptFile[:strings.LastIndex(in.promptFile, `\`)] + b.WriteString(fmt.Sprintf("if (!(Test-Path '%s')) { New-Item -ItemType Directory -Path '%s' -Force | Out-Null }\n", promptDir, promptDir)) + escaped := strings.ReplaceAll(in.systemPrompt, "'", "''") + b.WriteString(fmt.Sprintf("[IO.File]::WriteAllText('%s', @'\n%s\n'@, %s)\n", in.promptFile, escaped, noBOM)) + in.args = append(in.args, "--append-system-prompt-file", in.promptFile) + } + + b.WriteString("claude -p") + for _, arg := range in.args { + b.WriteString(fmt.Sprintf(" '%s'", strings.ReplaceAll(arg, "'", "''"))) + } + + return b.String(), []byte(in.inputJSONL + "\n") +} diff --git a/agent/sandbox/v2/claude/platform.go b/agent/sandbox/v2/claude/platform.go new file mode 100644 index 00000000..d2b60db9 --- /dev/null +++ b/agent/sandbox/v2/claude/platform.go @@ -0,0 +1,136 @@ +package claude + +import ( + "fmt" + "path" + "strings" + + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// platform encapsulates all OS-dependent behaviors for the target environment. +// Upper-layer business code uses this interface exclusively, with zero platform +// branching (no if isWindows() checks). +type platform interface { + OS() string + Shell() string + HomeEnv(workDir string) map[string]string + EnvPromptNote() string + PathJoin(parts ...string) string + RootDir() string + ShellCmd(script string) []string + KillCmd(pattern string) []string + ListDirCmd(dir string) []string + ConfigDir() string + XauthoritySetup(workDir string) string + BuildScript(input scriptInput) (script string, stdin []byte) +} + +type scriptInput struct { + args []string + systemPrompt string + inputJSONL string + workDir string + promptFile string +} + +// posixBase provides shared POSIX-compatible implementation (~80% of methods) +// for macOS and Linux. darwinPlatform and linuxPlatform embed this. +type posixBase struct { + os string + workDir string + shell string + tempDir string + userHome string +} + +func (b *posixBase) OS() string { return b.os } +func (b *posixBase) Shell() string { return b.shell } +func (b *posixBase) PathJoin(parts ...string) string { return path.Join(parts...) } +func (b *posixBase) RootDir() string { return "/" } +func (b *posixBase) ConfigDir() string { return ".config/claude" } +func (b *posixBase) XauthoritySetup(_ string) string { return "" } + +func (b *posixBase) HomeEnv(workDir string) map[string]string { + return map[string]string{"HOME": workDir} +} + +func (b *posixBase) ShellCmd(script string) []string { + return []string{"bash", "-c", script} +} + +func (b *posixBase) KillCmd(pattern string) []string { + return []string{"sh", "-c", fmt.Sprintf("pkill -f '%s' || true", pattern)} +} + +func (b *posixBase) ListDirCmd(dir string) []string { + return []string{"ls", dir} +} + +// buildBashScript is the shared bash script builder for macOS/Linux. +func (b *posixBase) buildBashScript(in scriptInput, xauthCmd string) string { + var s strings.Builder + + if xauthCmd != "" { + s.WriteString(xauthCmd) + } + + if in.systemPrompt != "" { + s.WriteString(fmt.Sprintf("mkdir -p \"$(dirname %q)\"\n", in.promptFile)) + s.WriteString(fmt.Sprintf("cat << 'PROMPTEOF' > %s\n", in.promptFile)) + s.WriteString(in.systemPrompt) + s.WriteString("\nPROMPTEOF\n") + in.args = append(in.args, "--append-system-prompt-file", in.promptFile) + } + + s.WriteString("cat << 'INPUTEOF' | claude -p") + for _, arg := range in.args { + s.WriteString(fmt.Sprintf(" %q", arg)) + } + s.WriteString("\n") + s.WriteString(in.inputJSONL) + s.WriteString("\nINPUTEOF") + + return s.String() +} + +// resolvePlatform creates the appropriate platform implementation based on +// the Computer's reported OS. +func resolvePlatform(computer infra.Computer) platform { + sys := computer.ComputerInfo().System + osName := strings.ToLower(sys.OS) + workDir := computer.GetWorkDir() + shell := sys.Shell + tempDir := sys.TempDir + + base := posixBase{ + os: osName, workDir: workDir, + shell: shell, tempDir: tempDir, + } + if base.shell == "" { + base.shell = "bash" + } + if base.tempDir == "" { + base.tempDir = path.Join(workDir, ".tmp") + } + + // DISPLAY and system HOME are not yet available via SystemInfo. + // For Desktop Linux containers (VNC/noVNC), these will be populated + // once infra reports environment variables. Until then, the Xauthority + // copy will be a no-op — same as the old code's behavior. + hasDisplay := false + sysHome := "" + + switch osName { + case "darwin": + return &darwinPlatform{posixBase: base} + case "windows": + return newWindowsPlatform(workDir, shell, tempDir) + default: + return &linuxPlatform{ + posixBase: base, + hasDisplay: hasDisplay, + sysHome: sysHome, + } + } +} diff --git a/agent/sandbox/v2/claude/platform_test.go b/agent/sandbox/v2/claude/platform_test.go new file mode 100644 index 00000000..f0b34253 --- /dev/null +++ b/agent/sandbox/v2/claude/platform_test.go @@ -0,0 +1,304 @@ +package claude + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestPosixBase(os string) posixBase { + return posixBase{ + os: os, + workDir: "/workspace", + shell: "bash", + tempDir: "/tmp", + } +} + +func TestPosixBase_Accessors(t *testing.T) { + b := newTestPosixBase("linux") + assert.Equal(t, "linux", b.OS()) + assert.Equal(t, "bash", b.Shell()) + assert.Equal(t, "/", b.RootDir()) + assert.Equal(t, ".config/claude", b.ConfigDir()) +} + +func TestPosixBase_PathJoin(t *testing.T) { + b := newTestPosixBase("linux") + assert.Equal(t, "/workspace/.yao/config", b.PathJoin("/workspace", ".yao", "config")) + assert.Equal(t, "a/b/c", b.PathJoin("a", "b", "c")) +} + +func TestPosixBase_HomeEnv(t *testing.T) { + b := newTestPosixBase("linux") + env := b.HomeEnv("/workspace/data") + assert.Equal(t, "/workspace/data", env["HOME"]) + assert.Len(t, env, 1) +} + +func TestPosixBase_ShellCmd(t *testing.T) { + b := newTestPosixBase("linux") + cmd := b.ShellCmd("echo hello") + require.Len(t, cmd, 3) + assert.Equal(t, "bash", cmd[0]) + assert.Equal(t, "-c", cmd[1]) + assert.Equal(t, "echo hello", cmd[2]) +} + +func TestPosixBase_KillCmd(t *testing.T) { + b := newTestPosixBase("linux") + cmd := b.KillCmd("claude") + require.Len(t, cmd, 3) + assert.Equal(t, "sh", cmd[0]) + assert.Contains(t, cmd[2], "pkill") + assert.Contains(t, cmd[2], "claude") +} + +func TestPosixBase_ListDirCmd(t *testing.T) { + b := newTestPosixBase("linux") + cmd := b.ListDirCmd("/workspace/.claude") + require.Len(t, cmd, 2) + assert.Equal(t, "ls", cmd[0]) + assert.Equal(t, "/workspace/.claude", cmd[1]) +} + +func TestPosixBase_XauthoritySetup(t *testing.T) { + b := newTestPosixBase("linux") + assert.Empty(t, b.XauthoritySetup("/workspace")) +} + +func TestPosixBase_BuildBashScript_NoPrompt(t *testing.T) { + b := newTestPosixBase("linux") + in := scriptInput{ + args: []string{"--verbose", "--output-format", "stream-json"}, + inputJSONL: `{"type":"user","message":{"role":"user","content":"hello"}}`, + workDir: "/workspace", + promptFile: "/workspace/.yao/.system-prompt.txt", + } + script := b.buildBashScript(in, "") + assert.Contains(t, script, "cat << 'INPUTEOF' | claude -p") + assert.Contains(t, script, "--verbose") + assert.Contains(t, script, "INPUTEOF") + assert.NotContains(t, script, "PROMPTEOF") +} + +func TestPosixBase_BuildBashScript_WithPrompt(t *testing.T) { + b := newTestPosixBase("linux") + in := scriptInput{ + args: []string{"--verbose"}, + systemPrompt: "You are a helpful assistant.", + inputJSONL: `{"type":"user"}`, + workDir: "/workspace", + promptFile: "/workspace/.yao/assistants/test-id/system-prompt.txt", + } + script := b.buildBashScript(in, "") + assert.Contains(t, script, "mkdir -p") + assert.Contains(t, script, "PROMPTEOF") + assert.Contains(t, script, "You are a helpful assistant.") + assert.Contains(t, script, "--append-system-prompt-file") +} + +func TestPosixBase_BuildBashScript_WithXauth(t *testing.T) { + b := newTestPosixBase("linux") + in := scriptInput{ + args: []string{"--verbose"}, + inputJSONL: `{"type":"user"}`, + workDir: "/workspace", + promptFile: "/workspace/.yao/.system-prompt.txt", + } + xauth := "cp /root/.Xauthority /workspace/.Xauthority\n" + script := b.buildBashScript(in, xauth) + assert.True(t, strings.HasPrefix(script, "cp /root/.Xauthority"), + "script should start with xauth command") +} + +// --- Darwin --- + +func TestDarwin_EnvPromptNote(t *testing.T) { + p := &darwinPlatform{posixBase: newTestPosixBase("darwin")} + note := p.EnvPromptNote() + assert.Contains(t, note, "macOS desktop") + assert.Contains(t, note, "GUI applications") +} + +func TestDarwin_BuildScript(t *testing.T) { + p := &darwinPlatform{posixBase: newTestPosixBase("darwin")} + script, stdin := p.BuildScript(scriptInput{ + args: []string{"--verbose"}, + inputJSONL: `{"type":"user"}`, + workDir: "/workspace", + promptFile: "/workspace/.yao/.system-prompt.txt", + }) + assert.Contains(t, script, "claude -p") + assert.Nil(t, stdin) +} + +// --- Linux --- + +func TestLinux_XauthoritySetup_Headless(t *testing.T) { + p := &linuxPlatform{ + posixBase: newTestPosixBase("linux"), + hasDisplay: false, + sysHome: "/root", + } + assert.Empty(t, p.XauthoritySetup("/workspace")) +} + +func TestLinux_XauthoritySetup_Desktop(t *testing.T) { + p := &linuxPlatform{ + posixBase: newTestPosixBase("linux"), + hasDisplay: true, + sysHome: "/root", + } + cmd := p.XauthoritySetup("/workspace") + assert.Contains(t, cmd, "/root/.Xauthority") + assert.Contains(t, cmd, "/workspace/.Xauthority") + assert.Contains(t, cmd, "cp") +} + +func TestLinux_XauthoritySetup_NoSysHome(t *testing.T) { + p := &linuxPlatform{ + posixBase: newTestPosixBase("linux"), + hasDisplay: true, + sysHome: "", + } + assert.Empty(t, p.XauthoritySetup("/workspace")) +} + +func TestLinux_EnvPromptNote_Desktop(t *testing.T) { + p := &linuxPlatform{posixBase: newTestPosixBase("linux"), hasDisplay: true} + note := p.EnvPromptNote() + assert.Contains(t, note, "VNC") +} + +func TestLinux_EnvPromptNote_Headless(t *testing.T) { + p := &linuxPlatform{posixBase: newTestPosixBase("linux"), hasDisplay: false} + assert.Empty(t, p.EnvPromptNote()) +} + +func TestLinux_BuildScript_Desktop(t *testing.T) { + p := &linuxPlatform{ + posixBase: newTestPosixBase("linux"), + hasDisplay: true, + sysHome: "/root", + } + script, stdin := p.BuildScript(scriptInput{ + args: []string{"--verbose"}, + inputJSONL: `{"type":"user"}`, + workDir: "/workspace", + promptFile: "/workspace/.yao/.system-prompt.txt", + }) + assert.Contains(t, script, ".Xauthority") + assert.Nil(t, stdin) +} + +// --- Windows --- + +func TestWindows_NewDefaults(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "", "") + assert.Equal(t, "pwsh", w.Shell()) + assert.Equal(t, `C:\workspace\.tmp`, w.tempDir) +} + +func TestWindows_Accessors(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "pwsh", `C:\temp`) + assert.Equal(t, "windows", w.OS()) + assert.Equal(t, `C:\`, w.RootDir()) + assert.Equal(t, `.claude`, w.ConfigDir()) + assert.Empty(t, w.XauthoritySetup(`C:\workspace`)) +} + +func TestWindows_PathJoin(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "pwsh", "") + assert.Equal(t, `C:\workspace\.yao\config`, w.PathJoin(`C:\workspace`, ".yao", "config")) +} + +func TestWindows_HomeEnv(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "pwsh", "") + env := w.HomeEnv(`C:\workspace`) + assert.Equal(t, `C:\workspace`, env["HOME"]) + assert.Equal(t, `C:\workspace`, env["USERPROFILE"]) + assert.Equal(t, `C:`, env["HOMEDRIVE"]) + assert.Equal(t, `\workspace`, env["HOMEPATH"]) +} + +func TestWindows_HomeEnv_ShortPath(t *testing.T) { + w := newWindowsPlatform("X", "pwsh", "") + env := w.HomeEnv("X") + assert.Equal(t, "X", env["HOME"]) + assert.Equal(t, "X", env["USERPROFILE"]) + _, hasDrive := env["HOMEDRIVE"] + assert.False(t, hasDrive, "should not set HOMEDRIVE for paths without drive letter") +} + +func TestWindows_ShellCmd_Pwsh(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "pwsh", "") + cmd := w.ShellCmd("echo hello") + assert.Equal(t, []string{"pwsh", "-NoProfile", "-Command", "echo hello"}, cmd) +} + +func TestWindows_ShellCmd_Powershell(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "powershell", "") + cmd := w.ShellCmd("echo hello") + assert.Equal(t, "powershell", cmd[0]) +} + +func TestWindows_ShellCmd_Cmd(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "cmd.exe", "") + cmd := w.ShellCmd("echo hello") + assert.Equal(t, []string{"cmd.exe", "/C", "echo hello"}, cmd) +} + +func TestWindows_KillCmd(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "pwsh", "") + cmd := w.KillCmd("claude") + require.Len(t, cmd, 4) + assert.Equal(t, "pwsh", cmd[0]) + assert.Contains(t, cmd[3], "claude") + assert.Contains(t, cmd[3], "taskkill") +} + +func TestWindows_ListDirCmd(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "pwsh", "") + cmd := w.ListDirCmd(`C:\ws\.claude`) + require.Len(t, cmd, 4) + assert.Contains(t, cmd[3], "Get-ChildItem") +} + +func TestWindows_BuildScript_NoPrompt(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "pwsh", `C:\temp`) + script, stdin := w.BuildScript(scriptInput{ + args: []string{"--verbose"}, + inputJSONL: `{"type":"user"}`, + workDir: `C:\workspace`, + promptFile: `C:\workspace\.yao\.system-prompt.txt`, + }) + assert.Contains(t, script, "UTF8") + assert.Contains(t, script, "claude -p") + assert.Contains(t, script, "'--verbose'") + require.NotNil(t, stdin) + assert.Contains(t, string(stdin), `{"type":"user"}`) +} + +func TestWindows_BuildScript_WithPrompt(t *testing.T) { + w := newWindowsPlatform(`C:\workspace`, "pwsh", `C:\temp`) + script, stdin := w.BuildScript(scriptInput{ + args: []string{"--verbose"}, + systemPrompt: "You are helpful", + inputJSONL: `{"type":"user"}`, + workDir: `C:\workspace`, + promptFile: `C:\workspace\.yao\assistants\test\system-prompt.txt`, + }) + assert.Contains(t, script, "WriteAllText") + assert.Contains(t, script, "You are helpful") + assert.Contains(t, script, "--append-system-prompt-file") + require.NotNil(t, stdin) +} + +func TestWindows_EnvPromptNote(t *testing.T) { + w := newWindowsPlatform(`C:\ws`, "pwsh", "") + note := w.EnvPromptNote() + assert.Contains(t, note, "Windows desktop") +} diff --git a/agent/sandbox/v2/claude/runner.go b/agent/sandbox/v2/claude/runner.go index 4c50f524..dbfe2bfe 100644 --- a/agent/sandbox/v2/claude/runner.go +++ b/agent/sandbox/v2/claude/runner.go @@ -2,32 +2,22 @@ package claude import ( "context" - "encoding/json" - "errors" "fmt" - "io" "os" - "path" - "strings" - "time" - "github.com/yaoapp/gou/connector" + agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/agent/sandbox/v2/types" infra "github.com/yaoapp/yao/sandbox/v2" ) -const defaultProxyPort = 3456 - // ClaudeRunner implements the Runner interface for Claude CLI (mode=cli). type ClaudeRunner struct { - mode string - hasMCP bool - mcpToolPattern string // e.g. "mcp__yao__*,mcp__github__*" - servicePort int - servicePath string - serviceProtocol string - streamCompleted bool // set when Stream received "result"; Cleanup skips kill + mode string + hasMCP bool + mcpToolPattern string + lastCompleted bool + logger *agentContext.RequestLogger } // New creates a new ClaudeRunner. @@ -44,13 +34,19 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e r.mode = "cli" } + assistantID := req.Config.ID + prefix := ".yao/assistants/" + assistantID + if assistantID == "" { + prefix = ".claude" + } + steps := append([]types.PrepareStep{}, req.Config.Prepare...) if req.SkillsDir != "" { ws := req.Computer.Workplace() if ws != nil { src := "local:///" + req.SkillsDir - dst := ".claude/skills" + dst := prefix + "/skills" if _, err := ws.Copy(src, dst); err != nil { fmt.Fprintf(os.Stderr, "[claude] warn: copy skills %s -> %s: %v\n", src, dst, err) } @@ -63,7 +59,7 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e mcpJSON := buildMCPConfig(req.MCPServers) steps = append(steps, types.PrepareStep{ Action: "file", - Path: ".claude/mcp.json", + Path: prefix + "/mcp.json", Content: mcpJSON, }) } @@ -84,11 +80,10 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han return fmt.Errorf("computer is nil") } - oe := resolveOSEnv(computer, req.Config) + p := resolvePlatform(computer) if req.ChatID != "" { - ws := computer.Workplace() - if ws != nil { + if ws := computer.Workplace(); ws != nil { processed, err := prepareAttachments(ctx, req.Messages, req.ChatID, ws) if err != nil { return fmt.Errorf("prepareAttachments: %w", err) @@ -97,332 +92,41 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han } } - isContinuation := hasExistingSession(ctx, computer, oe) + cmd := r.buildCommand(ctx, req, p) - cmd, env, stdin := r.buildCLICommand(req, oe, isContinuation) - - streamOpts := []infra.ExecOption{infra.WithWorkDir(oe.WorkDir), infra.WithEnv(env)} - if len(stdin) > 0 { - streamOpts = append(streamOpts, infra.WithStdin(stdin)) + r.logger = req.Logger + if r.logger == nil { + r.logger = agentContext.NoopLogger() } - fmt.Fprintf(os.Stderr, "[claude] Stream cmd=%v hasMCP=%v isContinuation=%v stdinLen=%d workDir=%q\n", cmd, r.hasMCP, isContinuation, len(stdin), oe.WorkDir) - - execStream, err := computer.Stream(ctx, cmd, streamOpts...) + sess, err := startSession(ctx, computer, p, cmd, r.logger) if err != nil { - return fmt.Errorf("computer.Stream: %w", err) + return err } - streamCtx, streamCancel := context.WithCancel(ctx) - defer streamCancel() - - // Kill claude processes only when the context is cancelled externally - // (upstream timeout, user interrupt) — NOT on normal return. - go func() { - <-streamCtx.Done() - if ctx.Err() == nil { - fmt.Fprintf(os.Stderr, "[claude] streamCtx done: normal return, skipping kill (ctx.Err=nil)\n") - return - } - fmt.Fprintf(os.Stderr, "[claude] streamCtx done: context cancelled externally (ctx.Err=%v), killing processes\n", ctx.Err()) - killCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - computer.Exec(killCtx, oe.killProcessCmd("claude")) - execStream.Cancel() - }() - - var stderrBuf strings.Builder - go func() { - buf := make([]byte, 4096) - for { - n, err := execStream.Stderr.Read(buf) - if n > 0 { - stderrBuf.Write(buf[:n]) - chunk := string(buf[:n]) - if strings.Contains(strings.ToLower(chunk), "error") { - streamCancel() - io.Copy(&stderrBuf, execStream.Stderr) - return - } - } - if err != nil { - return - } - } - }() - - parseErr := parseStreamJSON(streamCtx, execStream.Stdout, handler) - fmt.Fprintf(os.Stderr, "[claude] parseStreamJSON returned: %v\n", parseErr) - - // Received "result" — Claude finished normally. Return immediately. - if errors.Is(parseErr, errStreamCompleted) { - r.streamCompleted = true - fmt.Fprintf(os.Stderr, "[claude] stream completed normally, returning nil\n") - return nil - } - - // Parse failed or stream ended without "result" — wait for process. - fmt.Fprintf(os.Stderr, "[claude] stream did NOT complete normally, waiting for process exit...\n") - exitCode, waitErr := execStream.Wait() - stderrStr := strings.TrimSpace(stderrBuf.String()) - - if parseErr != nil { - if stderrStr != "" { - return fmt.Errorf("%w (stderr: %s)", parseErr, stderrStr) - } - return parseErr - } - if waitErr != nil { - if stderrStr != "" { - return fmt.Errorf("%w (stderr: %s)", waitErr, stderrStr) - } - return waitErr - } - if exitCode != 0 { - fmt.Fprintf(os.Stderr, "[claude] exit code=%d stderr=%q\n", exitCode, stderrStr) - if stderrStr != "" { - return fmt.Errorf("claude CLI exited with code %d: %s", exitCode, stderrStr) - } - return fmt.Errorf("claude CLI exited with code %d", exitCode) - } - return nil + completed, err := sess.runStream(handler) + r.lastCompleted = completed + return err } // Cleanup kills any remaining claude processes. If the stream completed -// normally (received "result"), child processes are preserved — the user -// may have asked Claude to launch a browser, server, etc. +// normally (received "result"), child processes are preserved. func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) error { if computer == nil { return nil } - if r.streamCompleted { - fmt.Fprintf(os.Stderr, "[claude] cleanup: stream completed normally, skipping process kill (child processes preserved)\n") + if r.lastCompleted { + if r.logger != nil { + r.logger.Info("cleanup: stream completed normally, preserving child processes") + } return nil } if r.mode != "service" { - oe := resolveOSEnv(computer, nil) - computer.Exec(ctx, oe.killProcessCmd("claude")) + p := resolvePlatform(computer) + computer.Exec(ctx, p.KillCmd("claude")) } return nil } - -// hasExistingSession checks if a Claude CLI session exists in the workspace. -func hasExistingSession(ctx context.Context, computer infra.Computer, oe *osEnv) bool { - sessionDir := oe.pathJoin(oe.WorkDir, ".claude", "projects") - result, err := computer.Exec(ctx, oe.listDirCmd(sessionDir)) - if err != nil || result.ExitCode != 0 { - return false - } - return strings.TrimSpace(result.Stdout) != "" -} - -// buildCLICommand constructs the Claude CLI command, environment variables, and optional stdin bytes. -func (r *ClaudeRunner) buildCLICommand(req *types.StreamRequest, oe *osEnv, isContinuation bool) ([]string, map[string]string, []byte) { - env := make(map[string]string) - - if oe.isWindows() { - env["USERPROFILE"] = oe.WorkDir - if len(oe.WorkDir) >= 2 && oe.WorkDir[1] == ':' { - env["HOMEDRIVE"] = oe.WorkDir[:2] - env["HOMEPATH"] = oe.WorkDir[2:] - } - } else { - env["HOME"] = oe.WorkDir - if oe.UserHome != "" { - env["XAUTHORITY"] = path.Join(oe.UserHome, ".Xauthority") - } - } - - if req.Connector != nil { - setting := req.Connector.Setting() - host, _ := setting["host"].(string) - key, _ := setting["key"].(string) - model, _ := setting["model"].(string) - - if req.Connector.Is(connector.ANTHROPIC) { - env["ANTHROPIC_BASE_URL"] = host - env["ANTHROPIC_API_KEY"] = key - } else { - env["ANTHROPIC_BASE_URL"] = fmt.Sprintf("http://127.0.0.1:%d", defaultProxyPort) - env["ANTHROPIC_API_KEY"] = "dummy" - } - - if model != "" { - env["ANTHROPIC_MODEL"] = model - env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model - env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model - env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model - env["CLAUDE_CODE_SUBAGENT_MODEL"] = model - } - - if thinking, ok := setting["thinking"].(map[string]interface{}); ok { - thinkType, _ := thinking["type"].(string) - switch thinkType { - case "disabled": - env["MAX_THINKING_TOKENS"] = "0" - case "enabled": - if budget, ok := thinking["budget_tokens"].(float64); ok && budget > 0 { - env["MAX_THINKING_TOKENS"] = fmt.Sprintf("%d", int(budget)) - } - } - } - } - - if req.Config != nil && len(req.Config.Secrets) > 0 { - for k, v := range req.Config.Secrets { - env[k] = v - } - } - - if req.Token != nil { - if req.Token.Token != "" { - env["YAO_TOKEN"] = req.Token.Token - } - if req.Token.RefreshToken != "" { - env["YAO_REFRESH_TOKEN"] = req.Token.RefreshToken - } - } - - var systemPrompt string - envPrompt := buildSandboxEnvPrompt(oe) - if !isContinuation && req.SystemPrompt != "" { - systemPrompt = req.SystemPrompt + "\n\n" + envPrompt - } else if !isContinuation { - systemPrompt = envPrompt - } - - var inputJSONL string - if isContinuation { - inputJSONL = buildLastUserMessageJSONL(req.Messages) - } else { - inputJSONL = buildFirstRequestJSONL(req.Messages) - } - - var args []string - - permMode := "" - if req.Config != nil && req.Config.Runner.Options != nil { - if v, ok := req.Config.Runner.Options["permission_mode"]; ok { - permMode = fmt.Sprintf("%v", v) - } - } - if permMode == "bypassPermissions" { - args = append(args, "--dangerously-skip-permissions") - args = append(args, "--permission-mode", permMode) - } - - args = append(args, "--input-format", "stream-json") - args = append(args, "--output-format", "stream-json") - args = append(args, "--include-partial-messages") - args = append(args, "--verbose") - - if isContinuation { - args = append(args, "--continue") - } - - if req.Config != nil && req.Config.Runner.Options != nil { - for key, val := range req.Config.Runner.Options { - if flag, ok := claudeArgWhitelist[key]; ok { - args = append(args, flag, fmt.Sprintf("%v", val)) - } - } - } - - if r.hasMCP { - mcpPath := oe.pathJoin(oe.WorkDir, ".claude", "mcp.json") - args = append(args, "--mcp-config", mcpPath) - if r.mcpToolPattern != "" { - args = append(args, "--allowedTools", r.mcpToolPattern) - } - } - - script, stdin := oe.buildCLIScript(args, systemPrompt, inputJSONL) - return oe.shellCmd(script), env, stdin -} - -// buildMCPConfig creates the .mcp.json for Claude CLI based on declared servers. -// Each server delegates to "tai mcp" which implements the standard MCP protocol -// over stdio and bridges to Yao gRPC with authentication. -// Connection is configured via env vars (YAO_GRPC_ADDR, YAO_TOKEN, etc.) -// injected by the sandbox infrastructure at container start. -func buildMCPConfig(servers []types.MCPServer) []byte { - mcpServers := make(map[string]any, len(servers)) - for _, s := range servers { - name := s.ServerID - if name == "" { - continue - } - mcpServers[name] = map[string]any{ - "command": "tai", - "args": []string{"mcp", name}, - } - } - if len(mcpServers) == 0 { - mcpServers["yao"] = map[string]any{ - "command": "tai", - "args": []string{"mcp"}, - } - } - config := map[string]any{"mcpServers": mcpServers} - data, _ := json.Marshal(config) - return data -} - -// buildMCPAllowedTools generates the --allowedTools pattern from server IDs. -func buildMCPAllowedTools(servers []types.MCPServer) string { - patterns := make([]string, 0, len(servers)) - for _, s := range servers { - if s.ServerID != "" { - patterns = append(patterns, fmt.Sprintf("mcp__%s__*", s.ServerID)) - } - } - if len(patterns) == 0 { - return "mcp__yao__*" - } - return strings.Join(patterns, ",") -} - -// buildSandboxEnvPrompt generates the sandbox environment prompt with system info and working directory. -func buildSandboxEnvPrompt(oe *osEnv) string { - workDir := oe.WorkDir - - osName := oe.OS - if osName == "" { - osName = "linux" - } - shell := oe.Shell - if shell == "" { - shell = "bash" - } - - shellNote := "" - if oe.isWindows() { - shellNote = ` -- **Desktop Environment**: You have full access to the Windows desktop (GUI applications, browsers, etc.) -- **Important**: When you launch GUI applications (browsers, editors, etc.), do NOT close them unless explicitly asked — the user expects them to remain open` - } - - return fmt.Sprintf(`## Sandbox Environment - -- **Operating System**: %[2]s -- **Shell**: %[3]s -- **Working Directory**: %[1]s -- **File Access**: You have full read/write access to %[1]s%[4]s - -## User Attachments - -User-uploaded files (images, documents, code files, etc.) are placed in %[1]s/.attachments/{chatID}/ -Each chat session has its own subdirectory to avoid conflicts. -When the user references an attached file, read it from this directory using the Read or Bash tool. -For image files, you can view them directly as Claude supports vision on local files. -`, workDir, osName, shell, shellNote) -} - -var claudeArgWhitelist = map[string]string{ - "max_turns": "--max-turns", - "disallowed_tools": "--disallowed-tools", - "allowed_tools": "--allowedTools", -} diff --git a/agent/sandbox/v2/claude/runner_test.go b/agent/sandbox/v2/claude/runner_test.go index 97afa979..71b64259 100644 --- a/agent/sandbox/v2/claude/runner_test.go +++ b/agent/sandbox/v2/claude/runner_test.go @@ -35,6 +35,14 @@ var cases = []e2eCase{ }, } +var toolCallCases = []e2eCase{ + { + ID: "tests.sandbox-v2.oneshot-cli", + Prompt: "Run the command 'echo refactor-ok' and tell me the output.", + Timeout: 3 * time.Minute, + }, +} + func TestSandboxV2_Claude_E2E(t *testing.T) { sandboxtestutils.Prepare(t) defer sandboxtestutils.Clean(t) @@ -277,3 +285,87 @@ func mapKeys(m map[string]interface{}) []string { } return keys } + +// TestSandboxV2_Claude_ToolCallE2E verifies that tool call execution emits +// "execute" messages and that usage/result_summary metadata is propagated. +func TestSandboxV2_Claude_ToolCallE2E(t *testing.T) { + sandboxtestutils.Prepare(t) + defer sandboxtestutils.Clean(t) + + require.NotNil(t, caller.AgentGetterFunc, "AgentGetterFunc should be registered after Prepare") + + for _, tc := range toolCallCases { + tc := tc + t.Run(tc.ID+"_tool_call", func(t *testing.T) { + agent, err := caller.AgentGetterFunc(tc.ID) + require.NoError(t, err, "should load assistant %s", tc.ID) + + timeout := tc.Timeout + if timeout == 0 { + timeout = 3 * time.Minute + } + + chatID := fmt.Sprintf("e2e-tool-%s-%d", tc.ID, time.Now().UnixMilli()) + ctx := agentcontext.New( + context.Background(), + &oauthtypes.AuthorizedInfo{ + TeamID: "test-team-e2e", + UserID: "test-user-e2e", + }, + chatID, + ) + + messages := []agentcontext.Message{ + {Role: "user", Content: tc.Prompt}, + } + + done := make(chan struct{}) + var resp *agentcontext.Response + var streamErr error + + go func() { + defer close(done) + resp, streamErr = agent.Stream(ctx, messages) + }() + + select { + case <-done: + case <-time.After(timeout): + t.Fatalf("timeout after %v", timeout) + } + + require.NoError(t, streamErr, "Stream should not return error") + require.NotNil(t, resp, "response should not be nil") + require.NotNil(t, resp.Completion, "completion should not be nil") + + contentStr, ok := resp.Completion.Content.(string) + require.True(t, ok, "Content should be a string, got %T", resp.Completion.Content) + t.Logf("Response (%d chars): %s", len(contentStr), contentStr) + + assert.Contains(t, strings.ToLower(contentStr), "refactor-ok", + "response should contain the command output") + + // ── Verify buffer has execute messages ── + require.NotNil(t, ctx.Buffer, "ctx.Buffer should not be nil") + msgs := ctx.Buffer.GetMessages() + t.Logf("buffer message count: %d", len(msgs)) + + var executeCount int + for _, m := range msgs { + t.Logf(" seq=%d role=%s type=%s streaming=%v props_keys=%v", + m.Sequence, m.Role, m.Type, m.IsStreaming, mapKeys(m.Props)) + if m.Type == "execute" { + executeCount++ + assert.NotNil(t, m.Props, "execute message should have props") + if m.Props != nil { + if toolName, ok := m.Props["tool"].(string); ok { + t.Logf(" execute tool=%s status=%v", toolName, m.Props["status"]) + } + } + } + } + assert.GreaterOrEqual(t, executeCount, 1, + "should have at least 1 execute message (tool call)") + }) + } +} diff --git a/agent/sandbox/v2/claude/session.go b/agent/sandbox/v2/claude/session.go new file mode 100644 index 00000000..bb032b35 --- /dev/null +++ b/agent/sandbox/v2/claude/session.go @@ -0,0 +1,166 @@ +package claude + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// session encapsulates a single Claude CLI execution lifecycle: +// process start, stderr collection, kill on cancel, Wait with timeout. +type session struct { + ctx context.Context + computer infra.Computer + plat platform + exec *infra.ExecStream + stderr strings.Builder + stderrMu sync.Mutex + logger *agentContext.RequestLogger +} + +func startSession(ctx context.Context, computer infra.Computer, p platform, cmd command, logger *agentContext.RequestLogger) (*session, error) { + opts := []infra.ExecOption{infra.WithWorkDir(cmd.workDir), infra.WithEnv(cmd.env)} + if len(cmd.stdin) > 0 { + opts = append(opts, infra.WithStdin(cmd.stdin)) + } + + logger.Info("claude session starting: cmd=%s workDir=%s platform=%s stdinLen=%d", + cmd.shell, cmd.workDir, p.OS(), len(cmd.stdin)) + + execStream, err := computer.Stream(ctx, cmd.shell, opts...) + if err != nil { + return nil, fmt.Errorf("computer.Stream: %w", err) + } + + return &session{ + ctx: ctx, + computer: computer, + plat: p, + exec: execStream, + logger: logger, + }, nil +} + +// runStream executes the main stream processing loop. +// Returns (completed, error) where completed=true means Claude CLI sent +// a "result" message and the stream finished normally. +func (s *session) runStream(handler message.StreamFunc) (completed bool, err error) { + s.collectStderr() + + cleanup := s.watchCancel() + defer cleanup() + + parser := newStreamParser(handler) + parseErr := parser.parse(s.ctx, s.exec.Stdout) + + if parser.completed { + s.logger.Info("claude stream completed normally") + return true, nil + } + + return false, s.waitForExit(parseErr) +} + +// collectStderr reads stderr in a background goroutine. +// Unlike the old code, this NEVER triggers stream cancellation. +// stderr is purely informational — logged and collected for error reporting. +func (s *session) collectStderr() { + go func() { + buf := make([]byte, 4096) + for { + n, err := s.exec.Stderr.Read(buf) + if n > 0 { + chunk := string(buf[:n]) + s.stderrMu.Lock() + s.stderr.WriteString(chunk) + s.stderrMu.Unlock() + s.logger.Debug("claude stderr: %s", chunk) + } + if err != nil { + return + } + } + }() +} + +// watchCancel monitors context cancellation and kills the Claude process. +// Returns a cleanup function that must be deferred. +func (s *session) watchCancel() func() { + done := make(chan struct{}) + go func() { + select { + case <-s.ctx.Done(): + s.logger.Info("context cancelled, killing claude: %v", s.ctx.Err()) + killCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + s.computer.Exec(killCtx, s.plat.KillCmd("claude")) + s.exec.Cancel() + case <-done: + } + }() + return func() { close(done) } +} + +// waitForExit waits for the Claude process to exit with timeout protection. +// This fixes the old code's issue where Wait() could block forever. +func (s *session) waitForExit(parseErr error) error { + s.logger.Info("claude stream did not complete normally, waiting for exit") + + type waitResult struct { + exitCode int + err error + } + ch := make(chan waitResult, 1) + go func() { + code, err := s.exec.Wait() + ch <- waitResult{code, err} + }() + + var exitCode int + var waitErr error + + select { + case wr := <-ch: + exitCode, waitErr = wr.exitCode, wr.err + case <-s.ctx.Done(): + select { + case wr := <-ch: + exitCode, waitErr = wr.exitCode, wr.err + case <-time.After(10 * time.Second): + s.exec.Cancel() + s.logger.Error("claude did not exit after kill, timeout") + return fmt.Errorf("claude did not exit after kill (timeout)") + } + } + + s.stderrMu.Lock() + stderrStr := strings.TrimSpace(s.stderr.String()) + s.stderrMu.Unlock() + + if parseErr != nil { + if stderrStr != "" { + return fmt.Errorf("%w (stderr: %s)", parseErr, stderrStr) + } + return parseErr + } + if waitErr != nil { + if stderrStr != "" { + return fmt.Errorf("%w (stderr: %s)", waitErr, stderrStr) + } + return waitErr + } + if exitCode != 0 { + s.logger.Warn("claude exited with non-zero code: exitCode=%d stderr=%s", exitCode, stderrStr) + if stderrStr != "" { + return fmt.Errorf("claude CLI exited with code %d: %s", exitCode, stderrStr) + } + return fmt.Errorf("claude CLI exited with code %d", exitCode) + } + return nil +} diff --git a/agent/sandbox/v2/stream.go b/agent/sandbox/v2/stream.go index 4816af04..c6615ceb 100644 --- a/agent/sandbox/v2/stream.go +++ b/agent/sandbox/v2/stream.go @@ -124,7 +124,7 @@ func ExecuteSandboxStream( loadingClosed := false wrappedHandler := func(chunkType message.StreamChunkType, data []byte) int { if !loadingClosed && req.LoadingMsgID != "" { - if chunkType == message.ChunkText || chunkType == message.ChunkToolCall || chunkType == message.ChunkMessageStart { + if chunkType == message.ChunkText || chunkType == message.ChunkToolCall || chunkType == message.ChunkExecute || chunkType == message.ChunkMessageStart { closeLoading(ctx, req.LoadingMsgID) loadingClosed = true } diff --git a/agent/sandbox/v2/types/runner.go b/agent/sandbox/v2/types/runner.go index 6268efbb..0166e4c7 100644 --- a/agent/sandbox/v2/types/runner.go +++ b/agent/sandbox/v2/types/runner.go @@ -53,5 +53,6 @@ type StreamRequest struct { Messages []agentContext.Message SystemPrompt string ChatID string - Token *SandboxToken // current user's sandbox token for MCP callbacks + Token *SandboxToken // current user's sandbox token for MCP callbacks + Logger *agentContext.RequestLogger // request-scoped logger propagated from agent context }