opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1507 from trheyi/main

Browse source 
feat(claude): enhance session management and argument building
This commit is contained in:
Max 2026-03-26 09:01:52 +08:00 • committed by GitHub
commit e98e370fb2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 328 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.gitignore vendored
View file

@ -80,3 +80,5 @@ tai/docs/refactor-registration.md
agent/robot/ROBOT-WATCHER-IMPROVEMENT.md
agent/robot/ROBOT-IM-INTEGRATION-IMPROVEMENT.md
agent/robot/ROBOT-CACHE-IMPROVEMENT.md
sandbox/v2/PID-KILL-UPGRADE.md
sandbox/v2/*.md

2
agent/assistant/handlers/stream.go
View file

@ -117,7 +117,6 @@ func (s *streamState) handleMessageStart(data []byte) int {
startData.ThreadID = s.ctx.Stack.ID
}
// Initialize message state with the correct message ID
s.inGroup = true
s.currentGroupID = messageID
s.buffer = []byte{}
@ -381,7 +380,6 @@ func (s *streamState) handleMessageEnd(data []byte) int {
return 0
}
// Calculate duration
durationMs := time.Since(s.groupStartTime).Milliseconds()
// Use the tracked message type (thinking, text, tool_call, etc.)

55
agent/robot/executor/standard/agent.go
View file

@ -2,8 +2,10 @@ package standard
import (
"fmt"
"time"
"github.com/yaoapp/gou/text"
kunlog "github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/assistant"
agentcontext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/output/message"
@ -190,27 +192,31 @@ func (c *AgentCaller) Call(ctx *robottypes.Context, assistantID string, messages
Connector: c.Connector,
}
// Convert robot context to agent context
agentCtx := c.buildAgentContext(ctx)
defer agentCtx.Release() // IMPORTANT: Release agent context to prevent resource leaks
agentCtx := c.buildAgentContext(ctx, assistantID)
defer func() {
kunlog.Trace("[robot-agent] releasing context: assistantID=%s chatID=%s", assistantID, c.ChatID)
agentCtx.Release()
}()
callStart := time.Now()
kunlog.Trace("[robot-agent] Call started: assistantID=%s chatID=%s", assistantID, c.ChatID)
// Call assistant with streaming
response, err := ast.Stream(agentCtx, messages, opts)
if err != nil {
kunlog.Trace("[robot-agent] Call failed: assistantID=%s elapsed=%v err=%v", assistantID, time.Since(callStart).Round(time.Second), err)
return nil, fmt.Errorf("assistant call failed: %w", err)
}
// Build result
kunlog.Trace("[robot-agent] Call completed: assistantID=%s elapsed=%v", assistantID, time.Since(callStart).Round(time.Second))
result := &CallResult{
Response: response,
}
// Extract Next hook data
if response.Next != nil {
result.Next = response.Next
}
// Extract Content from Completion
if response.Completion != nil {
if content, ok := response.Completion.Content.(string); ok {
result.Content = content
@ -294,14 +300,23 @@ func (c *AgentCaller) CallStream(ctx *robottypes.Context, assistantID string, me
}
}
agentCtx := c.buildAgentContext(ctx)
defer agentCtx.Release()
agentCtx := c.buildAgentContext(ctx, assistantID)
defer func() {
kunlog.Trace("[robot-agent] releasing context (CallStream): assistantID=%s chatID=%s", assistantID, c.ChatID)
agentCtx.Release()
}()
callStart := time.Now()
kunlog.Trace("[robot-agent] CallStream started: assistantID=%s chatID=%s", assistantID, c.ChatID)
response, err := ast.Stream(agentCtx, messages, opts)
if err != nil {
kunlog.Trace("[robot-agent] CallStream failed: assistantID=%s elapsed=%v err=%v", assistantID, time.Since(callStart).Round(time.Second), err)
return nil, fmt.Errorf("assistant call failed: %w", err)
}
kunlog.Trace("[robot-agent] CallStream completed: assistantID=%s elapsed=%v", assistantID, time.Since(callStart).Round(time.Second))
result := &CallResult{Response: response}
if response.Next != nil {
result.Next = response.Next
@ -353,14 +368,23 @@ func (c *AgentCaller) CallStreamRaw(ctx *robottypes.Context, assistantID string,
opts.OnMessage = onMessage
}
agentCtx := c.buildAgentContext(ctx)
defer agentCtx.Release()
agentCtx := c.buildAgentContext(ctx, assistantID)
defer func() {
kunlog.Trace("[robot-agent] releasing context (CallStreamRaw): assistantID=%s chatID=%s", assistantID, c.ChatID)
agentCtx.Release()
}()
callStart := time.Now()
kunlog.Trace("[robot-agent] CallStreamRaw started: assistantID=%s chatID=%s", assistantID, c.ChatID)
response, err := ast.Stream(agentCtx, messages, opts)
if err != nil {
kunlog.Trace("[robot-agent] CallStreamRaw failed: assistantID=%s elapsed=%v err=%v", assistantID, time.Since(callStart).Round(time.Second), err)
return nil, fmt.Errorf("assistant call failed: %w", err)
}
kunlog.Trace("[robot-agent] CallStreamRaw completed: assistantID=%s elapsed=%v", assistantID, time.Since(callStart).Round(time.Second))
result := &CallResult{Response: response}
if response.Next != nil {
result.Next = response.Next
@ -390,7 +414,7 @@ func (c *AgentCaller) CallWithMessagesStreamRaw(ctx *robottypes.Context, assista
}
// buildAgentContext converts robot context to agent context
func (c *AgentCaller) buildAgentContext(ctx *robottypes.Context) *agentcontext.Context {
func (c *AgentCaller) buildAgentContext(ctx *robottypes.Context, assistantID string) *agentcontext.Context {
// Build authorized info for agent context
var authorized *oauthtypes.AuthorizedInfo
if ctx.Auth != nil {
@ -403,10 +427,14 @@ func (c *AgentCaller) buildAgentContext(ctx *robottypes.Context) *agentcontext.C
// Create a new agent context
// Use ChatID for multi-turn conversations, empty for single calls
agentCtx := agentcontext.New(ctx.Context, authorized, c.ChatID)
agentCtx.AssistantID = assistantID
// Set locale if available
// Propagate locale to agent context; fall back to "en" so that
// i18n.Tr / buildBoxDisplayName always resolve {{name}} templates.
if ctx.Locale != "" {
agentCtx.Locale = ctx.Locale
} else {
agentCtx.Locale = "en"
}
// Use noop logger to suppress LLM debug output for robot executions
@ -416,6 +444,7 @@ func (c *AgentCaller) buildAgentContext(ctx *robottypes.Context) *agentcontext.C
}
agentCtx.Logger = agentcontext.Noop()
kunlog.Trace("[robot-agent] context built: assistantID=%s chatID=%s contextID=%s", assistantID, c.ChatID, agentCtx.ID)
return agentCtx
}

6
agent/robot/executor/standard/runner.go
View file

@ -8,6 +8,7 @@ import (
"github.com/yaoapp/gou/mcp"
"github.com/yaoapp/gou/process"
kunlog "github.com/yaoapp/kun/log"
agentcontext "github.com/yaoapp/yao/agent/context"
robottypes "github.com/yaoapp/yao/agent/robot/types"
)
@ -146,6 +147,9 @@ func (r *Runner) executeAssistantTask(task *robottypes.Task, taskCtx *RunnerCont
input = "## Context\n\n" + taskCtx.SystemPrompt + "\n\n## Task\n\n" + input
}
kunlog.Trace("[robot-runner] executeAssistantTask: task=%s assistant=%s promptLen=%d prevResults=%d",
task.ID, task.ExecutorID, len(input), len(taskCtx.PreviousResults))
r.log.logTaskInput(task, input)
result, err := caller.CallWithMessages(r.ctx, task.ExecutorID, input)
@ -338,5 +342,7 @@ func (r *Runner) FormatPreviousResultsAsContext(results []robottypes.TaskResult)
sb.WriteString("\n")
}
contextLen := sb.Len()
kunlog.Trace("[robot-runner] FormatPreviousResultsAsContext: results=%d totalLen=%d", len(results), contextLen)
return sb.String()
}

1
agent/robot/manager/manager.go
View file

@ -460,6 +460,7 @@ func (m *Manager) TriggerManual(ctx *types.Context, memberID string, trigger typ
// Create a new context with the cancellable context from ExecutionController
// This allows Stop() to propagate cancellation to the executor
execCtx := types.NewContext(ctrlExec.Context(), ctx.Auth)
execCtx.Locale = ctx.Locale
// Submit to pool with the cancellable context and execution control
// The control interface allows executor to check pause state and wait if paused

5
agent/robot/manager/manager_test.go
View file

@ -146,9 +146,8 @@ func TestManagerTick(t *testing.T) {
// Should not have triggered (times mode robot only triggers at 09:00, 14:00)
execCount := m.Executor().ExecCount()
// Note: interval mode robot might trigger if enough time passed
// We just verify the times mode robot didn't trigger
assert.LessOrEqual(t, execCount, 1, "Times mode robot should not trigger at non-matching time")
// daemon always triggers, interval may trigger (LastRun=zero) => up to 2, but NOT 3
assert.LessOrEqual(t, execCount, 2, "Times mode robot should not trigger at non-matching time")
})
t.Run("tick with interval mode", func(t *testing.T) {

55
agent/sandbox/v2/claude/command.go
View file

@ -4,9 +4,13 @@ import (
"context"
"encoding/json"
"fmt"
"regexp"
"strings"
"time"
"github.com/google/uuid"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/store"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/sandbox/v2/types"
infra "github.com/yaoapp/yao/sandbox/v2"
@ -14,6 +18,34 @@ import (
const defaultA2OPort = 3099
var yaoSessionNS = uuid.MustParse("f47ac10b-58cc-4372-a567-0e02b2c3d479")
var safeNameRe = regexp.MustCompile(`[^a-zA-Z0-9_\-.]`)
func chatIDToSessionUUID(assistantID, chatID string) string {
return uuid.NewSHA1(yaoSessionNS, []byte(assistantID+":"+chatID)).String()
}
func sanitizeSessionName(chatID string) string {
return "yao-" + safeNameRe.ReplaceAllString(chatID, "_")
}
func chatSessionExists(storeKey string) bool {
s, err := store.Get("__yao.store")
if err != nil {
return false
}
return s.Has(storeKey)
}
func markChatSession(storeKey, sessionUUID string, ttl time.Duration) {
s, err := store.Get("__yao.store")
if err != nil {
return
}
s.Set(storeKey, sessionUUID, ttl)
}
type command struct {
shell []string
env map[string]string
@ -28,10 +60,18 @@ func (r *ClaudeRunner) buildCommand(ctx context.Context, req *types.StreamReques
if req.Config != nil {
assistantID = req.Config.ID
}
chatID := req.ChatID
var isContinuation bool
if chatID != "" {
storeKey := "claude-session:" + assistantID + ":" + chatID
isContinuation = chatSessionExists(storeKey)
} else {
isContinuation = hasExistingSession(ctx, computer, p, assistantID)
}
isContinuation := hasExistingSession(ctx, computer, p, assistantID)
env := buildEnv(req, p)
args := buildArgs(req, r, p, isContinuation, assistantID)
args := buildArgs(req, r, p, isContinuation, assistantID, chatID)
inputJSONL := buildInput(req.Messages, isContinuation)
var systemPrompt string
@ -141,7 +181,7 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
return env
}
func buildArgs(req *types.StreamRequest, r *ClaudeRunner, p platform, isContinuation bool, assistantID string) []string {
func buildArgs(req *types.StreamRequest, r *ClaudeRunner, p platform, isContinuation bool, assistantID, chatID string) []string {
var args []string
permMode := ""
@ -160,7 +200,16 @@ func buildArgs(req *types.StreamRequest, r *ClaudeRunner, p platform, isContinua
args = append(args, "--include-partial-messages")
args = append(args, "--verbose")
if chatID != "" {
sessionUUID := chatIDToSessionUUID(assistantID, chatID)
sessionName := sanitizeSessionName(chatID)
if isContinuation {
args = append(args, "--resume", sessionUUID)
} else {
args = append(args, "--session-id", sessionUUID)
}
args = append(args, "--name", sessionName)
} else if isContinuation {
args = append(args, "--continue")
}

104
agent/sandbox/v2/claude/command_test.go
View file

@ -92,7 +92,7 @@ func TestBuildArgs_Default(t *testing.T) {
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "")
args := buildArgs(req, r, p, false, "", "")
assert.Contains(t, args, "--input-format")
assert.Contains(t, args, "stream-json")
assert.Contains(t, args, "--output-format")
@ -106,7 +106,7 @@ func TestBuildArgs_Continuation(t *testing.T) {
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, true, "")
args := buildArgs(req, r, p, true, "", "")
assert.Contains(t, args, "--continue")
}
@ -124,7 +124,7 @@ func TestBuildArgs_PermissionMode(t *testing.T) {
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "")
args := buildArgs(req, r, p, false, "", "")
assert.Contains(t, args, "--dangerously-skip-permissions")
assert.Contains(t, args, "--permission-mode")
}
@ -135,7 +135,7 @@ func TestBuildArgs_MCP(t *testing.T) {
r := &ClaudeRunner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
p := testPlatform()
args := buildArgs(req, r, p, false, "test-assistant")
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__*")
@ -166,7 +166,7 @@ func TestBuildArgs_WhitelistOptions(t *testing.T) {
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "")
args := buildArgs(req, r, p, false, "", "")
assert.Contains(t, args, "--max-turns")
}
@ -382,3 +382,97 @@ func (f *fakeComputer) Stream(_ context.Context, _ []string, _ ...infra.ExecOpti
}
func (f *fakeComputer) VNC(_ context.Context) (string, error) { return "", nil }
func (f *fakeComputer) Proxy(_ context.Context, _ int, _ string) (string, error) { return "", nil }
// --- chatIDToSessionUUID ---
func TestChatIDToSessionUUID_Deterministic(t *testing.T) {
u1 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
u2 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
assert.Equal(t, u1, u2, "same inputs should produce same UUID")
}
func TestChatIDToSessionUUID_DifferentAssistant(t *testing.T) {
u1 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
u2 := chatIDToSessionUUID("asst-2", "robot_m1_e1")
assert.NotEqual(t, u1, u2, "different assistantID should produce different UUID")
}
func TestChatIDToSessionUUID_ValidFormat(t *testing.T) {
u := chatIDToSessionUUID("asst-1", "robot_m1_e1")
assert.Regexp(t, `^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`, u)
}
// --- sanitizeSessionName ---
func TestSanitizeSessionName_Normal(t *testing.T) {
assert.Equal(t, "yao-robot_m1_e1", sanitizeSessionName("robot_m1_e1"))
}
func TestSanitizeSessionName_SpecialChars(t *testing.T) {
assert.Equal(t, "yao-user_s__chat_", sanitizeSessionName("user's \"chat\""))
}
func TestSanitizeSessionName_Empty(t *testing.T) {
assert.Equal(t, "yao-", sanitizeSessionName(""))
}
// --- buildArgs with session ---
func TestBuildArgs_SessionID_NewSession(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "asst-1", "robot_m1_e1")
assert.Contains(t, args, "--session-id")
assert.Contains(t, args, "--name")
assert.Contains(t, args, "yao-robot_m1_e1")
assert.NotContains(t, args, "--resume")
assert.NotContains(t, args, "--continue")
sidIdx := -1
for i, a := range args {
if a == "--session-id" {
sidIdx = i
break
}
}
require.Greater(t, sidIdx, -1)
assert.Regexp(t, `^[0-9a-f]{8}-`, args[sidIdx+1])
}
func TestBuildArgs_SessionID_Continuation(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
p := testPlatform()
args := buildArgs(req, r, p, true, "asst-1", "robot_m1_e1")
assert.Contains(t, args, "--resume")
assert.Contains(t, args, "--name")
assert.NotContains(t, args, "--session-id")
assert.NotContains(t, args, "--continue")
resumeIdx := -1
for i, a := range args {
if a == "--resume" {
resumeIdx = i
break
}
}
require.Greater(t, resumeIdx, -1)
assert.Regexp(t, `^[0-9a-f]{8}-`, args[resumeIdx+1])
}
func TestBuildArgs_EmptyChatID_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")
assert.NotContains(t, args, "--session-id")
assert.NotContains(t, args, "--name")
}

36
agent/sandbox/v2/claude/parse.go
View file

@ -68,11 +68,29 @@ func (p *streamParser) parse(ctx context.Context, stdout io.ReadCloser) error {
scanner := bufio.NewScanner(stdout)
scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024)
startTime := time.Now()
lineCount := 0
lastHeartbeat := time.Now()
lastEventType := ""
log.Trace("[claude-parse] stream started")
for scanner.Scan() {
line := scanner.Text()
if line == "" {
continue
}
lineCount++
if time.Since(lastHeartbeat) > 30*time.Second {
builderLen := 0
if p.curTool != nil {
builderLen = p.curTool.inputJSON.Len()
}
log.Trace("[claude-parse] heartbeat: lines=%d elapsed=%v lastEvent=%s toolBuilderLen=%d",
lineCount, time.Since(startTime).Round(time.Second), lastEventType, builderLen)
lastHeartbeat = time.Now()
}
var msg map[string]any
if err := json.Unmarshal([]byte(line), &msg); err != nil {
@ -85,6 +103,7 @@ func (p *streamParser) parse(ctx context.Context, stdout io.ReadCloser) error {
}
msgType, _ := msg["type"].(string)
lastEventType = msgType
var stopped bool
switch msgType {
@ -97,16 +116,22 @@ func (p *streamParser) parse(ctx context.Context, stdout io.ReadCloser) error {
case "user":
stopped = p.handleUser(msg)
case "result":
log.Trace("[claude-parse] stream ended: lines=%d elapsed=%v completed=true", lineCount, time.Since(startTime).Round(time.Second))
return p.handleResult(msg)
case "error":
log.Trace("[claude-parse] stream ended with error: lines=%d elapsed=%v", lineCount, time.Since(startTime).Round(time.Second))
return p.handleError(msg)
}
if stopped {
log.Trace("[claude-parse] stream stopped by handler: lines=%d elapsed=%v", lineCount, time.Since(startTime).Round(time.Second))
return nil
}
}
log.Trace("[claude-parse] stream ended: lines=%d elapsed=%v completed=%v scanErr=%v",
lineCount, time.Since(startTime).Round(time.Second), p.completed, scanner.Err())
if err := scanner.Err(); err != nil {
log.Trace("[claude-parse] scanner error: %v (ctx.Err=%v)", err, ctx.Err())
if ctx.Err() != nil {
@ -346,6 +371,10 @@ func (p *streamParser) onContentBlockDelta(event map[string]any) (stopped bool)
return false
}
p.curTool.inputJSON.WriteString(partial)
builderLen := p.curTool.inputJSON.Len()
if builderLen > 0 && builderLen%100000 < len(partial) {
log.Trace("[claude-parse] WARN: tool %s inputJSON growing: %d bytes", p.curTool.name, builderLen)
}
if p.handler != nil {
return p.emitExecute(map[string]any{
"input_delta": p.curTool.inputJSON.String(),
@ -381,10 +410,15 @@ func (p *streamParser) handleAssistant(msg map[string]any) (stopped bool) {
itemType, _ := ci["type"].(string)
if itemType == "tool_use" && p.handler != nil {
toolID, _ := ci["id"].(string)
if _, alreadyStreamed := p.toolNames[toolID]; alreadyStreamed && toolID != "" {
continue
}
p.closeTextMessage()
toolName, _ := ci["name"].(string)
toolID, _ := ci["id"].(string)
if toolID == "" {
toolID = fmt.Sprintf("tool_%d_%d", p.toolIndex, time.Now().UnixNano())
}

9
agent/sandbox/v2/claude/plat_win.go
View file

@ -76,6 +76,15 @@ func (w *windowsPlatform) KillCmd(pattern string) []string {
return w.ShellCmd(script)
}
func (w *windowsPlatform) KillSessionCmd(sessionName string) []string {
script := fmt.Sprintf(
"Get-Process -ErrorAction SilentlyContinue | "+
"Where-Object { $_.CommandLine -like '*%s*' } | "+
"ForEach-Object { taskkill /F /T /PID $_.Id 2>$null }",
sessionName)
return w.ShellCmd(script)
}
func (w *windowsPlatform) ListDirCmd(dir string) []string {
return w.ShellCmd(fmt.Sprintf("Get-ChildItem -Name '%s'", dir))
}

5
agent/sandbox/v2/claude/platform.go
View file

@ -20,6 +20,7 @@ type platform interface {
RootDir() string
ShellCmd(script string) []string
KillCmd(pattern string) []string
KillSessionCmd(sessionName string) []string
ListDirCmd(dir string) []string
ConfigDir() string
XauthoritySetup(workDir string) string
@ -63,6 +64,10 @@ func (b *posixBase) KillCmd(pattern string) []string {
return []string{"sh", "-c", fmt.Sprintf("pkill -f '%s' || true", pattern)}
}
func (b *posixBase) KillSessionCmd(sessionName string) []string {
return []string{"sh", "-c", fmt.Sprintf("pkill -9 -f '%s' || true", sessionName)}
}
func (b *posixBase) ListDirCmd(dir string) []string {
return []string{"ls", dir}
}

21
agent/sandbox/v2/claude/platform_test.go
View file

@ -262,6 +262,27 @@ func TestWindows_ShellCmd_Cmd(t *testing.T) {
assert.Equal(t, []string{"cmd.exe", "/C", "echo hello"}, cmd)
}
func TestPosixBase_KillSessionCmd(t *testing.T) {
b := newTestPosixBase("linux")
cmd := b.KillSessionCmd("yao-robot_m1_e1")
require.Len(t, cmd, 3)
assert.Equal(t, "sh", cmd[0])
assert.Equal(t, "-c", cmd[1])
assert.Contains(t, cmd[2], "pkill -9 -f")
assert.Contains(t, cmd[2], "yao-robot_m1_e1")
assert.Contains(t, cmd[2], "|| true")
}
func TestWindows_KillSessionCmd(t *testing.T) {
w := newWindowsPlatform(`C:\ws`, "pwsh", "")
cmd := w.KillSessionCmd("yao-robot_m1_e1")
require.Len(t, cmd, 4)
assert.Equal(t, "pwsh", cmd[0])
assert.Contains(t, cmd[3], "CommandLine")
assert.Contains(t, cmd[3], "yao-robot_m1_e1")
assert.Contains(t, cmd[3], "taskkill")
}
func TestWindows_KillCmd(t *testing.T) {
w := newWindowsPlatform(`C:\ws`, "pwsh", "")
cmd := w.KillCmd("claude")

33
agent/sandbox/v2/claude/runner.go
View file

@ -5,6 +5,7 @@ import (
"encoding/json"
"fmt"
"strings"
"time"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/log"
@ -20,6 +21,7 @@ type ClaudeRunner struct {
hasMCP bool
mcpToolPattern string
lastCompleted bool
lastChatID string
logger *agentContext.RequestLogger
}
@ -107,16 +109,37 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
r.logger = agentContext.NoopLogger()
}
sess, err := startSession(ctx, computer, p, cmd, r.logger)
chatID := req.ChatID
r.lastChatID = chatID
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
log.Trace("[claude-runner] Stream started: assistantID=%s chatID=%s promptLen=%d", assistantID, chatID, len(cmd.shell))
sess, err := startSession(ctx, computer, p, cmd, chatID, r.logger)
if err != nil {
return err
}
streamStart := time.Now()
completed, err := sess.runStream(handler)
r.lastCompleted = completed
r.logger.Debug("Stream: runStream returned completed=%v err=%v", completed, err)
elapsed := time.Since(streamStart).Round(time.Second)
log.Trace("[claude-runner] Stream finished: assistantID=%s chatID=%s completed=%v elapsed=%v err=%v", assistantID, chatID, completed, elapsed, err)
r.logger.Debug("Stream: runStream returned completed=%v err=%v elapsed=%v", completed, err, elapsed)
if completed {
sess.shutdown()
if chatID != "" {
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
storeKey := "claude-session:" + assistantID + ":" + chatID
sessionUUID := chatIDToSessionUUID(assistantID, chatID)
markChatSession(storeKey, sessionUUID, 90*24*time.Hour)
}
}
return err
}
@ -128,6 +151,8 @@ func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) err
return nil
}
log.Trace("[claude-runner] Cleanup: chatID=%s lastCompleted=%v", r.lastChatID, r.lastCompleted)
if r.lastCompleted {
if r.logger != nil {
r.logger.Info("cleanup: stream completed normally, preserving child processes")
@ -137,8 +162,12 @@ func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) err
if r.mode != "service" {
p := resolvePlatform(computer)
if r.lastChatID != "" {
computer.Exec(ctx, p.KillSessionCmd(sanitizeSessionName(r.lastChatID)))
} else {
computer.Exec(ctx, p.KillCmd("claude"))
}
}
return nil
}

28
agent/sandbox/v2/claude/session.go
View file

@ -22,16 +22,17 @@ type session struct {
stderr strings.Builder
stderrMu sync.Mutex
logger *agentContext.RequestLogger
chatID string
}
func startSession(ctx context.Context, computer infra.Computer, p platform, cmd command, logger *agentContext.RequestLogger) (*session, error) {
func startSession(ctx context.Context, computer infra.Computer, p platform, cmd command, chatID string, 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))
logger.Info("claude session starting: cmd=%s workDir=%s platform=%s stdinLen=%d chatID=%s",
cmd.shell, cmd.workDir, p.OS(), len(cmd.stdin), chatID)
execStream, err := computer.Stream(ctx, cmd.shell, opts...)
if err != nil {
@ -44,6 +45,7 @@ func startSession(ctx context.Context, computer infra.Computer, p platform, cmd
plat: p,
exec: execStream,
logger: logger,
chatID: chatID,
}, nil
}
@ -111,6 +113,19 @@ func (s *session) collectStderr() {
}()
}
// killProcess terminates the Claude CLI process. When chatID is available,
// uses KillSessionCmd for precise matching; otherwise falls back to KillCmd.
func (s *session) killProcess(ctx context.Context) {
if s.chatID != "" {
name := sanitizeSessionName(s.chatID)
result, err := s.computer.Exec(ctx, s.plat.KillSessionCmd(name))
s.logger.Debug("killProcess: KillSessionCmd(%s) exitCode=%d err=%v", name, result.ExitCode, err)
return
}
result, err := s.computer.Exec(ctx, s.plat.KillCmd("claude"))
s.logger.Debug("killProcess: KillCmd(claude) exitCode=%d err=%v", result.ExitCode, err)
}
// watchCancel monitors context cancellation and kills the Claude process.
// Returns a cleanup function that must be deferred.
func (s *session) watchCancel() func() {
@ -121,7 +136,7 @@ func (s *session) watchCancel() func() {
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.killProcess(killCtx)
s.exec.Cancel()
case <-done:
}
@ -140,12 +155,11 @@ func (s *session) watchCancel() func() {
// would actively terminate child processes (web servers, etc.). Those children
// survive because they run in separate process groups/sessions.
func (s *session) shutdown() {
s.logger.Info("shutting down completed claude exec session")
s.logger.Info("shutting down completed claude exec session: chatID=%s", s.chatID)
killCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
result, err := s.computer.Exec(killCtx, []string{"sh", "-c", "pkill -9 -x claude || true"})
s.logger.Debug("shutdown: pkill -9 -x claude exitCode=%d err=%v", result.ExitCode, err)
s.killProcess(killCtx)
s.exec.Cancel()
}