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

@@ -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")
 	}
 

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")
+}

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))
 }

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}
 }

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")

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

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.killProcess(killCtx)
-	s.logger.Debug("shutdown: pkill -9 -x claude exitCode=%d err=%v", result.ExitCode, err)
 	s.exec.Cancel()
 }