From 2084e204796e8080dfa71542a38f2ff3522d12c8 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 13 Apr 2026 14:47:08 +0800 Subject: [PATCH] feat(workspace): add workspace ID retrieval and integration into environment configuration - Implemented GetWorkspaceID method in the context package to retrieve the workspace ID. - Updated buildContextVariables to include WORKSPACE_ID in the context variables if available. - Enhanced buildEnv function to set CTX_WORKSPACE_ID in the environment configuration based on the workspace ID. - Refactored role connector handling to support multiple connectors in the environment setup. - Added tests for new role connector functionality and workspace ID integration. --- agent/assistant/build.go | 5 + agent/assistant/sandbox_v2.go | 1 + agent/context/context.go | 8 + agent/sandbox/v2/claude/command.go | 122 ++++++++- agent/sandbox/v2/claude/command_test.go | 347 ++++++++++++++++++++++++ agent/sandbox/v2/claude/runner.go | 90 +++++- agent/sandbox/v2/types/config.go | 36 ++- agent/sandbox/v2/types/config_test.go | 94 +++++++ agent/sandbox/v2/types/runner.go | 1 + tai/workspace/workspace.go | 7 + workspace/manager.go | 14 +- 11 files changed, 708 insertions(+), 17 deletions(-) create mode 100644 agent/sandbox/v2/types/config_test.go diff --git a/agent/assistant/build.go b/agent/assistant/build.go index 260b22ba..e7d7f35c 100644 --- a/agent/assistant/build.go +++ b/agent/assistant/build.go @@ -203,6 +203,11 @@ func (ast *Assistant) buildContextVariables(ctx *context.Context) map[string]str } } + // Workspace info + if workspaceID, err := ctx.GetWorkspaceID(); err == nil { + vars["WORKSPACE_ID"] = workspaceID + } + if ctx == nil { return vars } diff --git a/agent/assistant/sandbox_v2.go b/agent/assistant/sandbox_v2.go index d01e33ba..28c7cf45 100644 --- a/agent/assistant/sandbox_v2.go +++ b/agent/assistant/sandbox_v2.go @@ -214,6 +214,7 @@ func (ast *Assistant) executeSandboxV2Stream( ChatID: ctx.ChatID, Token: tok, Logger: ctx.Logger, + UserExplicit: p.Options != nil && p.Options.Connector != "", } execReq := &sandboxv2.ExecuteRequest{ diff --git a/agent/context/context.go b/agent/context/context.go index 05582029..3926c369 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -559,3 +559,11 @@ func (ctx *Context) MergeMetadata(metadata map[string]interface{}) { ctx.Metadata[k] = v } } + +// GetWorkspaceID returns the ID of the workspace +func (ctx *Context) GetWorkspaceID() (string, error) { + if ctx.workspace == nil { + return "", nil + } + return ctx.workspace.GetID() +} diff --git a/agent/sandbox/v2/claude/command.go b/agent/sandbox/v2/claude/command.go index 6e75d674..0ff2d323 100644 --- a/agent/sandbox/v2/claude/command.go +++ b/agent/sandbox/v2/claude/command.go @@ -105,6 +105,15 @@ 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 } @@ -127,6 +136,12 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string { key, _ := setting["key"].(string) model, _ := setting["model"].(string) + roleConnectors := getRoleConnectors(req) + getConn := func(id string) connector.Connector { + c, _ := connector.Connectors[id] + return c + } + if req.Connector.Is(connector.ANTHROPIC) { env["ANTHROPIC_BASE_URL"] = host env["ANTHROPIC_API_KEY"] = key @@ -137,18 +152,48 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string { env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model env["CLAUDE_CODE_SUBAGENT_MODEL"] = model } + + if len(roleConnectors) > 0 { + primaryHost := host + for role, rm := range claudeRoleEnvMap { + if role == "primary" { + continue + } + rc := resolveRoleConnector(role, roleConnectors, req.UserExplicit, getConn) + if rc == nil { + continue + } + rcHost := connectorHost(rc) + if rcHost == primaryHost && supportsProtocol(rc, "anthropic") { + rcModel, _ := rc.Setting()["model"].(string) + if rcModel != "" { + env[rm.EnvVar] = rcModel + } + } + } + } } else { connectorID := req.Connector.ID() env["ANTHROPIC_BASE_URL"] = fmt.Sprintf("http://127.0.0.1:%d/c/%s", defaultA2OPort, connectorID) env["ANTHROPIC_API_KEY"] = "dummy" - // Use a valid Anthropic model name to pass Claude CLI's local - // validation. The a2o proxy ignores this and substitutes the - // real backend model from its connector config. 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" + + if len(roleConnectors) > 0 { + for role, rm := range claudeRoleEnvMap { + if role == "primary" { + continue + } + rc := resolveRoleConnector(role, roleConnectors, req.UserExplicit, getConn) + if rc == nil { + continue + } + env[rm.EnvVar] = rm.ModelName + } + } } if thinking, ok := setting["thinking"].(map[string]interface{}); ok { @@ -344,6 +389,77 @@ func buildLastUserMessageJSONL(messages []agentContext.Message) string { return "" } +// claudeRoleEnvMap maps abstract Yao model roles to Claude CLI environment +// variables and virtual model name identifiers used as A2O route keys. +// ModelName uniqueness is only required among roles that have independent +// connectors (i.e. are added to the A2O routes map). +var claudeRoleEnvMap = map[string]struct { + EnvVar string + ModelName string +}{ + "primary": {EnvVar: "ANTHROPIC_MODEL", ModelName: "claude-sonnet-4-6"}, + "heavy": {EnvVar: "ANTHROPIC_DEFAULT_OPUS_MODEL", ModelName: "claude-opus-4-6"}, + "light": {EnvVar: "ANTHROPIC_DEFAULT_HAIKU_MODEL", ModelName: "claude-haiku-4-5"}, + "subagent": {EnvVar: "CLAUDE_CODE_SUBAGENT_MODEL", ModelName: "claude-subagent-4-6"}, + "vision": {EnvVar: "ANTHROPIC_DEFAULT_SONNET_MODEL", ModelName: "claude-vision-4-5"}, +} + +func connectorHost(c connector.Connector) string { + if c == nil { + return "" + } + 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 +} + +// resolveRoleConnector determines which connector to use for a given role. +// Returns nil when the role should use the primary connector (caller decides). +func resolveRoleConnector( + role string, + roleConnectors map[string]*types.RoleConnector, + userExplicit bool, + getConnector func(id string) connector.Connector, +) connector.Connector { + rc, ok := roleConnectors[role] + if !ok || rc == nil { + return nil + } + if rc.Override == "user" && userExplicit { + return nil + } + return getConnector(rc.Connector) +} + +func getRoleConnectors(req *types.StreamRequest) map[string]*types.RoleConnector { + if req.Config == nil { + return nil + } + return req.Config.Runner.Connectors +} + var claudeArgWhitelist = map[string]string{ "max_turns": "--max-turns", "disallowed_tools": "--disallowed-tools", diff --git a/agent/sandbox/v2/claude/command_test.go b/agent/sandbox/v2/claude/command_test.go index 5500cf6d..8b4d8d54 100644 --- a/agent/sandbox/v2/claude/command_test.go +++ b/agent/sandbox/v2/claude/command_test.go @@ -8,6 +8,10 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/connector" + gouTypes "github.com/yaoapp/gou/types" + "github.com/yaoapp/xun/dbal/query" + "github.com/yaoapp/xun/dbal/schema" agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/sandbox/v2/types" infra "github.com/yaoapp/yao/sandbox/v2" @@ -434,3 +438,346 @@ func TestBuildArgs_EmptyChatID_Continuation(t *testing.T) { assert.NotContains(t, args, "--session-id") assert.NotContains(t, args, "--name") } + +// --- fakeConnector implements connector.Connector for unit tests --- + +type fakeConnector struct { + id string + typ int + settings map[string]interface{} +} + +func (f *fakeConnector) Register(string, string, []byte) error { return nil } +func (f *fakeConnector) Query() (query.Query, error) { return nil, nil } +func (f *fakeConnector) Schema() (schema.Schema, error) { return nil, nil } +func (f *fakeConnector) Close() error { return nil } +func (f *fakeConnector) ID() string { return f.id } +func (f *fakeConnector) Is(t int) bool { return f.typ == t } +func (f *fakeConnector) Setting() map[string]interface{} { return f.settings } +func (f *fakeConnector) GetMetaInfo() gouTypes.MetaInfo { return gouTypes.MetaInfo{} } + +func newOpenAIConnector(id, host, model, key string) *fakeConnector { + return &fakeConnector{ + id: id, + typ: connector.OPENAI, + settings: map[string]interface{}{ + "host": host, + "model": model, + "key": key, + }, + } +} + +func newAnthropicConnector(id, host, model, key string) *fakeConnector { + return &fakeConnector{ + id: id, + typ: connector.ANTHROPIC, + settings: map[string]interface{}{ + "host": host, + "model": model, + "key": key, + }, + } +} + +func newDualProtoConnector(id, host, model, key string) *fakeConnector { + return &fakeConnector{ + id: id, + typ: connector.OPENAI, + settings: map[string]interface{}{ + "host": host, + "model": model, + "key": key, + "protocols": []string{"openai", "anthropic"}, + }, + } +} + +// --- helper functions tests --- + +func TestConnectorHost(t *testing.T) { + c := newOpenAIConnector("test", "https://api.openai.com", "gpt-4", "k") + assert.Equal(t, "https://api.openai.com", connectorHost(c)) + assert.Equal(t, "", connectorHost(nil)) +} + +func TestConnectorProtocols(t *testing.T) { + oai := newOpenAIConnector("oai", "https://api.openai.com", "gpt-4", "k") + assert.Equal(t, []string{"openai"}, connectorProtocols(oai)) + + anth := newAnthropicConnector("anth", "https://api.anthropic.com", "claude", "k") + assert.Equal(t, []string{"anthropic"}, connectorProtocols(anth)) + + dual := newDualProtoConnector("dual", "https://api.yao.run", "model", "k") + assert.Equal(t, []string{"openai", "anthropic"}, connectorProtocols(dual)) + + assert.Nil(t, connectorProtocols(nil)) +} + +func TestSupportsProtocol(t *testing.T) { + dual := newDualProtoConnector("dual", "https://api.yao.run", "model", "k") + assert.True(t, supportsProtocol(dual, "anthropic")) + assert.True(t, supportsProtocol(dual, "openai")) + assert.False(t, supportsProtocol(dual, "grpc")) + + oai := newOpenAIConnector("oai", "https://api.openai.com", "gpt-4", "k") + assert.False(t, supportsProtocol(oai, "anthropic")) + assert.True(t, supportsProtocol(oai, "openai")) +} + +func TestResolveRoleConnector_Undeclared(t *testing.T) { + roles := map[string]*types.RoleConnector{} + result := resolveRoleConnector("heavy", roles, false, func(id string) connector.Connector { return nil }) + assert.Nil(t, result) +} + +func TestResolveRoleConnector_Force(t *testing.T) { + heavyConn := newOpenAIConnector("thinking", "https://api.thinking.com", "think-model", "k") + roles := map[string]*types.RoleConnector{ + "heavy": {Connector: "thinking", Override: "force"}, + } + result := resolveRoleConnector("heavy", roles, true, func(id string) connector.Connector { + if id == "thinking" { + return heavyConn + } + return nil + }) + assert.Equal(t, heavyConn, result) +} + +func TestResolveRoleConnector_UserExplicit(t *testing.T) { + roles := map[string]*types.RoleConnector{ + "heavy": {Connector: "thinking", Override: "user"}, + } + result := resolveRoleConnector("heavy", roles, true, func(id string) connector.Connector { + return newOpenAIConnector("thinking", "h", "m", "k") + }) + assert.Nil(t, result, "override=user + userExplicit=true => use user's connector") +} + +func TestResolveRoleConnector_UserNotExplicit(t *testing.T) { + heavyConn := newOpenAIConnector("thinking", "h", "m", "k") + roles := map[string]*types.RoleConnector{ + "heavy": {Connector: "thinking", Override: "user"}, + } + result := resolveRoleConnector("heavy", roles, false, func(id string) connector.Connector { + if id == "thinking" { + return heavyConn + } + return nil + }) + assert.Equal(t, heavyConn, result, "override=user + userExplicit=false => use sandbox connector") +} + +// --- buildEnv with multi-connector --- + +func registerTestConnectors(t *testing.T, connectors map[string]connector.Connector) func() { + t.Helper() + for id, c := range connectors { + connector.Connectors[id] = c + } + return func() { + for id := range connectors { + delete(connector.Connectors, id) + } + } +} + +func TestBuildEnv_OpenAI_SingleConnector(t *testing.T) { + oai := newOpenAIConnector("kimi", "https://api.moonshot.cn", "kimi-k2.5", "sk-test") + req := &types.StreamRequest{ + Config: &types.SandboxConfig{}, + Connector: oai, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Contains(t, env["ANTHROPIC_BASE_URL"], "127.0.0.1") + assert.Contains(t, env["ANTHROPIC_BASE_URL"], "kimi") + assert.Equal(t, "claude-sonnet-4-6", env["ANTHROPIC_MODEL"]) + assert.Equal(t, "claude-sonnet-4-6", env["ANTHROPIC_DEFAULT_OPUS_MODEL"]) +} + +func TestBuildEnv_OpenAI_MultiConnector(t *testing.T) { + primary := newOpenAIConnector("kimi", "https://api.moonshot.cn", "kimi-k2.5", "sk-test") + vision := newOpenAIConnector("vision-conn", "https://api.vision.com", "vis-model", "sk-v") + + cleanup := registerTestConnectors(t, map[string]connector.Connector{ + "vision-conn": vision, + }) + defer cleanup() + + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Connectors: map[string]*types.RoleConnector{ + "vision": {Connector: "vision-conn", Override: "force"}, + }, + }, + }, + Connector: primary, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "claude-vision-4-5", env["ANTHROPIC_DEFAULT_SONNET_MODEL"], + "vision role should get its virtual model name for A2O routing") + assert.Equal(t, "claude-sonnet-4-6", env["ANTHROPIC_MODEL"], + "primary should keep default virtual model") +} + +func TestBuildEnv_Anthropic_SingleConnector(t *testing.T) { + anth := newAnthropicConnector("claude", "https://api.anthropic.com", "claude-sonnet-4-20250514", "sk-ant-test") + req := &types.StreamRequest{ + Config: &types.SandboxConfig{}, + Connector: anth, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "https://api.anthropic.com", env["ANTHROPIC_BASE_URL"]) + assert.Equal(t, "sk-ant-test", env["ANTHROPIC_API_KEY"]) + assert.Equal(t, "claude-sonnet-4-20250514", env["ANTHROPIC_MODEL"]) +} + +func TestBuildEnv_Anthropic_MultiConnector_Compatible(t *testing.T) { + primary := newAnthropicConnector("claude", "https://api.yao.run", "claude-sonnet-4-20250514", "sk-ant") + lightConn := newDualProtoConnector("light-conn", "https://api.yao.run", "claude-haiku-3-5-20241022", "sk-light") + + cleanup := registerTestConnectors(t, map[string]connector.Connector{ + "light-conn": lightConn, + }) + defer cleanup() + + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Connectors: map[string]*types.RoleConnector{ + "light": {Connector: "light-conn", Override: "force"}, + }, + }, + }, + Connector: primary, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "claude-haiku-3-5-20241022", env["ANTHROPIC_DEFAULT_HAIKU_MODEL"], + "compatible connector: use real model name") + assert.Equal(t, "https://api.yao.run", env["ANTHROPIC_BASE_URL"], + "primary base URL unchanged") +} + +// --- buildSingleA2OConfig + injectA2OConfigWithRoutes --- + +func TestBuildSingleA2OConfig_Basic(t *testing.T) { + conn := newOpenAIConnector("test", "https://api.openai.com", "gpt-4", "sk-test") + cfg := buildSingleA2OConfig(conn) + require.NotNil(t, cfg) + assert.Contains(t, cfg.Backend, "api.openai.com") + assert.Contains(t, cfg.Backend, "chat/completions") + assert.Equal(t, "gpt-4", cfg.Model) + assert.Equal(t, "sk-test", cfg.APIKey) +} + +func TestBuildSingleA2OConfig_Nil(t *testing.T) { + conn := &fakeConnector{id: "empty", typ: connector.OPENAI, settings: map[string]interface{}{}} + cfg := buildSingleA2OConfig(conn) + assert.Nil(t, cfg, "no host => nil config") +} + +func TestInjectA2OConfigWithRoutes_BuildsCorrectJSON(t *testing.T) { + primary := newOpenAIConnector("kimi", "https://api.moonshot.cn", "kimi-k2.5", "sk-kimi") + vision := newOpenAIConnector("vision", "https://api.vision.com", "vis-model", "sk-v") + + roleConnectors := map[string]connector.Connector{ + "claude-vision-4-5": vision, + } + + primaryCfg := buildSingleA2OConfig(primary) + require.NotNil(t, primaryCfg) + + 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) + require.NoError(t, err) + + var parsed map[string]interface{} + require.NoError(t, json.Unmarshal(data, &parsed)) + + routesMap, ok := parsed["routes"].(map[string]interface{}) + require.True(t, ok, "routes should be present in JSON") + assert.Len(t, routesMap, 1) + + visionRoute, ok := routesMap["claude-vision-4-5"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "vis-model", visionRoute["model"]) + assert.Contains(t, visionRoute["backend"], "api.vision.com") +} + +func TestResolveAllRoleConnectors_Empty(t *testing.T) { + req := &types.StreamRequest{ + Config: &types.SandboxConfig{}, + Connector: newOpenAIConnector("test", "h", "m", "k"), + } + result := resolveAllRoleConnectors(req) + assert.Nil(t, result) +} + +func TestResolveAllRoleConnectors_WithRoles(t *testing.T) { + vision := newOpenAIConnector("vis", "https://vis.com", "vis-m", "sk") + cleanup := registerTestConnectors(t, map[string]connector.Connector{"vis": vision}) + defer cleanup() + + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Connectors: map[string]*types.RoleConnector{ + "vision": {Connector: "vis", Override: "force"}, + }, + }, + }, + Connector: newOpenAIConnector("primary", "h", "m", "k"), + } + result := resolveAllRoleConnectors(req) + assert.Len(t, result, 1) + assert.Equal(t, vision, result["claude-vision-4-5"]) +} + +func TestBuildEnv_Anthropic_MultiConnector_Incompatible(t *testing.T) { + primary := newAnthropicConnector("claude", "https://api.anthropic.com", "claude-sonnet-4-20250514", "sk-ant") + visionConn := newOpenAIConnector("vision-oai", "https://api.openai.com", "gpt-4o", "sk-oai") + + cleanup := registerTestConnectors(t, map[string]connector.Connector{ + "vision-oai": visionConn, + }) + defer cleanup() + + req := &types.StreamRequest{ + Config: &types.SandboxConfig{ + Runner: types.RunnerConfig{ + Connectors: map[string]*types.RoleConnector{ + "vision": {Connector: "vision-oai", Override: "force"}, + }, + }, + }, + Connector: primary, + } + req.Computer = newFakeComputer("/workspace") + p := testPlatform() + + env := buildEnv(req, p) + assert.Equal(t, "claude-sonnet-4-20250514", env["ANTHROPIC_DEFAULT_SONNET_MODEL"], + "incompatible connector: vision should keep primary model (different host, no anthropic protocol)") +} diff --git a/agent/sandbox/v2/claude/runner.go b/agent/sandbox/v2/claude/runner.go index d57f4ba7..944eab3e 100644 --- a/agent/sandbox/v2/claude/runner.go +++ b/agent/sandbox/v2/claude/runner.go @@ -90,7 +90,12 @@ func (r *Runner) Stream(ctx context.Context, req *types.StreamRequest, handler m // Inject connector config into a2o proxy (best-effort, errors ignored). if req.Connector != nil && req.Connector.Is(connector.OPENAI) { - injectA2OConfig(ctx, computer, req.Connector) + roleConnectors := resolveAllRoleConnectors(req) + if len(roleConnectors) > 0 { + injectA2OConfigWithRoutes(ctx, computer, req.Connector, roleConnectors) + } else { + injectA2OConfig(ctx, computer, req.Connector) + } } if req.ChatID != "" { @@ -179,10 +184,11 @@ func (r *Runner) Cleanup(ctx context.Context, computer infra.Computer) error { } type a2oConnectorConfig struct { - Backend string `json:"backend"` - Model string `json:"model"` - APIKey string `json:"api_key"` - Options map[string]interface{} `json:"options,omitempty"` + Backend string `json:"backend"` + Model string `json:"model"` + APIKey string `json:"api_key"` + Options map[string]interface{} `json:"options,omitempty"` + Routes map[string]*a2oConnectorConfig `json:"routes,omitempty"` } func buildSingleA2OConfig(conn connector.Connector) *a2oConnectorConfig { @@ -224,6 +230,80 @@ func buildSingleA2OConfig(conn connector.Connector) *a2oConnectorConfig { 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. diff --git a/agent/sandbox/v2/types/config.go b/agent/sandbox/v2/types/config.go index 6a790004..e1345d10 100644 --- a/agent/sandbox/v2/types/config.go +++ b/agent/sandbox/v2/types/config.go @@ -83,9 +83,39 @@ type ComputerConfig struct { // RunnerConfig identifies which Runner to use and how. type RunnerConfig struct { - Name string `json:"name" yaml:"name"` - Mode string `json:"mode,omitempty" yaml:"mode,omitempty"` - Options map[string]any `json:"options,omitempty" yaml:"options,omitempty"` + Name string `json:"name" yaml:"name"` + Mode string `json:"mode,omitempty" yaml:"mode,omitempty"` + Options map[string]any `json:"options,omitempty" yaml:"options,omitempty"` + Connectors map[string]*RoleConnector `json:"connectors,omitempty" yaml:"connectors,omitempty"` +} + +// RoleConnector maps an abstract model role to a connector. +// Accepts both a plain string and a structured object in JSON/YAML: +// +// "thinking" → {Connector: "thinking", Override: "force"} +// {"connector": "thinking", "override": "user"} → as-is +type RoleConnector struct { + Connector string `json:"connector"` + Override string `json:"override,omitempty"` // "force" (default) or "user" +} + +func (r *RoleConnector) UnmarshalJSON(data []byte) error { + var s string + if err := json.Unmarshal(data, &s); err == nil { + r.Connector = s + r.Override = "force" + return nil + } + type alias RoleConnector + var a alias + if err := json.Unmarshal(data, &a); err != nil { + return fmt.Errorf("RoleConnector: expected a string or {connector, override} object") + } + *r = RoleConnector(a) + if r.Override == "" { + r.Override = "force" + } + return nil } // PrepareStep is a single action executed during Runner.Prepare. diff --git a/agent/sandbox/v2/types/config_test.go b/agent/sandbox/v2/types/config_test.go new file mode 100644 index 00000000..6864f75c --- /dev/null +++ b/agent/sandbox/v2/types/config_test.go @@ -0,0 +1,94 @@ +package types + +import ( + "encoding/json" + "testing" +) + +func TestRoleConnector_UnmarshalJSON_String(t *testing.T) { + var rc RoleConnector + if err := json.Unmarshal([]byte(`"thinking"`), &rc); err != nil { + t.Fatal(err) + } + if rc.Connector != "thinking" { + t.Errorf("expected connector=thinking, got %s", rc.Connector) + } + if rc.Override != "force" { + t.Errorf("expected override=force, got %s", rc.Override) + } +} + +func TestRoleConnector_UnmarshalJSON_Object(t *testing.T) { + var rc RoleConnector + if err := json.Unmarshal([]byte(`{"connector":"thinking","override":"user"}`), &rc); err != nil { + t.Fatal(err) + } + if rc.Connector != "thinking" { + t.Errorf("expected connector=thinking, got %s", rc.Connector) + } + if rc.Override != "user" { + t.Errorf("expected override=user, got %s", rc.Override) + } +} + +func TestRoleConnector_UnmarshalJSON_ObjectDefaultForce(t *testing.T) { + var rc RoleConnector + if err := json.Unmarshal([]byte(`{"connector":"thinking"}`), &rc); err != nil { + t.Fatal(err) + } + if rc.Connector != "thinking" { + t.Errorf("expected connector=thinking, got %s", rc.Connector) + } + if rc.Override != "force" { + t.Errorf("expected default override=force, got %s", rc.Override) + } +} + +func TestRoleConnector_UnmarshalJSON_InvalidJSON(t *testing.T) { + var rc RoleConnector + if err := json.Unmarshal([]byte(`123`), &rc); err == nil { + t.Fatal("expected error for invalid JSON type") + } +} + +func TestRunnerConfig_UnmarshalJSON_WithConnectors(t *testing.T) { + raw := `{ + "name": "claude", + "mode": "interactive", + "connectors": { + "heavy": "thinking", + "light": {"connector": "fast", "override": "user"}, + "vision": "multimodal" + } + }` + + var rc RunnerConfig + if err := json.Unmarshal([]byte(raw), &rc); err != nil { + t.Fatal(err) + } + + if rc.Name != "claude" { + t.Errorf("expected name=claude, got %s", rc.Name) + } + if rc.Mode != "interactive" { + t.Errorf("expected mode=interactive, got %s", rc.Mode) + } + if len(rc.Connectors) != 3 { + t.Fatalf("expected 3 connectors, got %d", len(rc.Connectors)) + } + + heavy := rc.Connectors["heavy"] + if heavy.Connector != "thinking" || heavy.Override != "force" { + t.Errorf("heavy: expected {thinking, force}, got {%s, %s}", heavy.Connector, heavy.Override) + } + + light := rc.Connectors["light"] + if light.Connector != "fast" || light.Override != "user" { + t.Errorf("light: expected {fast, user}, got {%s, %s}", light.Connector, light.Override) + } + + vision := rc.Connectors["vision"] + if vision.Connector != "multimodal" || vision.Override != "force" { + t.Errorf("vision: expected {multimodal, force}, got {%s, %s}", vision.Connector, vision.Override) + } +} diff --git a/agent/sandbox/v2/types/runner.go b/agent/sandbox/v2/types/runner.go index fb7c1796..3929ab3c 100644 --- a/agent/sandbox/v2/types/runner.go +++ b/agent/sandbox/v2/types/runner.go @@ -57,4 +57,5 @@ type StreamRequest struct { ChatID string Token *SandboxToken // current user's sandbox token for MCP callbacks Logger *agentContext.RequestLogger // request-scoped logger propagated from agent context + UserExplicit bool // true when the user explicitly selected the primary connector } diff --git a/tai/workspace/workspace.go b/tai/workspace/workspace.go index 08c518b1..344e5e8c 100644 --- a/tai/workspace/workspace.go +++ b/tai/workspace/workspace.go @@ -34,6 +34,9 @@ type FS interface { // GetRoot returns the absolute path of this workspace's root directory on the host filesystem. GetRoot() (string, error) + + // GetID returns the ID of this workspace. + GetID() (string, error) } // New creates an FS backed by the given Volume for the specified session. @@ -126,6 +129,10 @@ func (w *workspaceFS) GetRoot() (string, error) { return w.vol.Abs(context.Background(), w.session, ".") } +func (w *workspaceFS) GetID() (string, error) { + return w.session, nil +} + func (w *workspaceFS) Close() error { return nil } // --- fs.FileInfo adapter --- diff --git a/workspace/manager.go b/workspace/manager.go index 3864dbbd..abfb29a3 100644 --- a/workspace/manager.go +++ b/workspace/manager.go @@ -216,9 +216,10 @@ func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEnt result := make([]DirEntry, len(entries)) for i, e := range entries { result[i] = DirEntry{ - Name: e.Path, - IsDir: e.IsDir, - Size: e.Size, + Name: e.Path, + IsDir: e.IsDir, + Size: e.Size, + ModTime: e.Mtime, } } return result, nil @@ -314,7 +315,8 @@ func listNodes() []taitypes.NodeMeta { // DirEntry represents a file or directory entry in a workspace listing. type DirEntry struct { - Name string `json:"name"` - IsDir bool `json:"is_dir"` - Size int64 `json:"size"` + Name string `json:"name"` + IsDir bool `json:"is_dir"` + Size int64 `json:"size"` + ModTime time.Time `json:"mod_time,omitempty"` }