feat(claude): refactor runner and enhance request handling
- Renamed ClaudeRunner to Runner for consistency across the codebase, aligning with the sandbox Runner interface. - Updated buildCommand and Stream methods to utilize AssistantID directly from StreamRequest, improving clarity and reducing dependency on Config. - Modified tests to reflect the changes in runner instantiation and argument handling, ensuring compatibility with the new Runner structure. - Enhanced environment variable setup in buildEnv to include WORKDIR, streamlining the execution context for the runner.
This commit is contained in:
Max
parent
6e003875b4
commit
e90808e861
7 changed files with 107 additions and 69 deletions
|
|
@ -13,6 +13,7 @@ import (
|
|||
"github.com/yaoapp/yao/agent/output/message"
|
||||
sandboxv2 "github.com/yaoapp/yao/agent/sandbox/v2"
|
||||
sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types"
|
||||
store "github.com/yaoapp/yao/agent/store/types"
|
||||
"github.com/yaoapp/yao/config"
|
||||
infraV2 "github.com/yaoapp/yao/sandbox/v2"
|
||||
traceTypes "github.com/yaoapp/yao/trace/types"
|
||||
|
|
@ -124,6 +125,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
|
|||
Computer: computer,
|
||||
Config: cfg,
|
||||
Connector: conn,
|
||||
AssistantID: ast.ID,
|
||||
SkillsDir: skillsDir,
|
||||
AssistantDir: assistantDir,
|
||||
MCPServers: mcpServers,
|
||||
|
|
@ -172,10 +174,12 @@ func (ast *Assistant) executeSandboxV2Stream(
|
|||
cfg := ast.SandboxV2
|
||||
manager := infraV2.M()
|
||||
|
||||
// Build system prompt.
|
||||
// Build system prompt (parse $CTX variables the same way as buildSystemPrompts).
|
||||
var systemPrompt string
|
||||
if len(ast.Prompts) > 0 {
|
||||
for _, pr := range ast.Prompts {
|
||||
ctxVars := ast.buildContextVariables(ctx)
|
||||
parsed := store.Prompts(ast.Prompts).Parse(ctxVars)
|
||||
for _, pr := range parsed {
|
||||
if pr.Role == "system" && pr.Content != "" {
|
||||
systemPrompt = pr.Content
|
||||
break
|
||||
|
|
@ -199,6 +203,7 @@ func (ast *Assistant) executeSandboxV2Stream(
|
|||
Computer: p.Computer,
|
||||
Config: cfg,
|
||||
Connector: conn,
|
||||
AssistantID: ast.ID,
|
||||
Messages: p.Messages,
|
||||
SystemPrompt: systemPrompt,
|
||||
ChatID: ctx.ChatID,
|
||||
|
|
|
|||
|
|
@ -53,13 +53,10 @@ type command struct {
|
|||
workDir string
|
||||
}
|
||||
|
||||
func (r *ClaudeRunner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command {
|
||||
func (r *Runner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command {
|
||||
computer := req.Computer
|
||||
workDir := computer.GetWorkDir()
|
||||
assistantID := ""
|
||||
if req.Config != nil {
|
||||
assistantID = req.Config.ID
|
||||
}
|
||||
assistantID := req.AssistantID
|
||||
chatID := req.ChatID
|
||||
|
||||
var isContinuation bool
|
||||
|
|
@ -110,14 +107,17 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
|
|||
for k, v := range p.HomeEnv(workDir) {
|
||||
env[k] = v
|
||||
}
|
||||
env["WORKDIR"] = workDir
|
||||
|
||||
assistantID := ""
|
||||
if req.Config != nil {
|
||||
assistantID = req.Config.ID
|
||||
}
|
||||
assistantID := req.AssistantID
|
||||
if assistantID != "" {
|
||||
configDir := p.PathJoin(workDir, ".yao", "assistants", assistantID)
|
||||
env["CLAUDE_CONFIG_DIR"] = configDir
|
||||
env["CTX_ASSISTANT_ID"] = assistantID
|
||||
// CTX_SKILLS_DIR: absolute path to the skills directory inside the sandbox.
|
||||
// Use this in skill scripts instead of constructing the path manually,
|
||||
// so it works correctly on all platforms (Linux, macOS, Windows).
|
||||
env["CTX_SKILLS_DIR"] = p.PathJoin(workDir, ".yao", "assistants", assistantID, "skills")
|
||||
}
|
||||
|
||||
if req.Connector != nil {
|
||||
|
|
@ -130,11 +130,11 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
|
|||
env["ANTHROPIC_BASE_URL"] = host
|
||||
env["ANTHROPIC_API_KEY"] = key
|
||||
if model != "" {
|
||||
env["ANTHROPIC_MODEL"] = model
|
||||
env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model
|
||||
env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model
|
||||
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model
|
||||
env["CLAUDE_CODE_SUBAGENT_MODEL"] = model
|
||||
env["ANTHROPIC_MODEL"] = "claude-sonnet-4-6"
|
||||
env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = "claude-sonnet-4-6"
|
||||
env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = "claude-sonnet-4-6"
|
||||
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = "claude-sonnet-4-6"
|
||||
env["CLAUDE_CODE_SUBAGENT_MODEL"] = "claude-sonnet-4-6"
|
||||
}
|
||||
} else {
|
||||
connectorID := req.Connector.ID()
|
||||
|
|
@ -181,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, chatID string) []string {
|
||||
func buildArgs(req *types.StreamRequest, r *Runner, p platform, isContinuation bool, assistantID, chatID string) []string {
|
||||
var args []string
|
||||
|
||||
permMode := ""
|
||||
|
|
@ -248,12 +248,18 @@ func buildSandboxEnvPrompt(p platform, workDir string) string {
|
|||
|
||||
shellNote := p.EnvPromptNote()
|
||||
|
||||
envVarSyntax := "$VAR_NAME"
|
||||
if osName == "windows" {
|
||||
envVarSyntax = "$env:VAR_NAME"
|
||||
}
|
||||
|
||||
return fmt.Sprintf(`## Sandbox Environment
|
||||
|
||||
- **Operating System**: %[2]s
|
||||
- **Shell**: %[3]s
|
||||
- **Working Directory**: %[1]s
|
||||
- **File Access**: You have full read/write access to %[1]s%[4]s
|
||||
- **File Access**: You have full read/write access to %[1]s
|
||||
- **Environment variable syntax**: `+"`%[5]s`"+` (e.g. `+"`$CTX_SKILLS_DIR`"+` on POSIX, `+"`$env:CTX_SKILLS_DIR`"+` on Windows)%[4]s
|
||||
|
||||
## User Attachments
|
||||
|
||||
|
|
@ -261,7 +267,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)
|
||||
`, workDir, osName, shell, shellNote, envVarSyntax)
|
||||
}
|
||||
|
||||
func hasExistingSession(ctx context.Context, computer infra.Computer, p platform, assistantID string) bool {
|
||||
|
|
|
|||
|
|
@ -33,7 +33,8 @@ func TestBuildEnv_HomeEnv(t *testing.T) {
|
|||
|
||||
func TestBuildEnv_ConfigDirIsolation(t *testing.T) {
|
||||
req := &types.StreamRequest{
|
||||
Config: &types.SandboxConfig{ID: "my-assistant"},
|
||||
Config: &types.SandboxConfig{ID: "my-assistant"},
|
||||
AssistantID: "my-assistant",
|
||||
}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
p := testPlatform()
|
||||
|
|
@ -89,7 +90,7 @@ func TestBuildEnv_Secrets(t *testing.T) {
|
|||
func TestBuildArgs_Default(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, false, "", "")
|
||||
|
|
@ -103,7 +104,7 @@ func TestBuildArgs_Default(t *testing.T) {
|
|||
func TestBuildArgs_Continuation(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, true, "", "")
|
||||
|
|
@ -121,7 +122,7 @@ func TestBuildArgs_PermissionMode(t *testing.T) {
|
|||
},
|
||||
}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, false, "", "")
|
||||
|
|
@ -132,7 +133,7 @@ func TestBuildArgs_PermissionMode(t *testing.T) {
|
|||
func TestBuildArgs_MCP(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
|
||||
r := &Runner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, false, "test-assistant", "")
|
||||
|
|
@ -163,7 +164,7 @@ func TestBuildArgs_WhitelistOptions(t *testing.T) {
|
|||
},
|
||||
}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, false, "", "")
|
||||
|
|
@ -378,7 +379,7 @@ func TestSanitizeSessionName_Empty(t *testing.T) {
|
|||
func TestBuildArgs_SessionID_NewSession(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, false, "asst-1", "robot_m1_e1")
|
||||
|
|
@ -402,7 +403,7 @@ func TestBuildArgs_SessionID_NewSession(t *testing.T) {
|
|||
func TestBuildArgs_SessionID_Continuation(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, true, "asst-1", "robot_m1_e1")
|
||||
|
|
@ -425,7 +426,7 @@ func TestBuildArgs_SessionID_Continuation(t *testing.T) {
|
|||
func TestBuildArgs_EmptyChatID_Continuation(t *testing.T) {
|
||||
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
|
||||
req.Computer = newFakeComputer("/workspace")
|
||||
r := &ClaudeRunner{}
|
||||
r := &Runner{}
|
||||
p := testPlatform()
|
||||
|
||||
args := buildArgs(req, r, p, true, "", "")
|
||||
|
|
|
|||
|
|
@ -172,6 +172,30 @@ func (p *streamParser) closeTextMessage() {
|
|||
}
|
||||
}
|
||||
|
||||
// closeCurrentTool closes the in-flight streaming tool message (if any),
|
||||
// flushing its accumulated input and emitting message_end. This must be
|
||||
// called before opening a new message group so that the downstream handler
|
||||
// never sees interleaved message_start/message_end pairs.
|
||||
func (p *streamParser) closeCurrentTool() {
|
||||
if p.curTool == nil {
|
||||
return
|
||||
}
|
||||
toolID := p.curTool.id
|
||||
inputStr := p.curTool.inputJSON.String()
|
||||
if inputStr != "" {
|
||||
p.toolInputs[toolID] = inputStr
|
||||
summary := extractSummary(p.curTool.name, inputStr)
|
||||
if summary != "" {
|
||||
p.toolSummaries[toolID] = summary
|
||||
p.emitExecute(map[string]any{
|
||||
"summary": summary,
|
||||
})
|
||||
}
|
||||
}
|
||||
p.endMessage()
|
||||
p.curTool = nil
|
||||
}
|
||||
|
||||
func (p *streamParser) ensureTextMessage() (stopped bool) {
|
||||
if !p.textActive {
|
||||
_, stopped = p.beginMessage("text")
|
||||
|
|
@ -318,22 +342,7 @@ func (p *streamParser) onContentBlockStart(event map[string]any) (stopped bool)
|
|||
}
|
||||
|
||||
func (p *streamParser) onContentBlockStop() (stopped bool) {
|
||||
if p.curTool != nil {
|
||||
toolID := p.curTool.id
|
||||
inputStr := p.curTool.inputJSON.String()
|
||||
if inputStr != "" {
|
||||
p.toolInputs[toolID] = inputStr
|
||||
summary := extractSummary(p.curTool.name, inputStr)
|
||||
if summary != "" {
|
||||
p.toolSummaries[toolID] = summary
|
||||
p.emitExecute(map[string]any{
|
||||
"summary": summary,
|
||||
})
|
||||
}
|
||||
}
|
||||
p.endMessage()
|
||||
p.curTool = nil
|
||||
}
|
||||
p.closeCurrentTool()
|
||||
return false
|
||||
}
|
||||
|
||||
|
|
@ -417,6 +426,7 @@ func (p *streamParser) handleAssistant(msg map[string]any) (stopped bool) {
|
|||
}
|
||||
|
||||
p.closeTextMessage()
|
||||
p.closeCurrentTool()
|
||||
|
||||
toolName, _ := ci["name"].(string)
|
||||
if toolID == "" {
|
||||
|
|
@ -488,12 +498,17 @@ func (p *streamParser) handleUser(msg map[string]any) (stopped bool) {
|
|||
continue
|
||||
}
|
||||
|
||||
// Close any open text message before opening an execute message,
|
||||
// otherwise textActive stays true while currentGroupID gets
|
||||
// overwritten by the execute message lifecycle, causing subsequent
|
||||
// text chunks to be emitted without a message_id.
|
||||
// Close any open text message before opening an execute message.
|
||||
p.closeTextMessage()
|
||||
|
||||
// When Claude CLI executes tools in parallel, tool_result messages
|
||||
// can arrive while a new tool_use is still streaming. The downstream
|
||||
// handler (stream.go) tracks only a single currentGroupID, so we
|
||||
// must close the in-flight streaming tool message before opening
|
||||
// the result message — otherwise the message_start/message_end
|
||||
// pairs become interleaved and chunks lose their message_id.
|
||||
p.closeCurrentTool()
|
||||
|
||||
toolUseID, _ := ci["tool_use_id"].(string)
|
||||
content := ci["content"]
|
||||
isError, _ := ci["is_error"].(bool)
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ import (
|
|||
infra "github.com/yaoapp/yao/sandbox/v2"
|
||||
)
|
||||
|
||||
// ClaudeRunner implements the Runner interface for Claude CLI (mode=cli).
|
||||
type ClaudeRunner struct {
|
||||
// Runner implements the sandbox Runner interface for Claude CLI (mode=cli).
|
||||
type Runner struct {
|
||||
mode string
|
||||
hasMCP bool
|
||||
mcpToolPattern string
|
||||
|
|
@ -25,21 +25,22 @@ type ClaudeRunner struct {
|
|||
logger *agentContext.RequestLogger
|
||||
}
|
||||
|
||||
// New creates a new ClaudeRunner.
|
||||
func New() *ClaudeRunner {
|
||||
return &ClaudeRunner{mode: "cli"}
|
||||
// New creates a new Runner.
|
||||
func New() *Runner {
|
||||
return &Runner{mode: "cli"}
|
||||
}
|
||||
|
||||
func (r *ClaudeRunner) Name() string { return "claude" }
|
||||
// Name returns the runner identifier.
|
||||
func (r *Runner) Name() string { return "claude" }
|
||||
|
||||
// Prepare executes user-defined and runner-specific prepare steps.
|
||||
func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) error {
|
||||
func (r *Runner) Prepare(ctx context.Context, req *types.PrepareRequest) error {
|
||||
r.mode = req.Config.Runner.Mode
|
||||
if r.mode == "" {
|
||||
r.mode = "cli"
|
||||
}
|
||||
|
||||
assistantID := req.Config.ID
|
||||
assistantID := req.AssistantID
|
||||
prefix := ".yao/assistants/" + assistantID
|
||||
if assistantID == "" {
|
||||
prefix = ".claude"
|
||||
|
|
@ -70,7 +71,7 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e
|
|||
}
|
||||
|
||||
if req.RunSteps != nil && len(steps) > 0 {
|
||||
if err := req.RunSteps(ctx, steps, req.Computer, req.Config.ID, req.ConfigHash, req.AssistantDir); err != nil {
|
||||
if err := req.RunSteps(ctx, steps, req.Computer, req.AssistantID, req.ConfigHash, req.AssistantDir); err != nil {
|
||||
return fmt.Errorf("claude prepare steps: %w", err)
|
||||
}
|
||||
}
|
||||
|
|
@ -79,7 +80,7 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e
|
|||
}
|
||||
|
||||
// Stream executes the Claude CLI and streams output to handler.
|
||||
func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, handler message.StreamFunc) error {
|
||||
func (r *Runner) Stream(ctx context.Context, req *types.StreamRequest, handler message.StreamFunc) error {
|
||||
computer := req.Computer
|
||||
if computer == nil {
|
||||
return fmt.Errorf("computer is nil")
|
||||
|
|
@ -111,10 +112,7 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
|
|||
|
||||
chatID := req.ChatID
|
||||
r.lastChatID = chatID
|
||||
assistantID := ""
|
||||
if req.Config != nil {
|
||||
assistantID = req.Config.ID
|
||||
}
|
||||
assistantID := req.AssistantID
|
||||
|
||||
log.Trace("[claude-runner] Stream started: assistantID=%s chatID=%s promptLen=%d", assistantID, chatID, len(cmd.shell))
|
||||
|
||||
|
|
@ -132,10 +130,6 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
|
|||
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)
|
||||
|
|
@ -146,7 +140,7 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
|
|||
|
||||
// Cleanup kills any remaining claude processes. If the stream completed
|
||||
// normally (received "result"), child processes are preserved.
|
||||
func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) error {
|
||||
func (r *Runner) Cleanup(ctx context.Context, computer infra.Computer) error {
|
||||
if computer == nil {
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ func runFileStep(ws workspace.FS, step types.PrepareStep) error {
|
|||
return fmt.Errorf("file step requires workspace")
|
||||
}
|
||||
|
||||
step.Path = expandTilde(step.Path)
|
||||
dir := path.Dir(step.Path)
|
||||
if dir != "." && dir != "/" {
|
||||
if err := ws.MkdirAll(dir, 0755); err != nil {
|
||||
|
|
@ -126,8 +127,9 @@ func runCopyStep(ws workspace.FS, step types.PrepareStep, assistantDir string) e
|
|||
src = "local:///" + pathpkg.Join(assistantDir, src)
|
||||
}
|
||||
|
||||
if _, err := ws.Copy(src, step.Dst); err != nil {
|
||||
return fmt.Errorf("copy %s -> %s: %w", src, step.Dst, err)
|
||||
dst := expandTilde(step.Dst)
|
||||
if _, err := ws.Copy(src, dst); err != nil {
|
||||
return fmt.Errorf("copy %s -> %s: %w", src, dst, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
@ -136,6 +138,19 @@ func isHostURI(s string) bool {
|
|||
return strings.HasPrefix(s, "local:///") || strings.HasPrefix(s, "tmp:///")
|
||||
}
|
||||
|
||||
// expandTilde replaces a leading "~/" with the empty string so the path
|
||||
// becomes relative to the workspace root (which is HOME inside the sandbox).
|
||||
// Paths without "~/" are returned unchanged.
|
||||
func expandTilde(p string) string {
|
||||
if strings.HasPrefix(p, "~/") {
|
||||
return p[2:]
|
||||
}
|
||||
if p == "~" {
|
||||
return "."
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func runExecStep(ctx context.Context, computer infra.Computer, step types.PrepareStep) error {
|
||||
if step.Cmd == "" {
|
||||
return fmt.Errorf("exec step requires cmd")
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ type PrepareRequest struct {
|
|||
Computer infra.Computer
|
||||
Config *SandboxConfig
|
||||
Connector connector.Connector
|
||||
AssistantID string // the assistant's own ID (e.g. "yao/postman")
|
||||
SkillsDir string
|
||||
AssistantDir string // absolute host path to the assistant source directory
|
||||
MCPServers []MCPServer
|
||||
|
|
@ -50,6 +51,7 @@ type StreamRequest struct {
|
|||
Computer infra.Computer
|
||||
Config *SandboxConfig
|
||||
Connector connector.Connector
|
||||
AssistantID string // the assistant's own ID (e.g. "yao/postman")
|
||||
Messages []agentContext.Message
|
||||
SystemPrompt string
|
||||
ChatID string
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue