Merge pull request #1497 from trheyi/main

feat(stream): enhance process management and error handling

agent/sandbox/v2/claude/oscompat.go

@@ -75,9 +75,15 @@ func (e *osEnv) listDirCmd(dir string) []string {
 }
 
 // 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() {
-		script := fmt.Sprintf("Get-Process | Where-Object {$_.ProcessName -like '*%s*'} | Stop-Process -Force -ErrorAction SilentlyContinue", pattern)
+		// 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)}
@@ -146,6 +152,23 @@ func (e *osEnv) buildPowerShellScript(args []string, systemPrompt, inputJSONL, w
 	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))
 

agent/sandbox/v2/claude/parse.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -15,9 +16,28 @@ import (
 	"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")
+
 // parseStreamJSON reads stream-json lines from Claude CLI stdout and
 // pushes them through handler as standard StreamChunkType events.
-func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.StreamFunc) error {
+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.
+	doneParsing := make(chan struct{})
+	defer close(doneParsing)
+	go func() {
+		select {
+		case <-ctx.Done():
+			stdout.Close()
+		case <-doneParsing:
+		}
+	}()
+
 	scanner := bufio.NewScanner(stdout)
 	buf := make([]byte, 0, 64*1024)
 	scanner.Buffer(buf, 1024*1024)
@@ -283,6 +303,12 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
 					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
 
 		case "error":
 			var errMsg string
@@ -305,7 +331,17 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
 		}
 	}
 
-	return scanner.Err()
+	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.
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		return err
+	}
+	return nil
 }
 
 // buildFirstRequestJSONL builds JSONL with all messages for the first request.

agent/sandbox/v2/claude/runner.go

@@ -3,6 +3,7 @@ package claude
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -26,6 +27,7 @@ type ClaudeRunner struct {
 	servicePort     int
 	servicePath     string
 	serviceProtocol string
+	streamCompleted bool // set when Stream received "result"; Cleanup skips kill
 }
 
 // New creates a new ClaudeRunner.
@@ -104,7 +106,7 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
 		streamOpts = append(streamOpts, infra.WithStdin(stdin))
 	}
 
-	fmt.Fprintf(os.Stderr, "[claude] Stream cmd=%v hasMCP=%v isContinuation=%v stdinLen=%d\n", cmd, r.hasMCP, isContinuation, len(stdin))
+	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...)
 	if err != nil {
@@ -114,8 +116,15 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
 	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"))
@@ -143,7 +152,17 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
 	}()
 
 	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())
 
@@ -160,7 +179,7 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
 		return waitErr
 	}
 	if exitCode != 0 {
-		fmt.Fprintf(os.Stderr, "[claude] exit code=%d parseErr=%v waitErr=%v stderr=%q\n", exitCode, parseErr, waitErr, stderrStr)
+		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)
 		}
@@ -169,12 +188,19 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
 	return nil
 }
 
-// Cleanup kills any remaining claude processes.
+// 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.
 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")
+		return nil
+	}
+
 	if r.mode != "service" {
 		oe := resolveOSEnv(computer, nil)
 		computer.Exec(ctx, oe.killProcessCmd("claude"))
@@ -261,7 +287,7 @@ func (r *ClaudeRunner) buildCLICommand(req *types.StreamRequest, oe *osEnv, isCo
 	}
 
 	var systemPrompt string
-	envPrompt := buildSandboxEnvPrompt(oe.WorkDir)
+	envPrompt := buildSandboxEnvPrompt(oe)
 	if !isContinuation && req.SystemPrompt != "" {
 		systemPrompt = req.SystemPrompt + "\n\n" + envPrompt
 	} else if !isContinuation {
@@ -359,31 +385,32 @@ func buildMCPAllowedTools(servers []types.MCPServer) string {
 	return strings.Join(patterns, ",")
 }
 
-// buildSandboxEnvPrompt generates the sandbox environment prompt with the actual working directory.
+// buildSandboxEnvPrompt generates the sandbox environment prompt with system info and working directory.
-func buildSandboxEnvPrompt(workDir string) string {
+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
 
-You are running in a sandboxed environment with the following setup:
+- **Operating System**: %[2]s
-
+- **Shell**: %[3]s
 - **Working Directory**: %[1]s
-- **Project Structure**: If this is a new project, create a dedicated project folder (e.g., %[1]s/my-project/) and work inside it
+- **File Access**: You have full read/write access to %[1]s%[4]s
-- **File Access**: You have full read/write access to %[1]s
-- **Output Files**: Save all output files to the working directory
-
-When creating new projects:
-1. Create a project directory with a descriptive name
-2. Initialize the project structure inside that directory
-3. Keep all related files organized within the project folder
-
-## IMPORTANT: Restricted Tools
-
-The following tools are NOT available in this environment and you must NOT use them:
-- EnterPlanMode, ExitPlanMode (use regular text to explain plans instead)
-- Task, TaskOutput, TaskStop (complete tasks directly without delegation)
-- AskUserQuestion (make reasonable assumptions instead of asking)
-- Skill, ToolSearch (not supported)
-
-Focus on using the core tools: Bash, Read, Write, Edit, Glob, Grep, WebSearch, WebFetch.
 
 ## User Attachments
 
@@ -391,14 +418,7 @@ User-uploaded files (images, documents, code files, etc.) are placed in %[1]s/.a
 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)
-## GitHub CLI (gh) Usage
-
-When working with GitHub and a token is provided:
-1. First authenticate gh CLI using the token: echo "TOKEN" | gh auth login --with-token
-2. Then use gh commands normally (gh repo create, gh pr create, etc.)
-3. Do NOT use curl to call GitHub API directly - always prefer gh CLI
-`, workDir)
 }
 
 var claudeArgWhitelist = map[string]string{