opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(agent): add heavy role support in system configuration and connector management

Browse source 
- Introduced a new "Heavy" role in the system configuration, allowing for complex reasoning tasks.
- Updated the initialization and environment resolution functions to handle the new Heavy role appropriately.
- Enhanced tests to validate the integration of the Heavy role, ensuring proper connector resolution and model handling.
- Adjusted OpenAPI settings to reflect the addition of the Heavy role, improving overall system capabilities.
This commit is contained in:
Max 2026-05-02 13:59:14 +08:00
parent 8c34aa12ef
commit c5da1c1ba1
11 changed files with 184 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
agent/assistant/load_system.go
View file

@ -39,6 +39,7 @@ type SystemConfig struct {
Light string // Default connector for the "light" role
Vision string // Default connector for the "vision" role
Audio string // Default connector for the "audio" role
Heavy string // Default connector for the "heavy" role (complex reasoning)
// Per-agent overrides (consumed by resolveSystemConnector → ast.Connector)
Keyword string // Connector for __yao.keyword agent

3
agent/load.go
View file

@ -231,6 +231,7 @@ func initAssistant() error {
Light: agentDSL.System.Light,
Vision: agentDSL.System.Vision,
Audio: agentDSL.System.Audio,
Heavy: agentDSL.System.Heavy,
Keyword: agentDSL.System.Keyword,
QueryDSL: agentDSL.System.QueryDSL,
Title: agentDSL.System.Title,
@ -494,6 +495,7 @@ func buildSystemRoles(sys *types.System) map[string]string {
add("light", sys.Light)
add("vision", sys.Vision)
add("audio", sys.Audio)
add("heavy", sys.Heavy)
return roles
}
@ -506,6 +508,7 @@ func resolveEnvStrings(setting *types.DSL) {
setting.System.Light = helper.EnvString(setting.System.Light)
setting.System.Vision = helper.EnvString(setting.System.Vision)
setting.System.Audio = helper.EnvString(setting.System.Audio)
setting.System.Heavy = helper.EnvString(setting.System.Heavy)
setting.System.Keyword = helper.EnvString(setting.System.Keyword)
setting.System.QueryDSL = helper.EnvString(setting.System.QueryDSL)
setting.System.Title = helper.EnvString(setting.System.Title)

9
agent/sandbox/claude/command.go
View file

@ -372,9 +372,12 @@ func buildEnvironment(opts *Options, systemPrompt string) map[string]string {
}
}
// Note: System prompt and max_turns are passed via CLI flags in BuildCommand
// CLAUDE_SYSTEM_PROMPT environment variable is NOT supported by Claude CLI
// --append-system-prompt or --system-prompt flags must be used instead
// Prevent Claude CLI from using an excessive max_tokens that the backend
// API will reject. In OpenAI-proxy mode the hardcoded model is
// claude-sonnet-4-6 whose limit is 16384.
if opts.ConnectorType != "anthropic" {
env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] = "16384"
}
return env
}

22
agent/sandbox/v2/claude/command.go
View file

@ -19,6 +19,7 @@ import (
)
const defaultA2OPort = 3099
const defaultA2OMaxOutputTokens = 16384
var yaoSessionNS = uuid.MustParse("f47ac10b-58cc-4372-a567-0e02b2c3d479")
@ -178,7 +179,7 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
if len(roleConnectors) > 0 {
primaryHost := host
for role, rm := range claudeRoleEnvMap {
if role == "primary" {
if role == "default" {
continue
}
rc := resolveRoleConnector(role, roleConnectors, req.UserExplicit, getConn)
@ -206,7 +207,7 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
if len(roleConnectors) > 0 {
for role, rm := range claudeRoleEnvMap {
if role == "primary" {
if role == "default" {
continue
}
rc := resolveRoleConnector(role, roleConnectors, req.UserExplicit, getConn)
@ -228,6 +229,9 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
}
}
}
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)
@ -424,17 +428,17 @@ func buildLastUserMessageJSONL(messages []agentContext.Message) string {
// 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).
// Only roles with matching Claude CLI env vars are listed here.
// ANTHROPIC_DEFAULT_SONNET_MODEL and CLAUDE_CODE_SUBAGENT_MODEL are set to
// the primary model in buildEnv (Claude CLI doesn't have vision/subagent as
// independent role concepts).
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"},
"default": {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"},
}
func connectorHost(c connector.Connector) string {

38
agent/sandbox/v2/claude/command_test.go
View file

@ -601,10 +601,10 @@ func TestBuildEnv_OpenAI_SingleConnector(t *testing.T) {
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")
heavyConn := newOpenAIConnector("heavy-conn", "https://api.heavy.com", "heavy-model", "sk-h")
cleanup := registerTestConnectors(t, map[string]connector.Connector{
"vision-conn": vision,
"heavy-conn": heavyConn,
})
defer cleanup()
@ -612,7 +612,7 @@ func TestBuildEnv_OpenAI_MultiConnector(t *testing.T) {
Config: &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"vision": {Connector: "vision-conn", Override: "force"},
"heavy": {Connector: "heavy-conn", Override: "force"},
},
},
},
@ -622,8 +622,8 @@ func TestBuildEnv_OpenAI_MultiConnector(t *testing.T) {
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-opus-4-6", env["ANTHROPIC_DEFAULT_OPUS_MODEL"],
"heavy 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")
}
@ -692,10 +692,10 @@ func TestBuildSingleA2OConfig_Nil(t *testing.T) {
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")
heavyConn := newOpenAIConnector("heavy", "https://api.heavy.com", "heavy-model", "sk-h")
roleConnectors := map[string]connector.Connector{
"claude-vision-4-5": vision,
"claude-opus-4-6": heavyConn,
}
primaryCfg := buildSingleA2OConfig(primary)
@ -720,10 +720,10 @@ func TestInjectA2OConfigWithRoutes_BuildsCorrectJSON(t *testing.T) {
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{})
heavyRoute, ok := routesMap["claude-opus-4-6"].(map[string]interface{})
require.True(t, ok)
assert.Equal(t, "vis-model", visionRoute["model"])
assert.Contains(t, visionRoute["backend"], "api.vision.com")
assert.Equal(t, "heavy-model", heavyRoute["model"])
assert.Contains(t, heavyRoute["backend"], "api.heavy.com")
}
func TestResolveAllRoleConnectors_Empty(t *testing.T) {
@ -736,15 +736,15 @@ func TestResolveAllRoleConnectors_Empty(t *testing.T) {
}
func TestResolveAllRoleConnectors_WithRoles(t *testing.T) {
vision := newOpenAIConnector("vis", "https://vis.com", "vis-m", "sk")
cleanup := registerTestConnectors(t, map[string]connector.Connector{"vis": vision})
heavyConn := newOpenAIConnector("hvy", "https://heavy.com", "heavy-m", "sk")
cleanup := registerTestConnectors(t, map[string]connector.Connector{"hvy": heavyConn})
defer cleanup()
req := &types.StreamRequest{
Config: &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"vision": {Connector: "vis", Override: "force"},
"heavy": {Connector: "hvy", Override: "force"},
},
},
},
@ -752,15 +752,15 @@ func TestResolveAllRoleConnectors_WithRoles(t *testing.T) {
}
result := resolveAllRoleConnectors(req)
assert.Len(t, result, 1)
assert.Equal(t, vision, result["claude-vision-4-5"])
assert.Equal(t, heavyConn, result["claude-opus-4-6"])
}
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")
heavyConn := newOpenAIConnector("heavy-oai", "https://api.openai.com", "gpt-4o", "sk-oai")
cleanup := registerTestConnectors(t, map[string]connector.Connector{
"vision-oai": visionConn,
"heavy-oai": heavyConn,
})
defer cleanup()
@ -768,7 +768,7 @@ func TestBuildEnv_Anthropic_MultiConnector_Incompatible(t *testing.T) {
Config: &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"vision": {Connector: "vision-oai", Override: "force"},
"heavy": {Connector: "heavy-oai", Override: "force"},
},
},
},
@ -778,6 +778,6 @@ func TestBuildEnv_Anthropic_MultiConnector_Incompatible(t *testing.T) {
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)")
assert.Equal(t, "claude-sonnet-4-20250514", env["ANTHROPIC_DEFAULT_OPUS_MODEL"],
"incompatible connector: heavy should keep primary model (different host, no anthropic protocol)")
}

6
agent/sandbox/v2/claude/runner.go
View file

@ -233,6 +233,10 @@ func buildSingleA2OConfig(conn connector.Connector) *a2oConnectorConfig {
cfg.Options = extra
}
if cfg.MaxOutputTokens == 0 {
cfg.MaxOutputTokens = defaultA2OMaxOutputTokens
}
if cfg.Backend == "" {
return nil
}
@ -249,7 +253,7 @@ func resolveAllRoleConnectors(req *types.StreamRequest) map[string]connector.Con
result := make(map[string]connector.Connector)
for role, rm := range claudeRoleEnvMap {
if role == "primary" {
if role == "default" {
continue
}
rc := resolveRoleConnector(role, roleConns, req.UserExplicit, func(id string) connector.Connector {

24
agent/sandbox/v2/opencode/command.go
View file

@ -122,14 +122,15 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
// like browsers should be nohup'd; this prevents accidental 2-min hangs.
env["OPENCODE_EXPERIMENTAL_BASH_DEFAULT_TIMEOUT_MS"] = "30000"
if req.Connector != nil {
setting := req.Connector.Setting()
primaryConn := resolvePrimaryConnector(req.Connector, req.Config)
if primaryConn != nil {
setting := primaryConn.Setting()
key, _ := setting["key"].(string)
if key != "" {
env["YAO_PROVIDER_KEY"] = key
}
if req.Connector.Is(connector.ANTHROPIC) {
if primaryConn.Is(connector.ANTHROPIC) {
apiKey, _ := setting["key"].(string)
if apiKey != "" {
env["ANTHROPIC_API_KEY"] = apiKey
@ -175,8 +176,9 @@ func buildArgs(req *types.StreamRequest, r *Runner, isContinuation bool, chatID
args = append(args, "--continue", "--session", sessionID)
}
if req.Connector != nil {
if mid := connectorModelID(req.Connector); mid != "" {
primaryConn := resolvePrimaryConnector(req.Connector, req.Config)
if primaryConn != nil {
if mid := connectorModelID(primaryConn); mid != "" {
args = append(args, "--model", mid)
}
}
@ -350,10 +352,16 @@ func shellQuotePowerShell(program string, args ...string) string {
// connectorModelID returns the "provider/model" string matching the
// provider ID used in opencode.json (see buildProviderConfig).
// Uses LLMConnector interface first (consistent with buildProviderConfig).
func connectorModelID(c connector.Connector) string {
setting := c.Setting()
modelName, _ := setting["model"].(string)
host, _ := setting["host"].(string)
host := connectorHost(c)
var modelName string
if lc, ok := c.(goullm.LLMConnector); ok {
modelName = lc.GetModel()
}
if modelName == "" {
modelName, _ = c.Setting()["model"].(string)
}
if c.Is(connector.ANTHROPIC) {
return "anthropic/" + modelName

72
agent/sandbox/v2/opencode/config.go
View file

@ -12,9 +12,11 @@ import (
type roleSpec struct {
EnvKeyPrefix string
TopLevel string
Modalities map[string][]string
}
// openCodeRoleMap lists roles that map to native OpenCode config concepts.
// "light" → top-level "small_model"; "vision" → env vars only (read.ts hack).
// "heavy" is handled via resolvePrimaryConnector (becomes the main model).
var openCodeRoleMap = map[string]roleSpec{
"light": {
EnvKeyPrefix: "YAO_LIGHT",
@ -22,19 +24,28 @@ var openCodeRoleMap = map[string]roleSpec{
},
"vision": {
EnvKeyPrefix: "YAO_VISION",
Modalities: map[string][]string{
"input": {"text", "image"},
"output": {"text"},
},
},
"heavy": {
EnvKeyPrefix: "YAO_HEAVY",
},
"subagent": {
EnvKeyPrefix: "YAO_SUBAGENT",
},
}
// resolvePrimaryConnector returns the heavy connector if configured,
// otherwise falls back to the caller-supplied primary (typically the
// assistant's default connector). This aligns with OpenCode's semantics
// where the top-level "model" handles complex coding tasks.
func resolvePrimaryConnector(primary connector.Connector, cfg *types.SandboxConfig) connector.Connector {
if cfg == nil || cfg.Runner.Connectors == nil {
return primary
}
rc, ok := cfg.Runner.Connectors["heavy"]
if !ok || rc == nil || rc.Connector == "" {
return primary
}
c, exists := connector.Connectors[rc.Connector]
if !exists || c == nil {
return primary
}
return c
}
// buildOpenCodeConfig generates the opencode.json project configuration.
// All provider configuration is direct (no a2o proxy).
func buildOpenCodeConfig(req *types.PrepareRequest, mcpServers []types.MCPServer) []byte {
@ -47,14 +58,15 @@ func buildOpenCodeConfig(req *types.PrepareRequest, mcpServers []types.MCPServer
"permission": map[string]any{"*": "allow"},
}
if req.Connector != nil {
providerID, providerCfg, modelStr := buildProviderConfig(req.Connector)
primaryConn := resolvePrimaryConnector(req.Connector, req.Config)
if primaryConn != nil {
providerID, providerCfg, modelStr := buildProviderConfig(primaryConn)
cfg["provider"] = map[string]any{providerID: providerCfg}
cfg["model"] = modelStr
cfg["enabled_providers"] = []string{providerID}
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
if len(mcpServers) > 0 {
cfg["mcp"] = buildMCPConfig(mcpServers)
@ -126,7 +138,8 @@ func buildProviderConfig(conn connector.Connector) (providerID string, cfg map[s
modelCfg["interleaved"] = map[string]any{"field": "reasoning_content"}
// Forward connector-level request body params (thinking, reasoning, etc.)
// to OpenCode model options so they reach the upstream API.
// to OpenCode model options, using the same FilterRequestBodyParams
// mechanism as buildRequestBody in yao/agent/llm.
connParams := connector.FilterRequestBodyParams(setting, conn)
if len(connParams) > 0 {
modelCfg["options"] = connParams
@ -179,9 +192,9 @@ func normalizeBaseURL(host string) string {
// injectRoleProviders iterates openCodeRoleMap and injects provider blocks
// for every role that has a configured connector. For the "light" role it
// also sets the top-level "small_model" field. This replaces the old
// buildSmallModel function and adds support for vision/heavy/subagent roles.
func injectRoleProviders(cfg map[string]any, req *types.PrepareRequest) {
// also sets the top-level "small_model" field. primaryConn is the resolved
// primary connector (may be heavy or default) used for sameProvider checks.
func injectRoleProviders(cfg map[string]any, req *types.PrepareRequest, primaryConn connector.Connector) {
if req.Config == nil || req.Config.Runner.Connectors == nil {
return
}
@ -200,9 +213,9 @@ func injectRoleProviders(cfg map[string]any, req *types.PrepareRequest) {
primaryHost := ""
primaryType := ""
if req.Connector != nil {
primaryHost = connectorHost(req.Connector)
if req.Connector.Is(connector.ANTHROPIC) {
if primaryConn != nil {
primaryHost = connectorHost(primaryConn)
if primaryConn.Is(connector.ANTHROPIC) {
primaryType = "anthropic"
} else {
primaryType = "openai"
@ -247,10 +260,10 @@ func injectRoleProviders(cfg map[string]any, req *types.PrepareRequest) {
if sameProvider {
providerID = resolveExistingProviderID(providers, primaryType)
modelRef = providerID + "/" + modelName
mergeModelIntoProvider(providers, providerID, modelName, spec.Modalities)
mergeModelIntoProvider(providers, providerID, modelName)
} else {
providerID = role
providerCfg := buildRoleProviderConfig(c, spec.EnvKeyPrefix, spec.Modalities)
providerCfg := buildRoleProviderConfig(c, spec.EnvKeyPrefix)
providers[providerID] = providerCfg
modelRef = providerID + "/" + modelName
}
@ -284,7 +297,7 @@ func resolveExistingProviderID(providers map[string]any, pType string) string {
}
// mergeModelIntoProvider adds a model entry to an existing provider block.
func mergeModelIntoProvider(providers map[string]any, providerID, modelName string, modalities map[string][]string) {
func mergeModelIntoProvider(providers map[string]any, providerID, modelName string) {
block, ok := providers[providerID].(map[string]any)
if !ok {
return
@ -294,17 +307,13 @@ func mergeModelIntoProvider(providers map[string]any, providerID, modelName stri
models = map[string]any{}
block["models"] = models
}
modelCfg := map[string]any{"name": modelName}
if len(modalities) > 0 {
modelCfg["modalities"] = modalities
}
models[modelName] = modelCfg
models[modelName] = map[string]any{"name": modelName}
}
// buildRoleProviderConfig creates a provider configuration block for a
// non-primary role connector. Uses the role's env key prefix for API key
// and base URL references.
func buildRoleProviderConfig(conn connector.Connector, envKeyPrefix string, modalities map[string][]string) map[string]any {
func buildRoleProviderConfig(conn connector.Connector, envKeyPrefix string) map[string]any {
setting := conn.Setting()
modelName, _ := setting["model"].(string)
host, _ := setting["host"].(string)
@ -314,9 +323,6 @@ func buildRoleProviderConfig(conn connector.Connector, envKeyPrefix string, moda
}
modelCfg := map[string]any{"name": modelName}
if len(modalities) > 0 {
modelCfg["modalities"] = modalities
}
if lc, ok := conn.(goullm.LLMConnector); ok {
if caps := lc.GetCapabilities(); caps != nil {

112
agent/sandbox/v2/opencode/roles_test.go
View file

@ -89,7 +89,7 @@ func TestInjectRoleProviders_VisionCustomProvider(t *testing.T) {
},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
providers := cfg["provider"].(map[string]any)
visionBlock, ok := providers["vision"]
@ -99,14 +99,8 @@ func TestInjectRoleProviders_VisionCustomProvider(t *testing.T) {
vBlock := visionBlock.(map[string]any)
models := vBlock["models"].(map[string]any)
modelCfg := models["gpt-4o-mini"].(map[string]any)
mods, ok := modelCfg["modalities"].(map[string][]string)
if !ok {
t.Fatal("vision model should have modalities declared")
}
if len(mods["input"]) != 2 || mods["input"][0] != "text" || mods["input"][1] != "image" {
t.Errorf("modalities.input = %v, want [text, image]", mods["input"])
if _, ok := models["gpt-4o-mini"]; !ok {
t.Fatal("vision provider should contain gpt-4o-mini model")
}
enabled := cfg["enabled_providers"].([]string)
@ -144,21 +138,12 @@ func TestInjectRoleProviders_VisionNativeOpenAI(t *testing.T) {
},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
providers := cfg["provider"].(map[string]any)
visionBlock, ok := providers["vision"]
if !ok {
if _, ok := providers["vision"]; !ok {
t.Fatal("should have separate 'vision' provider (different host from primary)")
}
vBlock := visionBlock.(map[string]any)
models := vBlock["models"].(map[string]any)
modelCfg := models["gpt-4o-mini"].(map[string]any)
if _, ok := modelCfg["modalities"]; !ok {
t.Error("native OpenAI vision model should still declare modalities")
}
}
func TestInjectRoleProviders_LightWithDifferentHost(t *testing.T) {
@ -184,7 +169,7 @@ func TestInjectRoleProviders_LightWithDifferentHost(t *testing.T) {
},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
providers := cfg["provider"].(map[string]any)
if _, ok := providers["light"]; !ok {
@ -236,7 +221,7 @@ func TestInjectRoleProviders_LightSameHostAsPrimary(t *testing.T) {
},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
providers := cfg["provider"].(map[string]any)
if _, ok := providers["light"]; ok {
@ -262,11 +247,13 @@ func TestInjectRoleProviders_NoConnectors(t *testing.T) {
"enabled_providers": []string{"openai"},
}
primaryConn := newFakeOpenAI("primary", "", "gpt-4o", "sk-oai")
req := &types.PrepareRequest{
Config: &types.SandboxConfig{},
Connector: primaryConn,
Config: &types.SandboxConfig{},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
enabled := cfg["enabled_providers"].([]string)
if len(enabled) != 1 || enabled[0] != "openai" {
@ -297,7 +284,7 @@ func TestInjectRoleProviders_AnthropicVision(t *testing.T) {
},
}
injectRoleProviders(cfg, req)
injectRoleProviders(cfg, req, primaryConn)
providers := cfg["provider"].(map[string]any)
visionBlock, ok := providers["vision"]
@ -381,7 +368,7 @@ func TestInjectRoleEnvVars_NoConnectors(t *testing.T) {
env := map[string]string{}
injectRoleEnvVars(env, req)
for _, prefix := range []string{"YAO_VISION", "YAO_LIGHT", "YAO_HEAVY", "YAO_SUBAGENT"} {
for _, prefix := range []string{"YAO_VISION", "YAO_LIGHT"} {
for _, suffix := range []string{"_KEY", "_BASE_URL", "_MODEL"} {
if v, ok := env[prefix+suffix]; ok {
t.Errorf("unexpected env %s=%s with no connectors", prefix+suffix, v)
@ -393,12 +380,10 @@ func TestInjectRoleEnvVars_NoConnectors(t *testing.T) {
func TestInjectRoleEnvVars_MultipleRoles(t *testing.T) {
visionConn := newFakeOpenAI("vis", "https://api.vision.com", "vis-model", "sk-vis")
lightConn := newFakeOpenAI("light-c", "https://api.light.com", "light-model", "sk-light")
heavyConn := newFakeOpenAI("heavy-c", "https://api.heavy.com", "heavy-model", "sk-heavy")
cleanup := registerFakeConnectors(t, map[string]connector.Connector{
"vis-c": visionConn,
"light-c": lightConn,
"heavy-c": heavyConn,
})
defer cleanup()
@ -408,7 +393,6 @@ func TestInjectRoleEnvVars_MultipleRoles(t *testing.T) {
Connectors: map[string]*types.RoleConnector{
"vision": {Connector: "vis-c", Override: "force"},
"light": {Connector: "light-c", Override: "force"},
"heavy": {Connector: "heavy-c", Override: "force"},
},
},
},
@ -423,9 +407,6 @@ func TestInjectRoleEnvVars_MultipleRoles(t *testing.T) {
if env["YAO_LIGHT_KEY"] != "sk-light" {
t.Errorf("YAO_LIGHT_KEY = %q", env["YAO_LIGHT_KEY"])
}
if env["YAO_HEAVY_KEY"] != "sk-heavy" {
t.Errorf("YAO_HEAVY_KEY = %q", env["YAO_HEAVY_KEY"])
}
}
// ---------------------------------------------------------------------------
@ -491,3 +472,70 @@ func TestBuildOpenCodeConfig_WithVisionAndLight(t *testing.T) {
}
}
}
// ---------------------------------------------------------------------------
// resolvePrimaryConnector tests
// ---------------------------------------------------------------------------
func TestResolvePrimaryConnector_HeavyConfigured(t *testing.T) {
defaultConn := newFakeOpenAI("default", "https://api.deepseek.com", "deepseek-chat", "sk-ds")
heavyConn := newFakeOpenAI("heavy", "https://api.openai.com", "o3-pro", "sk-oai")
cleanup := registerFakeConnectors(t, map[string]connector.Connector{"heavy-conn": heavyConn})
defer cleanup()
cfg := &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"heavy": {Connector: "heavy-conn", Override: "force"},
},
},
}
result := resolvePrimaryConnector(defaultConn, cfg)
if result != heavyConn {
t.Error("should return heavy connector when configured")
}
}
func TestResolvePrimaryConnector_NoHeavy(t *testing.T) {
defaultConn := newFakeOpenAI("default", "https://api.deepseek.com", "deepseek-chat", "sk-ds")
cfg := &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"light": {Connector: "light-conn", Override: "force"},
},
},
}
result := resolvePrimaryConnector(defaultConn, cfg)
if result != defaultConn {
t.Error("should fallback to default when heavy not configured")
}
}
func TestResolvePrimaryConnector_HeavyNotRegistered(t *testing.T) {
defaultConn := newFakeOpenAI("default", "https://api.deepseek.com", "deepseek-chat", "sk-ds")
cfg := &types.SandboxConfig{
Runner: types.RunnerConfig{
Connectors: map[string]*types.RoleConnector{
"heavy": {Connector: "nonexistent-conn", Override: "force"},
},
},
}
result := resolvePrimaryConnector(defaultConn, cfg)
if result != defaultConn {
t.Error("should fallback to default when heavy connector not registered")
}
}
func TestResolvePrimaryConnector_NilConfig(t *testing.T) {
defaultConn := newFakeOpenAI("default", "https://api.deepseek.com", "deepseek-chat", "sk-ds")
result := resolvePrimaryConnector(defaultConn, nil)
if result != defaultConn {
t.Error("should return default when config is nil")
}
}

1
agent/types/types.go
View file

@ -98,6 +98,7 @@ type System struct {
Light string `json:"light,omitempty" yaml:"light,omitempty"` // Default connector for the "light" role (titles, keywords, summaries)
Vision string `json:"vision,omitempty" yaml:"vision,omitempty"` // Default connector for the "vision" role
Audio string `json:"audio,omitempty" yaml:"audio,omitempty"` // Default connector for the "audio" role
Heavy string `json:"heavy,omitempty" yaml:"heavy,omitempty"` // Default connector for the "heavy" role (complex reasoning)
// Per-agent overrides (optional, highest priority — bypasses role resolution)
Keyword string `json:"keyword,omitempty" yaml:"keyword,omitempty"` // Connector for __yao.keyword agent

2
openapi/tests/setting/llm_test.go
View file

@ -278,7 +278,7 @@ func TestLLMProviderDelete(t *testing.T) {
rolesPayload := map[string]interface{}{
"default": map[string]interface{}{
"provider": "anthropic",
"model": "claude-sonnet-4-20250514",
"model": "claude-sonnet-4-6",
},
}
rolesResp := llmPut(t, llmURL(serverURL, "/roles"), token, rolesPayload)