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

package claude

import (
	"context"
	"encoding/json"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/yaoapp/gou/connector"
	goullm "github.com/yaoapp/gou/llm"
	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"
	"github.com/yaoapp/yao/tai/workspace"
)

func testPlatform() platform {
	return &darwinPlatform{posixBase: posixBase{
		os: "darwin", workDir: "/workspace", shell: "bash", tempDir: "/tmp",
	}}
}

// --- buildEnv ---

func TestBuildEnv_HomeEnv(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	assert.Equal(t, "/workspace", env["HOME"])
}

func TestBuildEnv_ConfigDirIsolation(t *testing.T) {
	req := &types.StreamRequest{
		Config:      &types.SandboxConfig{ID: "my-assistant"},
		AssistantID: "my-assistant",
	}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	expected := "/workspace/.yao/assistants/my-assistant"
	assert.Equal(t, expected, env["CLAUDE_CONFIG_DIR"])
}

func TestBuildEnv_NoAssistantID(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	_, hasConfigDir := env["CLAUDE_CONFIG_DIR"]
	assert.False(t, hasConfigDir, "should not set CLAUDE_CONFIG_DIR without assistantID")
}

func TestBuildEnv_Token(t *testing.T) {
	req := &types.StreamRequest{
		Config: &types.SandboxConfig{},
		Token: &types.SandboxToken{
			Token:        "tok123",
			RefreshToken: "ref456",
		},
	}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	assert.Equal(t, "tok123", env["YAO_TOKEN"])
	assert.Equal(t, "ref456", env["YAO_REFRESH_TOKEN"])
}

func TestBuildEnv_Secrets(t *testing.T) {
	req := &types.StreamRequest{
		Config: &types.SandboxConfig{
			Secrets: map[string]string{
				"MY_SECRET": "secret_val",
			},
		},
	}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	assert.Equal(t, "secret_val", env["MY_SECRET"])
}

// --- buildArgs ---

func TestBuildArgs_Default(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, false, "", "")
	assert.Contains(t, args, "--input-format")
	assert.Contains(t, args, "stream-json")
	assert.Contains(t, args, "--output-format")
	assert.Contains(t, args, "--verbose")
	assert.Contains(t, args, "--include-partial-messages")
}

func TestBuildArgs_Continuation(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, true, "", "")
	assert.Contains(t, args, "--continue")
}

func TestBuildArgs_PermissionMode(t *testing.T) {
	req := &types.StreamRequest{
		Config: &types.SandboxConfig{
			Runner: types.RunnerConfig{
				Options: map[string]interface{}{
					"permission_mode": "bypassPermissions",
				},
			},
		},
	}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, false, "", "")
	assert.Contains(t, args, "--dangerously-skip-permissions")
	assert.Contains(t, args, "--permission-mode")
}

func TestBuildArgs_MCP(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
	p := testPlatform()

	args := buildArgs(req, r, p, false, "test-assistant", "")
	assert.Contains(t, args, "--mcp-config")
	assert.Contains(t, args, "--allowedTools")
	assert.Contains(t, args, "mcp__yao__*")

	mcpIdx := -1
	for i, a := range args {
		if a == "--mcp-config" {
			mcpIdx = i
			break
		}
	}
	require.Greater(t, mcpIdx, -1)
	mcpPath := args[mcpIdx+1]
	assert.Contains(t, mcpPath, "test-assistant")
}

func TestBuildArgs_WhitelistOptions(t *testing.T) {
	req := &types.StreamRequest{
		Config: &types.SandboxConfig{
			Runner: types.RunnerConfig{
				Options: map[string]interface{}{
					"max_turns": 10,
				},
			},
		},
	}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, false, "", "")
	assert.Contains(t, args, "--max-turns")
}

// --- buildLastUserMessageJSONL (was buildInput / buildFirstRequestJSONL) ---

func TestBuildLastUserMessageJSONL_SkipsSystem(t *testing.T) {
	msgs := []agentContext.Message{
		{Role: "system", Content: "system prompt"},
		{Role: "user", Content: "hello"},
	}
	result := buildLastUserMessageJSONL(msgs)
	var parsed map[string]any
	err := json.Unmarshal([]byte(strings.TrimSpace(result)), &parsed)
	require.NoError(t, err)
	assert.Equal(t, "user", parsed["type"])
	msg, _ := parsed["message"].(map[string]any)
	assert.Equal(t, "hello", msg["content"])
}

func TestBuildLastUserMessageJSONL_OnlyLastUser(t *testing.T) {
	msgs := []agentContext.Message{
		{Role: "user", Content: "q1"},
		{Role: "assistant", Content: "a1"},
		{Role: "user", Content: "q2"},
	}
	result := buildLastUserMessageJSONL(msgs)
	var parsed map[string]any
	err := json.Unmarshal([]byte(strings.TrimSpace(result)), &parsed)
	require.NoError(t, err)
	assert.Equal(t, "user", parsed["type"])
	msg, _ := parsed["message"].(map[string]any)
	assert.Equal(t, "q2", msg["content"])
	assert.NotContains(t, result, "q1")
}

func TestBuildLastUserMessageJSONL_NilContent(t *testing.T) {
	msgs := []agentContext.Message{
		{Role: "user", Content: nil},
	}
	result := buildLastUserMessageJSONL(msgs)
	var parsed map[string]any
	err := json.Unmarshal([]byte(strings.TrimSpace(result)), &parsed)
	require.NoError(t, err)
	msg, _ := parsed["message"].(map[string]any)
	assert.Equal(t, "", msg["content"])
}

// --- buildLastUserMessageJSONL ---

func TestBuildLastUserMessageJSONL_Found(t *testing.T) {
	msgs := []agentContext.Message{
		{Role: "user", Content: "first"},
		{Role: "assistant", Content: "reply"},
		{Role: "user", Content: "second"},
	}
	result := buildLastUserMessageJSONL(msgs)
	var parsed map[string]any
	err := json.Unmarshal([]byte(result), &parsed)
	require.NoError(t, err)
	msg, _ := parsed["message"].(map[string]any)
	assert.Equal(t, "second", msg["content"])
}

func TestBuildLastUserMessageJSONL_NoUser(t *testing.T) {
	msgs := []agentContext.Message{
		{Role: "assistant", Content: "only assistant"},
	}
	assert.Empty(t, buildLastUserMessageJSONL(msgs))
}

func TestBuildLastUserMessageJSONL_Empty(t *testing.T) {
	assert.Empty(t, buildLastUserMessageJSONL(nil))
}

// --- buildMCPConfig ---

func TestBuildMCPConfig_WithServers(t *testing.T) {
	servers := []types.MCPServer{
		{ServerID: "server1"},
		{ServerID: "server2"},
	}
	data := buildMCPConfig(servers)
	var cfg map[string]any
	err := json.Unmarshal(data, &cfg)
	require.NoError(t, err)
	mcpServers, ok := cfg["mcpServers"].(map[string]any)
	require.True(t, ok)
	assert.Contains(t, mcpServers, "server1")
	assert.Contains(t, mcpServers, "server2")
}

func TestBuildMCPConfig_Empty(t *testing.T) {
	data := buildMCPConfig(nil)
	var cfg map[string]any
	err := json.Unmarshal(data, &cfg)
	require.NoError(t, err)
	mcpServers, ok := cfg["mcpServers"].(map[string]any)
	require.True(t, ok)
	assert.Contains(t, mcpServers, "yao", "should default to yao server")
}

func TestBuildMCPConfig_EmptyServerID(t *testing.T) {
	servers := []types.MCPServer{{ServerID: ""}}
	data := buildMCPConfig(servers)
	var cfg map[string]any
	json.Unmarshal(data, &cfg)
	mcpServers, _ := cfg["mcpServers"].(map[string]any)
	assert.Contains(t, mcpServers, "yao", "empty serverID should fall back to default")
}

// --- buildMCPAllowedTools ---

func TestBuildMCPAllowedTools_WithServers(t *testing.T) {
	servers := []types.MCPServer{
		{ServerID: "s1"},
		{ServerID: "s2"},
	}
	result := buildMCPAllowedTools(servers)
	assert.Contains(t, result, "mcp__s1__*")
	assert.Contains(t, result, "mcp__s2__*")
	assert.Contains(t, result, ",")
}

func TestBuildMCPAllowedTools_Empty(t *testing.T) {
	assert.Equal(t, "mcp__yao__*", buildMCPAllowedTools(nil))
}

// --- buildSandboxEnvPrompt ---

func TestBuildSandboxEnvPrompt(t *testing.T) {
	p := testPlatform()
	prompt := buildSandboxEnvPrompt(p, "/workspace")
	assert.Contains(t, prompt, "/workspace")
	assert.Contains(t, prompt, "darwin")
	assert.Contains(t, prompt, "bash")
	assert.Contains(t, prompt, "Sandbox Environment")
}

func TestBuildSandboxEnvPrompt_WindowsPlatform(t *testing.T) {
	p := newWindowsPlatform(`C:\ws`, "pwsh", "")
	prompt := buildSandboxEnvPrompt(p, `C:\ws`)
	assert.Contains(t, prompt, "windows")
	assert.Contains(t, prompt, "pwsh")
	assert.Contains(t, prompt, `C:\ws`)
	assert.Contains(t, prompt, "Windows desktop")
}

// --- fakeComputer implements infra.Computer for unit tests ---

type fakeComputer struct {
	workDir string
}

func newFakeComputer(workDir string) *fakeComputer {
	return &fakeComputer{workDir: workDir}
}

func (f *fakeComputer) GetWorkDir() string      { return f.workDir }
func (f *fakeComputer) BindWorkplace(string)    {}
func (f *fakeComputer) Workplace() workspace.FS { return nil }
func (f *fakeComputer) ComputerInfo() infra.ComputerInfo {
	return infra.ComputerInfo{System: infra.SystemInfo{OS: "linux", Shell: "bash"}}
}
func (f *fakeComputer) Exec(_ context.Context, _ []string, _ ...infra.ExecOption) (*infra.ExecResult, error) {
	return &infra.ExecResult{}, nil
}
func (f *fakeComputer) Stream(_ context.Context, _ []string, _ ...infra.ExecOption) (*infra.ExecStream, error) {
	return nil, nil
}
func (f *fakeComputer) VNC(_ context.Context) (string, error)                    { return "", nil }
func (f *fakeComputer) Proxy(_ context.Context, _ int, _ string) (string, error) { return "", nil }

// --- chatIDToSessionUUID ---

func TestChatIDToSessionUUID_Deterministic(t *testing.T) {
	u1 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
	u2 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
	assert.Equal(t, u1, u2, "same inputs should produce same UUID")
}

func TestChatIDToSessionUUID_DifferentAssistant(t *testing.T) {
	u1 := chatIDToSessionUUID("asst-1", "robot_m1_e1")
	u2 := chatIDToSessionUUID("asst-2", "robot_m1_e1")
	assert.NotEqual(t, u1, u2, "different assistantID should produce different UUID")
}

func TestChatIDToSessionUUID_ValidFormat(t *testing.T) {
	u := chatIDToSessionUUID("asst-1", "robot_m1_e1")
	assert.Regexp(t, `^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`, u)
}

// --- sanitizeSessionName ---

func TestSanitizeSessionName_Normal(t *testing.T) {
	assert.Equal(t, "yao-robot_m1_e1", sanitizeSessionName("robot_m1_e1"))
}

func TestSanitizeSessionName_SpecialChars(t *testing.T) {
	assert.Equal(t, "yao-user_s__chat_", sanitizeSessionName("user's \"chat\""))
}

func TestSanitizeSessionName_Empty(t *testing.T) {
	assert.Equal(t, "yao-", sanitizeSessionName(""))
}

// --- buildArgs with session ---

func TestBuildArgs_SessionID_NewSession(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, false, "asst-1", "robot_m1_e1")
	assert.Contains(t, args, "--session-id")
	assert.Contains(t, args, "--name")
	assert.Contains(t, args, "yao-robot_m1_e1")
	assert.NotContains(t, args, "--resume")
	assert.NotContains(t, args, "--continue")

	sidIdx := -1
	for i, a := range args {
		if a == "--session-id" {
			sidIdx = i
			break
		}
	}
	require.Greater(t, sidIdx, -1)
	assert.Regexp(t, `^[0-9a-f]{8}-`, args[sidIdx+1])
}

func TestBuildArgs_SessionID_Continuation(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, true, "asst-1", "robot_m1_e1")
	assert.Contains(t, args, "--resume")
	assert.Contains(t, args, "--name")
	assert.NotContains(t, args, "--session-id")
	assert.NotContains(t, args, "--continue")

	resumeIdx := -1
	for i, a := range args {
		if a == "--resume" {
			resumeIdx = i
			break
		}
	}
	require.Greater(t, resumeIdx, -1)
	assert.Regexp(t, `^[0-9a-f]{8}-`, args[resumeIdx+1])
}

func TestBuildArgs_EmptyChatID_Continuation(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	req.Computer = newFakeComputer("/workspace")
	r := &Runner{}
	p := testPlatform()

	args := buildArgs(req, r, p, true, "", "")
	assert.Contains(t, args, "--continue")
	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 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, "kimi-k2.5", env["ANTHROPIC_MODEL"])
	assert.Equal(t, "kimi-k2.5", env["ANTHROPIC_DEFAULT_OPUS_MODEL"])
}

func TestBuildEnv_OpenAI_MultiConnector(t *testing.T) {
	primary := newOpenAIConnector("kimi", "https://api.moonshot.cn", "kimi-k2.5", "sk-test")
	heavyConn := newOpenAIConnector("heavy-conn", "https://api.heavy.com", "heavy-model", "sk-h")

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"heavy":   heavyConn,
		},
	}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	assert.Equal(t, "heavy-model", env["ANTHROPIC_DEFAULT_OPUS_MODEL"],
		"heavy role should use actual model name from connector")
	assert.Equal(t, "kimi-k2.5", env["ANTHROPIC_MODEL"],
		"primary should use actual model name from connector")
	assert.Equal(t, "kimi-k2.5", env["ANTHROPIC_CUSTOM_MODEL_OPTION"],
		"non-standard model should set custom model option")
}

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")

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"light":   lightConn,
		},
	}
	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")
	heavyConn := newOpenAIConnector("heavy", "https://api.heavy.com", "heavy-model", "sk-h")

	roleConnectors := map[string]connector.Connector{
		"heavy-model": heavyConn,
	}

	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)

	heavyRoute, ok := routesMap["heavy-model"].(map[string]interface{})
	require.True(t, ok)
	assert.Equal(t, "heavy-model", heavyRoute["model"])
	assert.Contains(t, heavyRoute["backend"], "api.heavy.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) {
	primaryConn := newOpenAIConnector("primary", "https://primary.com", "primary-m", "k")
	heavyConn := newOpenAIConnector("hvy", "https://heavy.com", "heavy-m", "sk")

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primaryConn,
		Roles: map[string]connector.Connector{
			"default": primaryConn,
			"heavy":   heavyConn,
		},
	}
	result := resolveAllRoleConnectors(req)
	assert.Len(t, result, 2)
	assert.Equal(t, primaryConn, result["primary-m"])
	assert.Equal(t, heavyConn, result["heavy-m"])
}

func TestBuildEnv_Anthropic_MultiConnector_Incompatible(t *testing.T) {
	primary := newAnthropicConnector("claude", "https://api.anthropic.com", "claude-sonnet-4-20250514", "sk-ant")
	heavyConn := newOpenAIConnector("heavy-oai", "https://api.openai.com", "gpt-4o", "sk-oai")

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"heavy":   heavyConn,
		},
	}
	req.Computer = newFakeComputer("/workspace")
	p := testPlatform()

	env := buildEnv(req, p)
	assert.Equal(t, "claude-sonnet-4-20250514", env["ANTHROPIC_DEFAULT_OPUS_MODEL"],
		"incompatible connector: heavy should keep primary model (different host, no anthropic protocol)")
}

// --- fakeLLMConnector implements both connector.Connector and goullm.LLMConnector ---

type fakeLLMConnector struct {
	fakeConnector
	model string
	caps  *goullm.Capabilities
}

func (f *fakeLLMConnector) GetAuthMode() goullm.AuthMode { return goullm.AuthBearer }
func (f *fakeLLMConnector) GetURL() string               { return "" }
func (f *fakeLLMConnector) GetKey() string               { return "" }
func (f *fakeLLMConnector) GetModel() string             { return f.model }
func (f *fakeLLMConnector) GetSupportedParams() map[string]*goullm.ParamSpec {
	return nil
}
func (f *fakeLLMConnector) GetCapabilities() *goullm.Capabilities { return f.caps }

// --- buildModelCapabilityPrompt tests ---

func TestBuildModelCapabilityPrompt_NilConnector(t *testing.T) {
	req := &types.StreamRequest{Config: &types.SandboxConfig{}}
	result := buildModelCapabilityPrompt(req)
	assert.Empty(t, result)
}

func TestBuildModelCapabilityPrompt_NoRoles_NoSpecialCaps(t *testing.T) {
	primary := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "test", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-flash"}},
		model:         "deepseek-v4-flash",
		caps:          &goullm.Capabilities{ToolCalls: true, Streaming: true},
	}
	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
	}
	result := buildModelCapabilityPrompt(req)
	assert.Empty(t, result, "no roles and no special caps → empty")
}

func TestBuildModelCapabilityPrompt_WithHeavyAndLight(t *testing.T) {
	primary := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "ds-flash", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-flash"}},
		model:         "deepseek-v4-flash",
		caps:          &goullm.Capabilities{ToolCalls: true},
	}
	heavy := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "ds-pro", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-pro"}},
		model:         "deepseek-v4-pro",
		caps:          &goullm.Capabilities{Reasoning: true, ToolCalls: true},
	}
	light := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "ds-lite", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-flash"}},
		model:         "deepseek-v4-flash",
		caps:          &goullm.Capabilities{ToolCalls: true},
	}

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"heavy":   heavy,
			"light":   light,
		},
	}
	result := buildModelCapabilityPrompt(req)
	assert.Contains(t, result, "deepseek-v4-flash")
	assert.Contains(t, result, "deepseek-v4-pro")
	assert.Contains(t, result, "opus")
	assert.Contains(t, result, "haiku")
	assert.Contains(t, result, "thinking")
	assert.Contains(t, result, "tool_calls")
	assert.Contains(t, result, "sub-agent")
}

func TestBuildModelCapabilityPrompt_VisionGuidance(t *testing.T) {
	primary := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "ds-flash", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-flash"}},
		model:         "deepseek-v4-flash",
		caps:          &goullm.Capabilities{ToolCalls: true},
	}
	visionConn := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "vision", typ: connector.OPENAI, settings: map[string]interface{}{"model": "gpt-4o"}},
		model:         "gpt-4o",
		caps:          &goullm.Capabilities{Vision: true, ToolCalls: true},
	}

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"vision":  visionConn,
		},
	}
	result := buildModelCapabilityPrompt(req)
	assert.Contains(t, result, "image_read")
	assert.Contains(t, result, "Image/Vision")
}

func TestBuildModelCapabilityPrompt_PrimaryHasVision_NoGuidance(t *testing.T) {
	primary := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "gpt4o", typ: connector.OPENAI, settings: map[string]interface{}{"model": "gpt-4o"}},
		model:         "gpt-4o",
		caps:          &goullm.Capabilities{Vision: true, ToolCalls: true},
	}

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"vision":  primary,
		},
	}
	result := buildModelCapabilityPrompt(req)
	assert.NotContains(t, result, "image_read", "should not suggest image_read when primary has vision")
}

func TestBuildModelCapabilityPrompt_ThinkingFromSettings(t *testing.T) {
	primary := &fakeLLMConnector{
		fakeConnector: fakeConnector{
			id:  "ds-pro",
			typ: connector.OPENAI,
			settings: map[string]interface{}{
				"model": "deepseek-v4-pro",
				"thinking": map[string]interface{}{
					"type": "enabled",
				},
			},
		},
		model: "deepseek-v4-pro",
		caps:  &goullm.Capabilities{ToolCalls: true},
	}
	heavy := &fakeLLMConnector{
		fakeConnector: fakeConnector{id: "ds-pro2", typ: connector.OPENAI, settings: map[string]interface{}{"model": "deepseek-v4-pro-max"}},
		model:         "deepseek-v4-pro-max",
		caps:          &goullm.Capabilities{ToolCalls: true},
	}

	req := &types.StreamRequest{
		Config:    &types.SandboxConfig{},
		Connector: primary,
		Roles: map[string]connector.Connector{
			"default": primary,
			"heavy":   heavy,
		},
	}
	result := buildModelCapabilityPrompt(req)
	assert.Contains(t, result, "thinking", "should detect thinking from Setting()[\"thinking\"]")
}