From abde3adf130af1d271f60ad192d81121ac7e6a97 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 25 Mar 2026 17:25:55 +0800 Subject: [PATCH] feat(claude): enhance session management and argument building - Introduced session ID and name handling in the buildArgs function, allowing for better tracking of chat sessions. - Added chatIDToSessionUUID and sanitizeSessionName functions to generate and format session identifiers. - Updated the Stream method to store session information and manage session lifecycle more effectively. - Implemented KillSessionCmd for precise process termination based on session names in both Windows and POSIX platforms. - Enhanced tests to cover new session management features and ensure correct behavior in various scenarios. --- agent/sandbox/v2/claude/command.go | 57 ++++++++++++- agent/sandbox/v2/claude/command_test.go | 104 +++++++++++++++++++++-- agent/sandbox/v2/claude/plat_win.go | 9 ++ agent/sandbox/v2/claude/platform.go | 5 ++ agent/sandbox/v2/claude/platform_test.go | 21 +++++ agent/sandbox/v2/claude/runner.go | 22 ++++- agent/sandbox/v2/claude/session.go | 22 ++++- 7 files changed, 225 insertions(+), 15 deletions(-) diff --git a/agent/sandbox/v2/claude/command.go b/agent/sandbox/v2/claude/command.go index d15229a9..0ba92a02 100644 --- a/agent/sandbox/v2/claude/command.go +++ b/agent/sandbox/v2/claude/command.go @@ -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 isContinuation { + 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") } diff --git a/agent/sandbox/v2/claude/command_test.go b/agent/sandbox/v2/claude/command_test.go index 67e554fa..9ecdfdd5 100644 --- a/agent/sandbox/v2/claude/command_test.go +++ b/agent/sandbox/v2/claude/command_test.go @@ -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") +} diff --git a/agent/sandbox/v2/claude/plat_win.go b/agent/sandbox/v2/claude/plat_win.go index 79cdc4f5..f966315f 100644 --- a/agent/sandbox/v2/claude/plat_win.go +++ b/agent/sandbox/v2/claude/plat_win.go @@ -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)) } diff --git a/agent/sandbox/v2/claude/platform.go b/agent/sandbox/v2/claude/platform.go index c9a95a0f..b8f839e0 100644 --- a/agent/sandbox/v2/claude/platform.go +++ b/agent/sandbox/v2/claude/platform.go @@ -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} } diff --git a/agent/sandbox/v2/claude/platform_test.go b/agent/sandbox/v2/claude/platform_test.go index 71f4c453..326935a3 100644 --- a/agent/sandbox/v2/claude/platform_test.go +++ b/agent/sandbox/v2/claude/platform_test.go @@ -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") diff --git a/agent/sandbox/v2/claude/runner.go b/agent/sandbox/v2/claude/runner.go index d08a2c12..61cf6df2 100644 --- a/agent/sandbox/v2/claude/runner.go +++ b/agent/sandbox/v2/claude/runner.go @@ -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,7 +109,10 @@ 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 + + sess, err := startSession(ctx, computer, p, cmd, chatID, r.logger) if err != nil { return err } @@ -117,6 +122,15 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han r.logger.Debug("Stream: runStream returned completed=%v err=%v", completed, err) 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 } @@ -137,7 +151,11 @@ func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) err if r.mode != "service" { p := resolvePlatform(computer) - computer.Exec(ctx, p.KillCmd("claude")) + if r.lastChatID != "" { + computer.Exec(ctx, p.KillSessionCmd(sanitizeSessionName(r.lastChatID))) + } else { + computer.Exec(ctx, p.KillCmd("claude")) + } } return nil diff --git a/agent/sandbox/v2/claude/session.go b/agent/sandbox/v2/claude/session.go index 1f0450dc..377e081d 100644 --- a/agent/sandbox/v2/claude/session.go +++ b/agent/sandbox/v2/claude/session.go @@ -22,9 +22,10 @@ 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)) @@ -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: } @@ -144,8 +159,7 @@ func (s *session) shutdown() { 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() }