yao/agent/sandbox/v2/claude/command.go
Max 194faac9b7 feat(agent): enhance sandbox V2 initialization and role management
- Refactored the initSandboxV2 function to return a structured result, consolidating the runner, computer, configuration, cleanup function, loading message ID, and roles into a single return type.
- Updated the Stream method to utilize the new sandboxV2InitResult structure, improving clarity and reducing complexity in handling sandbox initialization.
- Introduced role management enhancements, allowing for pre-resolved role connectors to be passed through the request, streamlining connector resolution during execution.
- Adjusted various components to support the new roles structure, ensuring consistent handling across the agent's sandbox operations.
- Added logging for connector resolution and role management, improving diagnostics and traceability during sandbox execution.
- Updated .gitignore to include tools/TOOL-REGISTRATION.md for better project organization.
2026-05-03 19:11:29 +08:00

576 lines
16 KiB
Go

package claude
import (
"context"
"encoding/json"
"fmt"
"regexp"
"strings"
"time"
"github.com/google/uuid"
"github.com/yaoapp/gou/connector"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/str"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/sandbox/v2/types"
infra "github.com/yaoapp/yao/sandbox/v2"
)
const defaultA2OPort = 3099
const defaultA2OMaxOutputTokens = 16384
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
stdin []byte
workDir string
}
func (r *Runner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command {
computer := req.Computer
workDir := computer.GetWorkDir()
assistantID := req.AssistantID
chatID := req.ChatID
var isContinuation bool
if chatID != "" {
storeKey := "claude-session:" + assistantID + ":" + chatID
isContinuation = chatSessionExists(storeKey)
} else {
isContinuation = hasExistingSession(ctx, computer, p, assistantID)
}
env := buildEnv(req, p)
args := buildArgs(req, r, p, isContinuation, assistantID, chatID)
inputJSONL := buildLastUserMessageJSONL(req.Messages)
var systemPrompt string
envPrompt := buildSandboxEnvPrompt(p, workDir)
if !isContinuation && req.SystemPrompt != "" {
systemPrompt = req.SystemPrompt + "\n\n" + envPrompt
} else if !isContinuation {
systemPrompt = envPrompt
}
promptFile := p.PathJoin(workDir, ".yao", "assistants", assistantID, "system-prompt.txt")
if assistantID == "" {
promptFile = p.PathJoin(workDir, ".yao", ".system-prompt.txt")
}
// On continuation turns the system prompt is not re-sent, but the previously
// written prompt file is still on disk. Pass --append-system-prompt-file so
// Claude keeps the same constraints (e.g. workspace path rules) across the
// entire session without injecting a duplicate system turn.
if isContinuation {
args = append(args, "--append-system-prompt-file", promptFile)
}
script, stdin := p.BuildScript(scriptInput{
args: args,
systemPrompt: systemPrompt,
inputJSONL: inputJSONL,
workDir: workDir,
promptFile: promptFile,
})
return command{
shell: p.ShellCmd(script),
env: env,
stdin: stdin,
workDir: workDir,
}
}
func buildEnv(req *types.StreamRequest, p platform) map[string]string {
env := make(map[string]string)
workDir := req.Computer.GetWorkDir()
// Workspace ID
workspace := req.Computer.Workplace()
if workspace != nil {
workspaceID, err := workspace.GetID()
if err == nil {
env["CTX_WORKSPACE_ID"] = workspaceID
}
}
for k, v := range p.HomeEnv(workDir) {
env[k] = v
}
env["WORKDIR"] = workDir
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 {
setting := req.Connector.Setting()
var host, key, model string
if lc, ok := req.Connector.(goullm.LLMConnector); ok {
host = lc.GetURL()
key = lc.GetKey()
model = lc.GetModel()
}
if host == "" {
host, _ = setting["host"].(string)
}
if key == "" {
key, _ = setting["key"].(string)
}
if model == "" {
model, _ = setting["model"].(string)
}
isAnthropic := req.Connector.Is(connector.ANTHROPIC)
if isAnthropic {
setAnthropicModelEnv(env, host, key, model, req.Connector)
applyAnthropicRoleOverrides(env, host, req.Roles)
} else {
setA2OModelEnv(env, req.Connector.ID(), model, req.Connector)
applyA2ORoleOverrides(env, req.Roles)
}
if lc, ok := req.Connector.(goullm.LLMConnector); ok {
if caps := lc.GetCapabilities(); caps != nil {
if caps.MaxOutputTokens > 0 {
env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] = fmt.Sprintf("%d", caps.MaxOutputTokens)
}
if caps.MaxInputTokens > 0 {
env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = fmt.Sprintf("%d", caps.MaxInputTokens)
}
}
}
if _, ok := env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"]; !ok && !req.Connector.Is(connector.ANTHROPIC) {
env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] = fmt.Sprintf("%d", defaultA2OMaxOutputTokens)
}
if thinking, ok := setting["thinking"].(map[string]interface{}); ok {
thinkType, _ := thinking["type"].(string)
switch thinkType {
case "disabled":
env["MAX_THINKING_TOKENS"] = "0"
case "enabled":
if budget, ok := thinking["budget_tokens"].(float64); ok && budget > 0 {
env["MAX_THINKING_TOKENS"] = fmt.Sprintf("%d", int(budget))
}
}
}
}
if req.Config != nil && len(req.Config.Secrets) > 0 {
for k, v := range req.Config.Secrets {
env[k] = str.EnvVar(v)
}
}
if req.Token != nil {
if req.Token.Token != "" {
env["YAO_TOKEN"] = req.Token.Token
}
if req.Token.RefreshToken != "" {
env["YAO_REFRESH_TOKEN"] = req.Token.RefreshToken
}
}
logger := req.Logger
if logger == nil {
logger = agentContext.NoopLogger()
}
connectorID := ""
if req.Connector != nil {
connectorID = req.Connector.ID()
}
logger.Debug("claude-env: connector=%s isAnthropic=%v", connectorID, req.Connector != nil && req.Connector.Is(connector.ANTHROPIC))
logger.Debug("claude-env: ANTHROPIC_MODEL=%s", env["ANTHROPIC_MODEL"])
logger.Debug("claude-env: OPUS_MODEL=%s SONNET_MODEL=%s HAIKU_MODEL=%s",
env["ANTHROPIC_DEFAULT_OPUS_MODEL"],
env["ANTHROPIC_DEFAULT_SONNET_MODEL"],
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"])
logger.Debug("claude-env: CUSTOM_MODEL_OPTION=%s CAPABILITIES=%s",
env["ANTHROPIC_CUSTOM_MODEL_OPTION"],
env["ANTHROPIC_CUSTOM_MODEL_OPTION_SUPPORTED_CAPABILITIES"])
logger.Debug("claude-env: MAX_THINKING_TOKENS=%s", env["MAX_THINKING_TOKENS"])
return env
}
func buildArgs(req *types.StreamRequest, r *Runner, p platform, isContinuation bool, assistantID, chatID string) []string {
var args []string
permMode := ""
if req.Config != nil && req.Config.Runner.Options != nil {
if v, ok := req.Config.Runner.Options["permission_mode"]; ok {
permMode = fmt.Sprintf("%v", v)
}
}
if permMode == "bypassPermissions" {
args = append(args, "--dangerously-skip-permissions")
args = append(args, "--permission-mode", permMode)
}
args = append(args, "--input-format", "stream-json")
args = append(args, "--output-format", "stream-json")
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")
}
if req.Config != nil && req.Config.Runner.Options != nil {
for key, val := range req.Config.Runner.Options {
if flag, ok := claudeArgWhitelist[key]; ok {
args = append(args, flag, fmt.Sprintf("%v", val))
}
}
}
if r.hasMCP {
workDir := req.Computer.GetWorkDir()
mcpPath := p.PathJoin(workDir, ".yao", "assistants", assistantID, "mcp.json")
if assistantID == "" {
mcpPath = p.PathJoin(workDir, ".claude", "mcp.json")
}
args = append(args, "--mcp-config", mcpPath)
if r.mcpToolPattern != "" {
args = append(args, "--allowedTools", r.mcpToolPattern)
}
}
return args
}
func buildSandboxEnvPrompt(p platform, workDir string) string {
osName := p.OS()
if osName == "" {
osName = "linux"
}
shell := p.Shell()
if shell == "" {
shell = "bash"
}
shellNote := p.EnvPromptNote()
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`, workDir, osName, shell, shellNote)
}
func hasExistingSession(ctx context.Context, computer infra.Computer, p platform, assistantID string) bool {
workDir := computer.GetWorkDir()
var sessionDir string
if assistantID != "" {
configDir := p.PathJoin(workDir, ".yao", "assistants", assistantID)
sessionDir = p.PathJoin(configDir, "projects")
} else {
sessionDir = p.PathJoin(workDir, ".claude", "projects")
}
result, err := computer.Exec(ctx, p.ListDirCmd(sessionDir))
if err != nil || result.ExitCode != 0 {
return false
}
return strings.TrimSpace(result.Stdout) != ""
}
func buildMCPConfig(servers []types.MCPServer) []byte {
mcpServers := make(map[string]any, len(servers))
for _, s := range servers {
name := s.ServerID
if name == "" {
continue
}
mcpServers[name] = map[string]any{
"command": "tai",
"args": []string{"mcp", name},
}
}
if len(mcpServers) == 0 {
mcpServers["yao"] = map[string]any{
"command": "tai",
"args": []string{"mcp"},
}
}
config := map[string]any{"mcpServers": mcpServers}
data, _ := json.Marshal(config)
return data
}
func buildMCPAllowedTools(servers []types.MCPServer) string {
patterns := make([]string, 0, len(servers))
for _, s := range servers {
if s.ServerID != "" {
patterns = append(patterns, fmt.Sprintf("mcp__%s__*", s.ServerID))
}
}
if len(patterns) == 0 {
return "mcp__yao__*"
}
return strings.Join(patterns, ",")
}
func buildLastUserMessageJSONL(messages []agentContext.Message) string {
for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Role == "user" {
content := messages[i].Content
if content == nil {
content = ""
}
msg := map[string]any{
"type": "user",
"message": map[string]any{
"role": "user",
"content": content,
},
}
data, _ := json.Marshal(msg)
return string(data)
}
}
return ""
}
// claudeRoleEnvMap maps abstract Yao model roles to Claude CLI environment
// variables. Only roles with matching Claude CLI env vars are listed here.
// ANTHROPIC_DEFAULT_SONNET_MODEL is set to the primary model in buildEnv.
var claudeRoleEnvMap = map[string]struct{ EnvVar string }{
"default": {EnvVar: "ANTHROPIC_MODEL"},
"heavy": {EnvVar: "ANTHROPIC_DEFAULT_OPUS_MODEL"},
"light": {EnvVar: "ANTHROPIC_DEFAULT_HAIKU_MODEL"},
}
func connectorHost(c connector.Connector) string {
if c == nil {
return ""
}
if lc, ok := c.(goullm.LLMConnector); ok {
if u := lc.GetURL(); u != "" {
return u
}
}
host, _ := c.Setting()["host"].(string)
return host
}
func connectorProtocols(c connector.Connector) []string {
if c == nil {
return nil
}
setting := c.Setting()
if ps, ok := setting["protocols"].([]string); ok && len(ps) > 0 {
return ps
}
if c.Is(connector.ANTHROPIC) {
return []string{"anthropic"}
}
return []string{"openai"}
}
func supportsProtocol(c connector.Connector, proto string) bool {
for _, p := range connectorProtocols(c) {
if p == proto {
return true
}
}
return false
}
var claudeArgWhitelist = map[string]string{
"max_turns": "--max-turns",
"disallowed_tools": "--disallowed-tools",
"allowed_tools": "--allowedTools",
}
func isStandardAnthropicModel(model string) bool {
return strings.HasPrefix(model, "claude-") || strings.HasPrefix(model, "anthropic.")
}
func buildClaudeCodeCapabilities(conn connector.Connector) string {
if conn == nil {
return ""
}
setting := conn.Setting()
if setting == nil {
return ""
}
var caps []string
if thinking, ok := setting["thinking"].(map[string]interface{}); ok {
if thinkType, _ := thinking["type"].(string); thinkType == "enabled" {
caps = append(caps, "thinking")
}
}
return strings.Join(caps, ",")
}
func setAnthropicModelEnv(env map[string]string, host, key, model string, conn connector.Connector) {
env["ANTHROPIC_BASE_URL"] = host
env["ANTHROPIC_API_KEY"] = key
if model == "" {
return
}
env["ANTHROPIC_MODEL"] = model
env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model
env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model
if isStandardAnthropicModel(model) {
return
}
caps := buildClaudeCodeCapabilities(conn)
env["ANTHROPIC_CUSTOM_MODEL_OPTION"] = model
env["ANTHROPIC_CUSTOM_MODEL_OPTION_NAME"] = model
env["ANTHROPIC_DEFAULT_OPUS_MODEL_NAME"] = model
env["ANTHROPIC_DEFAULT_SONNET_MODEL_NAME"] = model
env["ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME"] = model
if caps != "" {
env["ANTHROPIC_CUSTOM_MODEL_OPTION_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES"] = caps
}
}
func applyAnthropicRoleOverrides(
env map[string]string,
primaryHost string,
roles map[string]connector.Connector,
) {
for role, rm := range claudeRoleEnvMap {
if role == "default" {
continue
}
rc, ok := roles[role]
if !ok || rc == nil {
continue
}
roleHost := connectorHost(rc)
if roleHost != primaryHost {
log.Warn("[claude] role %s: host mismatch (%s != %s), falling back to primary", role, roleHost, primaryHost)
continue
}
if !supportsProtocol(rc, "anthropic") {
log.Warn("[claude] role %s: not anthropic protocol, falling back to primary", role)
continue
}
rcModel, _ := rc.Setting()["model"].(string)
if rcModel == "" {
continue
}
env[rm.EnvVar] = rcModel
if isStandardAnthropicModel(rcModel) {
continue
}
env[rm.EnvVar+"_NAME"] = rcModel
if caps := buildClaudeCodeCapabilities(rc); caps != "" {
env[rm.EnvVar+"_SUPPORTED_CAPABILITIES"] = caps
}
}
}
func setA2OModelEnv(env map[string]string, connectorID, model string, conn connector.Connector) {
env["ANTHROPIC_BASE_URL"] = fmt.Sprintf("http://127.0.0.1:%d/c/%s", defaultA2OPort, connectorID)
env["ANTHROPIC_API_KEY"] = "dummy"
env["ANTHROPIC_MODEL"] = model
env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model
env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model
if !isStandardAnthropicModel(model) {
env["ANTHROPIC_CUSTOM_MODEL_OPTION"] = model
env["ANTHROPIC_CUSTOM_MODEL_OPTION_NAME"] = model
env["ANTHROPIC_DEFAULT_OPUS_MODEL_NAME"] = model
env["ANTHROPIC_DEFAULT_SONNET_MODEL_NAME"] = model
env["ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME"] = model
if caps := buildClaudeCodeCapabilities(conn); caps != "" {
env["ANTHROPIC_CUSTOM_MODEL_OPTION_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES"] = caps
env["ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES"] = caps
}
}
}
func applyA2ORoleOverrides(
env map[string]string,
roles map[string]connector.Connector,
) {
for role, rm := range claudeRoleEnvMap {
if role == "default" {
continue
}
rc, ok := roles[role]
if !ok || rc == nil {
continue
}
var rcModel string
if lc, ok := rc.(goullm.LLMConnector); ok {
rcModel = lc.GetModel()
}
if rcModel == "" {
rcModel, _ = rc.Setting()["model"].(string)
}
if rcModel == "" {
continue
}
env[rm.EnvVar] = rcModel
if !isStandardAnthropicModel(rcModel) {
env[rm.EnvVar+"_NAME"] = rcModel
if caps := buildClaudeCodeCapabilities(rc); caps != "" {
env[rm.EnvVar+"_SUPPORTED_CAPABILITIES"] = caps
}
}
}
}