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.
This commit is contained in:
parent
cc5e888f5a
commit
2084e20479
11 changed files with 708 additions and 17 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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{
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -90,8 +90,13 @@ 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) {
|
||||
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 {
|
||||
|
|
@ -183,6 +188,7 @@ type a2oConnectorConfig struct {
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -86,6 +86,36 @@ 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"`
|
||||
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.
|
||||
|
|
|
|||
94
agent/sandbox/v2/types/config_test.go
Normal file
94
agent/sandbox/v2/types/config_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 ---
|
||||
|
|
|
|||
|
|
@ -219,6 +219,7 @@ func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEnt
|
|||
Name: e.Path,
|
||||
IsDir: e.IsDir,
|
||||
Size: e.Size,
|
||||
ModTime: e.Mtime,
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
|
|
@ -317,4 +318,5 @@ type DirEntry struct {
|
|||
Name string `json:"name"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
Size int64 `json:"size"`
|
||||
ModTime time.Time `json:"mod_time,omitempty"`
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue