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

feat(claude): enhance session management and argument building

Browse source 
- 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.
This commit is contained in:
Max 2026-03-25 17:25:55 +08:00
parent 427c38bac7
commit abde3adf13
7 changed files with 225 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

22
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,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

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

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