yao/agent/sandbox/v2/claude/runner.go

package claude

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/yaoapp/gou/connector"
	goullm "github.com/yaoapp/gou/llm"
	"github.com/yaoapp/kun/log"
	agentContext "github.com/yaoapp/yao/agent/context"
	"github.com/yaoapp/yao/agent/output/message"
	"github.com/yaoapp/yao/agent/sandbox/v2/types"
	infra "github.com/yaoapp/yao/sandbox/v2"
)

// Runner implements the sandbox Runner interface for Claude CLI (mode=cli).
type Runner struct {
	mode           string
	hasMCP         bool
	mcpToolPattern string
	lastCompleted  bool
	lastChatID     string
	logger         *agentContext.RequestLogger
}

// New creates a new Runner.
func New() *Runner {
	return &Runner{mode: "cli"}
}

// Name returns the runner identifier.
func (r *Runner) Name() string { return "claude" }

// Prepare executes user-defined and runner-specific prepare steps.
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.AssistantID
	prefix := ".yao/assistants/" + assistantID
	if assistantID == "" {
		prefix = ".claude"
	}

	steps := append([]types.PrepareStep{}, req.Config.Prepare...)

	if req.SkillsDir != "" {
		ws := req.Computer.Workplace()
		if ws != nil {
			src := "local:///" + req.SkillsDir
			dst := prefix + "/skills"
			if _, err := ws.Copy(src, dst); err != nil {
				r.logger.Warn("copy skills %s -> %s: %v", src, dst, err)
			}
		}
	}

	if len(req.MCPServers) > 0 {
		r.hasMCP = true
		r.mcpToolPattern = buildMCPAllowedTools(req.MCPServers)
		mcpJSON := buildMCPConfig(req.MCPServers)
		steps = append(steps, types.PrepareStep{
			Action:  "file",
			Path:    prefix + "/mcp.json",
			Content: mcpJSON,
		})
	}

	if req.RunSteps != nil && len(steps) > 0 {
		if err := req.RunSteps(ctx, steps, req.Computer, req.AssistantID, req.ConfigHash, req.AssistantDir); err != nil {
			return fmt.Errorf("claude prepare steps: %w", err)
		}
	}

	return nil
}

// Stream executes the Claude CLI and streams output to handler.
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")
	}

	p := resolvePlatform(computer)

	// Inject connector config into a2o proxy (best-effort, errors ignored).
	if req.Connector != nil && req.Connector.Is(connector.OPENAI) {
		roleConnectors := resolveAllRoleConnectors(req)
		if len(roleConnectors) > 0 {
			injectA2OConfigWithRoutes(ctx, computer, req.Connector, roleConnectors)
		} else {
			injectA2OConfig(ctx, computer, req.Connector)
		}
	}

	if req.ChatID != "" {
		if ws := computer.Workplace(); ws != nil {
			processed, err := prepareAttachments(ctx, req.Messages, req.ChatID, ws)
			if err != nil {
				return fmt.Errorf("prepareAttachments: %w", err)
			}
			req.Messages = processed
		}
	}

	cmd := r.buildCommand(ctx, req, p)

	r.logger = req.Logger
	if r.logger == nil {
		r.logger = agentContext.NoopLogger()
	}

	chatID := req.ChatID
	r.lastChatID = chatID
	assistantID := req.AssistantID

	log.Trace("[claude-runner] Stream started: assistantID=%s chatID=%s promptLen=%d", assistantID, chatID, len(cmd.shell))
	r.logger.Debug("env vars passed to session (%d total):", len(cmd.env))
	for k, v := range cmd.env {
		if strings.HasPrefix(k, "CTX_") || k == "CLAUDE_CONFIG_DIR" || k == "HOME" || k == "WORKDIR" {
			r.logger.Debug("  %s=%s", k, v)
		} else {
			r.logger.Debug("  %s=(set, len=%d)", k, len(v))
		}
	}

	sess, err := startSession(ctx, computer, p, cmd, chatID, r.logger)
	if err != nil {
		return err
	}

	// Mark session in store immediately after CLI process starts.
	// Claude CLI creates session files on disk at startup, so subsequent
	// requests must use --resume (not --session-id) even if this stream fails.
	if chatID != "" {
		storeKey := "claude-session:" + assistantID + ":" + chatID
		sessionUUID := chatIDToSessionUUID(assistantID, chatID)
		markChatSession(storeKey, sessionUUID, 90*24*time.Hour)
	}

	streamStart := time.Now()
	completed, err := sess.runStream(handler)
	r.lastCompleted = completed
	elapsed := time.Since(streamStart).Round(time.Second)
	log.Trace("[claude-runner] Stream finished: assistantID=%s chatID=%s completed=%v elapsed=%v err=%v", assistantID, chatID, completed, elapsed, err)
	r.logger.Debug("Stream: runStream returned completed=%v err=%v elapsed=%v", completed, err, elapsed)
	if completed {
		sess.shutdown()
	}
	return err
}

// Cleanup kills any remaining claude processes. If the stream completed
// normally (received "result"), child processes are preserved.
func (r *Runner) Cleanup(ctx context.Context, computer infra.Computer) error {
	if computer == nil {
		return nil
	}

	log.Trace("[claude-runner] Cleanup: chatID=%s lastCompleted=%v", r.lastChatID, r.lastCompleted)

	if r.lastCompleted {
		if r.logger != nil {
			r.logger.Info("cleanup: stream completed normally, preserving child processes")
		}
		return nil
	}

	if r.mode != "service" {
		p := resolvePlatform(computer)
		if r.lastChatID != "" {
			computer.Exec(ctx, p.KillSessionCmd(sanitizeSessionName(r.lastChatID)))
		} else {
			computer.Exec(ctx, p.KillCmd("claude"))
		}
	}

	return nil
}

type a2oConnectorConfig struct {
	Backend         string                         `json:"backend"`
	Model           string                         `json:"model"`
	APIKey          string                         `json:"api_key"`
	AuthMode        string                         `json:"auth_mode,omitempty"`
	MaxOutputTokens int                            `json:"max_output_tokens,omitempty"`
	Options         map[string]interface{}         `json:"options,omitempty"`
	Routes          map[string]*a2oConnectorConfig `json:"routes,omitempty"`
}

func buildSingleA2OConfig(conn connector.Connector) *a2oConnectorConfig {
	settings := conn.Setting()
	if settings == nil {
		return nil
	}

	cfg := &a2oConnectorConfig{}

	// Extract standard fields via LLMConnector methods when available
	if lc, ok := conn.(goullm.LLMConnector); ok {
		if url := lc.GetURL(); url != "" {
			cfg.Backend = connector.BuildAPIURL(url, "/chat/completions")
		}
		cfg.Model = lc.GetModel()
		cfg.APIKey = lc.GetKey()
		cfg.AuthMode = string(lc.GetAuthMode())
		if caps := lc.GetCapabilities(); caps != nil && caps.MaxOutputTokens > 0 {
			cfg.MaxOutputTokens = caps.MaxOutputTokens
		}
	} else {
		if host, ok := settings["host"].(string); ok && host != "" {
			cfg.Backend = connector.BuildAPIURL(host, "/chat/completions")
		} else if proxy, ok := settings["proxy"].(string); ok && proxy != "" {
			cfg.Backend = connector.BuildAPIURL(proxy, "/chat/completions")
		}
		if model, ok := settings["model"].(string); ok && model != "" {
			cfg.Model = model
		}
		if key, ok := settings["key"].(string); ok && key != "" {
			cfg.APIKey = key
		}
	}

	// Whitelist-filter remaining settings for the options field
	extra := connector.FilterRequestBodyParams(settings, conn)
	if len(extra) > 0 {
		cfg.Options = extra
	}

	if cfg.Backend == "" {
		return nil
	}
	return cfg
}

// resolveAllRoleConnectors resolves all declared role connectors and returns
// a map of virtual model name -> connector for roles that have independent connectors.
func resolveAllRoleConnectors(req *types.StreamRequest) map[string]connector.Connector {
	roleConns := getRoleConnectors(req)
	if len(roleConns) == 0 {
		return nil
	}

	result := make(map[string]connector.Connector)
	for role, rm := range claudeRoleEnvMap {
		if role == "primary" {
			continue
		}
		rc := resolveRoleConnector(role, roleConns, req.UserExplicit, func(id string) connector.Connector {
			c, _ := connector.Connectors[id]
			return c
		})
		if rc == nil {
			continue
		}
		result[rm.ModelName] = rc
	}
	return result
}

// injectA2OConfigWithRoutes pushes the primary connector config along with
// role-based routes to the a2o proxy. Each route maps a virtual model name
// to a different backend connector.
func injectA2OConfigWithRoutes(ctx context.Context, computer infra.Computer, primaryConn connector.Connector, roleConnectors map[string]connector.Connector) {
	primaryCfg := buildSingleA2OConfig(primaryConn)
	if primaryCfg == nil {
		log.Trace("[claude] injectA2OConfigWithRoutes: no valid primary config for connector %s", primaryConn.ID())
		return
	}

	routes := make(map[string]*a2oConnectorConfig, len(roleConnectors))
	for modelName, rc := range roleConnectors {
		routeCfg := buildSingleA2OConfig(rc)
		if routeCfg != nil {
			routes[modelName] = routeCfg
		}
	}
	primaryCfg.Routes = routes

	data, err := json.Marshal(primaryCfg)
	if err != nil {
		log.Trace("[claude] injectA2OConfigWithRoutes: marshal error: %v", err)
		return
	}

	connID := primaryConn.ID()
	var result *infra.ExecResult

	info := computer.ComputerInfo()
	if info.Kind == "host" {
		result, err = computer.Exec(ctx, []string{"tai", "a2o", "config", "put", connID}, infra.WithStdin(data))
	} else {
		escaped := strings.ReplaceAll(string(data), "'", "'\\''")
		script := fmt.Sprintf("echo '%s' | tai a2o config put %s", escaped, connID)
		result, err = computer.Exec(ctx, []string{"sh", "-c", script})
	}

	if err != nil {
		log.Trace("[claude] injectA2OConfigWithRoutes: exec error (ignored): %v", err)
		return
	}
	if result.ExitCode != 0 {
		log.Trace("[claude] injectA2OConfigWithRoutes: exit %d stderr=%s (ignored)", result.ExitCode, result.Stderr)
		return
	}

	log.Trace("[claude] injectA2OConfigWithRoutes: connector=%s injected with %d routes", connID, len(routes))
}

// injectA2OConfig pushes the connector config to the a2o proxy.
// For box (Linux container): uses sh pipe since Docker exec stdin may not work.
// For host: uses WithStdin which works reliably on all platforms.
// Best-effort: errors are logged and ignored.
func injectA2OConfig(ctx context.Context, computer infra.Computer, conn connector.Connector) {
	cfg := buildSingleA2OConfig(conn)
	if cfg == nil {
		log.Trace("[claude] injectA2OConfig: no valid config for connector %s", conn.ID())
		return
	}

	data, err := json.Marshal(cfg)
	if err != nil {
		log.Trace("[claude] injectA2OConfig: marshal error: %v", err)
		return
	}

	connID := conn.ID()
	var result *infra.ExecResult

	info := computer.ComputerInfo()
	if info.Kind == "host" {
		result, err = computer.Exec(ctx, []string{"tai", "a2o", "config", "put", connID}, infra.WithStdin(data))
	} else {
		escaped := strings.ReplaceAll(string(data), "'", "'\\''")
		script := fmt.Sprintf("echo '%s' | tai a2o config put %s", escaped, connID)
		result, err = computer.Exec(ctx, []string{"sh", "-c", script})
	}

	if err != nil {
		log.Trace("[claude] injectA2OConfig: exec error (ignored): %v", err)
		return
	}
	if result.ExitCode != 0 {
		log.Trace("[claude] injectA2OConfig: exit %d stderr=%s (ignored)", result.ExitCode, result.Stderr)
		return
	}

	log.Trace("[claude] injectA2OConfig: connector=%s injected ok", connID)
}