From 28f2e3a6acd223d73b5cc1984de48912c8aa03e1 Mon Sep 17 00:00:00 2001 From: IronRookieCoder Date: Thu, 14 May 2026 14:50:30 +0800 Subject: [PATCH] fix(agent-team): handle idle teammate coordination failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Bug 详情 Agent Team 在 CoStrict 环境下可能出现子代理无法执行、完成后主代理持续 Idle、shutdown 后 TeamDelete 被阻塞，以及 TeamCreate 参数错误时抛出 name.replace 深层异常。 ## 根因 - Team teammate 默认模型路径未复用普通 subagent 的 CoStrict 继承逻辑，未配置 teammateDefaultModel 时会回退到 CoStrict 不可用的 claude-opus-4-6。 - in-process teammate 的 AsyncLocalStorage 上下文可能影响 team lead 的 inbox polling 身份判断。 - SendMessage 收到字符串形式的 shutdown_response JSON 时按普通消息转发，未进入 shutdown approval 分支。 - ExecuteExtraTool 调用 deferred tool 前未执行目标工具 validateInput，错误参数会绕过入口校验进入目标工具内部。 ## 修复方案 - 将 teammate 模型解析集中到 teammateModel，并在 CoStrict 默认场景继承 leader model。 - 抽出 inbox poller 身份判断，基于 AppState 识别 team lead polling。 - 兼容字符串 JSON 协议响应，将 shutdown_response / plan_approval_response 路由到结构化处理逻辑。 - ExecuteExtraTool 在权限检查和实际执行前调用目标工具 validateInput。 ## 变更要点 - 新增 teammate 模型解析和 CoStrict 继承测试。 - 新增 inbox poller identity helper 和 in-process 上下文回归测试。 - 新增 SendMessage 字符串 shutdown_response 回归测试。 - 新增 ExecuteExtraTool 目标工具校验回归测试。 ## 自测 - bun test packages/builtin-tools/src/tools/ExecuteTool/__tests__/ExecuteTool.test.ts - bun test packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts src/utils/swarm/__tests__/teammateModel.test.ts src/hooks/__tests__/useInboxPoller.test.ts - bunx biome check 相关变更文件 - 变更文件相关 tsc 过滤检查无输出 --- .../src/tools/ExecuteTool/ExecuteTool.ts | 18 ++++ .../__tests__/ExecuteTool.runner.ts | 50 +++++++++++ .../tools/SendMessageTool/SendMessageTool.ts | 88 ++++++++++++++++++- .../udsRecipientSanitization.test.ts | 84 +++++++++++++++++- .../src/tools/shared/spawnMultiAgent.ts | 35 +------- src/hooks/__tests__/useInboxPoller.test.ts | 52 +++++++++++ src/hooks/inboxPollerIdentity.ts | 46 ++++++++++ src/hooks/useInboxPoller.ts | 60 ++----------- .../swarm/__tests__/teammateModel.test.ts | 85 ++++++++++++++++++ src/utils/swarm/teammateModel.ts | 31 +++++++ 10 files changed, 461 insertions(+), 88 deletions(-) create mode 100644 src/hooks/__tests__/useInboxPoller.test.ts create mode 100644 src/hooks/inboxPollerIdentity.ts create mode 100644 src/utils/swarm/__tests__/teammateModel.test.ts diff --git a/packages/builtin-tools/src/tools/ExecuteTool/ExecuteTool.ts b/packages/builtin-tools/src/tools/ExecuteTool/ExecuteTool.ts index 52103bb1b..2cd78c03b 100644 --- a/packages/builtin-tools/src/tools/ExecuteTool/ExecuteTool.ts +++ b/packages/builtin-tools/src/tools/ExecuteTool/ExecuteTool.ts @@ -89,6 +89,24 @@ export const ExecuteTool = buildTool({ } } + const validationResult = await targetTool.validateInput?.( + input.params as Record, + context, + ) + if (validationResult && validationResult.result === false) { + return { + data: { + result: null, + tool_name: input.tool_name, + }, + newMessages: [ + createUserMessage({ + content: `Invalid input for tool "${input.tool_name}": ${validationResult.message}`, + }), + ], + } + } + // Check permissions on the target tool const permResult = await targetTool.checkPermissions?.( input.params as Record, diff --git a/packages/builtin-tools/src/tools/ExecuteTool/__tests__/ExecuteTool.runner.ts b/packages/builtin-tools/src/tools/ExecuteTool/__tests__/ExecuteTool.runner.ts index eb10facf1..eaead1a27 100644 --- a/packages/builtin-tools/src/tools/ExecuteTool/__tests__/ExecuteTool.runner.ts +++ b/packages/builtin-tools/src/tools/ExecuteTool/__tests__/ExecuteTool.runner.ts @@ -71,6 +71,14 @@ function makeMockTool(name: string, callResult: unknown = 'ok') { return { name, call: async () => ({ data: callResult }), + validateInput: undefined as + | (( + input: Record, + ) => Promise< + | { result: true } + | { result: false; message: string; errorCode: number } + >) + | undefined, checkPermissions: async () => ({ behavior: 'allow' as const }), prompt: async () => `Description for ${name}`, description: async () => `Description for ${name}`, @@ -154,6 +162,48 @@ describe('ExecuteTool', () => { expect(result.newMessages).toBeDefined() }) + test('validates target tool input before calling it', async () => { + let called = false + const mockTarget = makeMockTool('TeamCreate', { created: true }) + mockTarget.validateInput = async input => + 'team_name' in input + ? { result: true } + : { + result: false, + message: 'team_name is required for TeamCreate', + errorCode: 9, + } + mockTarget.call = async () => { + called = true + return { data: { created: true } } + } + const ctx = makeContext([mockTarget]) + + const result = await ExecuteTool.call( + { + tool_name: 'TeamCreate', + params: { + name: 'worker-a', + description: 'wrong shape', + task: 'do work', + }, + }, + ctx, + async () => ({ behavior: 'allow' }), + { type: 'assistant', content: [], uuid: 'msg1' } as never, + undefined, + ) + + expect(called).toBe(false) + expect(result.data).toEqual({ + result: null, + tool_name: 'TeamCreate', + }) + expect(result.newMessages?.[0]?.content).toContain( + 'team_name is required for TeamCreate', + ) + }) + test('has correct name', () => { expect(ExecuteTool.name).toBe(EXECUTE_TOOL_NAME) }) diff --git a/packages/builtin-tools/src/tools/SendMessageTool/SendMessageTool.ts b/packages/builtin-tools/src/tools/SendMessageTool/SendMessageTool.ts index 82c9a04c3..c0c086af5 100644 --- a/packages/builtin-tools/src/tools/SendMessageTool/SendMessageTool.ts +++ b/packages/builtin-tools/src/tools/SendMessageTool/SendMessageTool.ts @@ -17,6 +17,7 @@ import { logForDebugging } from 'src/utils/debug.js' import { errorMessage } from 'src/utils/errors.js' import { truncate } from 'src/utils/format.js' import { gracefulShutdown } from 'src/utils/gracefulShutdown.js' +import { safeParseJSON } from 'src/utils/json.js' import { lazySchema } from 'src/utils/lazySchema.js' import { parseAddress } from 'src/utils/peerAddress.js' import { semanticBoolean } from 'src/utils/semanticBoolean.js' @@ -86,6 +87,7 @@ const inputSchema = lazySchema(() => }), ) type InputSchema = ReturnType +type StructuredMessageInput = z.infer> export type Input = z.infer @@ -130,6 +132,24 @@ export type SendMessageToolOutput = | RequestOutput | ResponseOutput +function parseProtocolResponseString( + content: string, +): Extract< + StructuredMessageInput, + { type: 'shutdown_response' | 'plan_approval_response' } +> | null { + const parsed = safeParseJSON(content, false) + const result = StructuredMessage().safeParse(parsed) + if (!result.success) return null + if ( + result.data.type !== 'shutdown_response' && + result.data.type !== 'plan_approval_response' + ) { + return null + } + return result.data +} + const UDS_INLINE_TOKEN_MARKER = '#token=' function stripInlineUdsToken(target: string): string { @@ -725,7 +745,33 @@ export const SendMessageTool: Tool = return { result: true } } if (typeof input.message === 'string') { - if (!input.summary || input.summary.trim().length === 0) { + const protocolResponse = parseProtocolResponseString(input.message) + if ( + protocolResponse?.type === 'shutdown_response' && + input.to !== TEAM_LEAD_NAME + ) { + return { + result: false, + message: `shutdown_response must be sent to "${TEAM_LEAD_NAME}"`, + errorCode: 9, + } + } + if ( + protocolResponse?.type === 'shutdown_response' && + !protocolResponse.approve && + (!protocolResponse.reason || + protocolResponse.reason.trim().length === 0) + ) { + return { + result: false, + message: 'reason is required when rejecting a shutdown request', + errorCode: 9, + } + } + if ( + !protocolResponse && + (!input.summary || input.summary.trim().length === 0) + ) { return { result: false, message: 'summary is required when message is a string', @@ -799,6 +845,11 @@ export const SendMessageTool: Tool = }, async call(input, context, canUseTool, assistantMessage) { + const protocolResponse = + typeof input.message === 'string' + ? parseProtocolResponseString(input.message) + : null + if (typeof input.message === 'string') { const addr = parseAddress(input.to) if (addr.scheme === 'uds' && hasInlineUdsToken(input.to)) { @@ -908,7 +959,11 @@ export const SendMessageTool: Tool = // Route to in-process subagent by name or raw agentId before falling // through to ambient-team resolution. Stopped agents are auto-resumed. - if (typeof input.message === 'string' && input.to !== '*') { + if ( + typeof input.message === 'string' && + input.to !== '*' && + !protocolResponse + ) { const appState = context.getAppState() const registered = appState.agentNameRegistry.get(input.to) const agentId = registered ?? toAgentId(input.to) @@ -987,6 +1042,35 @@ export const SendMessageTool: Tool = } if (typeof input.message === 'string') { + if (protocolResponse) { + switch (protocolResponse.type) { + case 'shutdown_response': + if (protocolResponse.approve) { + return handleShutdownApproval( + protocolResponse.request_id, + context, + ) + } + return handleShutdownRejection( + protocolResponse.request_id, + protocolResponse.reason!, + ) + case 'plan_approval_response': + if (protocolResponse.approve) { + return handlePlanApproval( + input.to, + protocolResponse.request_id, + context, + ) + } + return handlePlanRejection( + input.to, + protocolResponse.request_id, + protocolResponse.feedback!, + context, + ) + } + } if (input.to === '*') { return handleBroadcast(input.message, input.summary, context) } diff --git a/packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts b/packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts index e0ce1a823..f0cfe0b28 100644 --- a/packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts +++ b/packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts @@ -1,4 +1,60 @@ -import { describe, expect, test } from 'bun:test' +import { describe, expect, mock, test } from 'bun:test' + +type MailboxWrite = { + recipient: string + message: { text: string; from?: string } + team?: string +} + +let lastMailboxWrite: MailboxWrite | undefined +const writeToMailboxMock = mock( + async ( + recipient: string, + message: { text: string; from?: string }, + team?: string, + ) => { + lastMailboxWrite = { recipient, message, team } + }, +) +mock.module('src/utils/teammateMailbox.js', () => ({ + createShutdownApprovedMessage: (params: { + requestId: string + from: string + paneId?: string + backendType?: string + }) => ({ + type: 'shutdown_approved', + requestId: params.requestId, + from: params.from, + timestamp: '2026-05-14T00:00:00.000Z', + paneId: params.paneId, + backendType: params.backendType, + }), + createShutdownRejectedMessage: (params: { + requestId: string + from: string + reason: string + }) => ({ + type: 'shutdown_rejected', + requestId: params.requestId, + from: params.from, + reason: params.reason, + timestamp: '2026-05-14T00:00:00.000Z', + }), + createShutdownRequestMessage: (params: { + requestId: string + from: string + reason?: string + }) => ({ + type: 'shutdown_request', + requestId: params.requestId, + from: params.from, + reason: params.reason, + timestamp: '2026-05-14T00:00:00.000Z', + }), + writeToMailbox: writeToMailboxMock, +})) + import { SendMessageTool } from '../SendMessageTool.js' describe('SendMessageTool UDS recipient handling', () => { @@ -178,4 +234,30 @@ describe('SendMessageTool UDS recipient handling', () => { expect(result.data.success).toBe(false) expect(JSON.stringify(result)).not.toContain('secret-token') }) + + test('handles shutdown_response sent as a JSON string protocol response', async () => { + writeToMailboxMock.mockClear() + lastMailboxWrite = undefined + + const result = await SendMessageTool.call( + { + to: 'team-lead', + summary: 'Approving shutdown', + message: + '{"type":"shutdown_response","request_id":"shutdown-worker-a","approve":true}', + }, + { + getAppState: () => ({ tasks: {} }), + } as never, + undefined as never, + undefined as never, + ) + + expect(result.data.success).toBe(true) + expect(result.data.message).toContain('Shutdown approved') + const mailboxWrite = lastMailboxWrite as MailboxWrite | undefined + expect(mailboxWrite?.recipient).toBe('team-lead') + expect(mailboxWrite?.message.from).toBe('teammate') + expect(mailboxWrite?.message.text).toContain('shutdown_approved') + }) }) diff --git a/packages/builtin-tools/src/tools/shared/spawnMultiAgent.ts b/packages/builtin-tools/src/tools/shared/spawnMultiAgent.ts index 813afd9b3..e1a022739 100644 --- a/packages/builtin-tools/src/tools/shared/spawnMultiAgent.ts +++ b/packages/builtin-tools/src/tools/shared/spawnMultiAgent.ts @@ -8,10 +8,8 @@ import React from 'react' import { getSessionId } from 'src/bootstrap/state.js' import type { ToolUseContext } from 'src/Tool.js' import { formatAgentId } from 'src/utils/agentId.js' -import { getGlobalConfig } from 'src/utils/config.js' import { getCwd } from 'src/utils/cwd.js' import { logForDebugging } from 'src/utils/debug.js' -import { parseUserSpecifiedModel } from 'src/utils/model/model.js' import { getTeammateExecutor } from 'src/utils/swarm/backends/registry.js' import type { BackendType, @@ -30,41 +28,10 @@ import { type TeamFile, } from 'src/utils/swarm/teamHelpers.js' import { assignTeammateColor } from 'src/utils/swarm/teammateLayoutManager.js' -import { getHardcodedTeammateModelFallback } from 'src/utils/swarm/teammateModel.js' +import { resolveTeammateModel } from 'src/utils/swarm/teammateModel.js' import type { CustomAgentDefinition } from '../AgentTool/loadAgentsDir.js' import { isCustomAgent } from '../AgentTool/loadAgentsDir.js' -function getDefaultTeammateModel(leaderModel: string | null): string { - const configured = getGlobalConfig().teammateDefaultModel - if (configured === null) { - // User picked "Default" in the /config picker — follow the leader. - return leaderModel ?? getHardcodedTeammateModelFallback() - } - if (configured !== undefined) { - return parseUserSpecifiedModel(configured) - } - return getHardcodedTeammateModelFallback() -} - -/** - * Resolve a teammate model value. Handles the 'inherit' alias (from agent - * frontmatter) by substituting the leader's model. gh-31069: 'inherit' was - * passed literally to --model, producing "It may not exist or you may not - * have access". If leader model is null (not yet set), falls through to the - * default. - * - * Exported for testing. - */ -export function resolveTeammateModel( - inputModel: string | undefined, - leaderModel: string | null, -): string { - if (inputModel === 'inherit') { - return leaderModel ?? getDefaultTeammateModel(leaderModel) - } - return inputModel ?? getDefaultTeammateModel(leaderModel) -} - // ============================================================================ // Types // ============================================================================ diff --git a/src/hooks/__tests__/useInboxPoller.test.ts b/src/hooks/__tests__/useInboxPoller.test.ts new file mode 100644 index 000000000..c4cc04fa9 --- /dev/null +++ b/src/hooks/__tests__/useInboxPoller.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, test } from 'bun:test' +import type { AppState } from '../../state/AppState' +import { + createTeammateContext, + runWithTeammateContext, +} from '../../utils/teammateContext' +import { getAgentNameToPoll, isPollingAsTeamLead } from '../inboxPollerIdentity' + +function createLeaderState(): AppState { + return { + teamContext: { + teamName: 'alpha', + teamFilePath: '/tmp/alpha/config.json', + leadAgentId: 'team-lead@alpha', + teammates: { + 'team-lead@alpha': { + name: 'team-lead', + tmuxSessionName: '', + tmuxPaneId: '', + cwd: '/tmp', + spawnedAt: 1, + }, + 'worker@alpha': { + name: 'worker', + tmuxSessionName: 'in-process', + tmuxPaneId: 'in-process', + cwd: '/tmp', + spawnedAt: 2, + }, + }, + }, + } as unknown as AppState +} + +describe('useInboxPoller identity helpers', () => { + test('keeps polling as team lead during in-process teammate async context', () => { + const appState = createLeaderState() + const teammateContext = createTeammateContext({ + agentId: 'worker@alpha', + agentName: 'worker', + teamName: 'alpha', + planModeRequired: false, + parentSessionId: 'leader-session', + abortController: new AbortController(), + }) + + runWithTeammateContext(teammateContext, () => { + expect(getAgentNameToPoll(appState)).toBe('team-lead') + expect(isPollingAsTeamLead(appState, 'team-lead')).toBe(true) + }) + }) +}) diff --git a/src/hooks/inboxPollerIdentity.ts b/src/hooks/inboxPollerIdentity.ts new file mode 100644 index 000000000..68940fa08 --- /dev/null +++ b/src/hooks/inboxPollerIdentity.ts @@ -0,0 +1,46 @@ +import type { AppState } from '../state/AppState.js' +import { TEAM_LEAD_NAME } from '../utils/swarm/constants.js' +import { getAgentName, isTeammate } from '../utils/teammate.js' +import { isInProcessTeammate } from '../utils/teammateContext.js' + +export function getAgentNameToPoll(appState: AppState): string | undefined { + const teamLeadName = getTeamLeadNameToPoll(appState) + + if (isTeammate() && !isInProcessTeammate()) { + return getAgentName() + } + + if (teamLeadName) { + return teamLeadName + } + + // In-process teammates should NOT use useInboxPoller - they have their own + // polling mechanism via waitForNextPromptOrShutdown() in inProcessRunner.ts. + // Using useInboxPoller would cause message routing issues since in-process + // teammates share the same React context and AppState with the leader. + // + // Note: This can be called when the leader's REPL re-renders while an + // in-process teammate's AsyncLocalStorage context is active (due to shared + // setAppState). We return undefined to gracefully skip polling rather than + // throwing, since this is a normal occurrence during concurrent execution. + if (isInProcessTeammate()) { + return undefined + } + return undefined +} + +function getTeamLeadNameToPoll(appState: AppState): string | undefined { + const teamContext = appState.teamContext + if (!teamContext?.leadAgentId) { + return undefined + } + + return teamContext.teammates[teamContext.leadAgentId]?.name || TEAM_LEAD_NAME +} + +export function isPollingAsTeamLead( + appState: AppState, + agentName: string, +): boolean { + return getTeamLeadNameToPoll(appState) === agentName +} diff --git a/src/hooks/useInboxPoller.ts b/src/hooks/useInboxPoller.ts index ac2c98555..5c256328b 100644 --- a/src/hooks/useInboxPoller.ts +++ b/src/hooks/useInboxPoller.ts @@ -44,10 +44,8 @@ import { unassignTeammateTasks } from '../utils/tasks.js' import { getAgentName, isPlanModeRequired, - isTeamLead, isTeammate, } from '../utils/teammate.js' -import { isInProcessTeammate } from '../utils/teammateContext.js' import { isModeSetRequest, isPermissionRequest, @@ -70,39 +68,10 @@ import { processMailboxPermissionResponse, processSandboxPermissionResponse, } from './useSwarmPermissionPoller.js' - -/** - * Get the agent name to poll for messages. - * - In-process teammates return undefined (they use waitForNextPromptOrShutdown instead) - * - Process-based teammates use their CLAUDE_CODE_AGENT_NAME - * - Team leads use their name from teamContext.teammates - * - Standalone sessions return undefined - */ -function getAgentNameToPoll(appState: AppState): string | undefined { - // In-process teammates should NOT use useInboxPoller - they have their own - // polling mechanism via waitForNextPromptOrShutdown() in inProcessRunner.ts. - // Using useInboxPoller would cause message routing issues since in-process - // teammates share the same React context and AppState with the leader. - // - // Note: This can be called when the leader's REPL re-renders while an - // in-process teammate's AsyncLocalStorage context is active (due to shared - // setAppState). We return undefined to gracefully skip polling rather than - // throwing, since this is a normal occurrence during concurrent execution. - if (isInProcessTeammate()) { - return undefined - } - if (isTeammate()) { - return getAgentName() - } - // Team lead polls using their agent name (not ID) - if (isTeamLead(appState.teamContext)) { - const leadAgentId = appState.teamContext!.leadAgentId - // Look up the lead's name from teammates map - const leadName = appState.teamContext!.teammates[leadAgentId]?.name - return leadName || 'team-lead' - } - return undefined -} +import { + getAgentNameToPoll, + isPollingAsTeamLead, +} from './inboxPollerIdentity.js' const INBOX_POLL_INTERVAL_MS = 1000 @@ -143,6 +112,7 @@ export function useInboxPoller({ const currentAppState = store.getState() const agentName = getAgentNameToPoll(currentAppState) if (!agentName) return + const pollingAsTeamLead = isPollingAsTeamLead(currentAppState, agentName) const unread = await readUnreadMessages( agentName, @@ -248,10 +218,7 @@ export function useInboxPoller({ } // Handle permission requests (leader side) - route to ToolUseConfirmQueue - if ( - permissionRequests.length > 0 && - isTeamLead(currentAppState.teamContext) - ) { + if (permissionRequests.length > 0 && pollingAsTeamLead) { logForDebugging( `[InboxPoller] Found ${permissionRequests.length} permission request(s)`, ) @@ -397,10 +364,7 @@ export function useInboxPoller({ } // Handle sandbox permission requests (leader side) - add to workerSandboxPermissions queue - if ( - sandboxPermissionRequests.length > 0 && - isTeamLead(currentAppState.teamContext) - ) { + if (sandboxPermissionRequests.length > 0 && pollingAsTeamLead) { logForDebugging( `[InboxPoller] Found ${sandboxPermissionRequests.length} sandbox permission request(s)`, ) @@ -597,10 +561,7 @@ export function useInboxPoller({ } // Handle plan approval requests (leader side) - auto-approve and write response to teammate inbox - if ( - planApprovalRequests.length > 0 && - isTeamLead(currentAppState.teamContext) - ) { + if (planApprovalRequests.length > 0 && pollingAsTeamLead) { logForDebugging( `[InboxPoller] Found ${planApprovalRequests.length} plan approval request(s), auto-approving`, ) @@ -675,10 +636,7 @@ export function useInboxPoller({ } // Handle shutdown approvals (leader side) - kill the teammate's pane - if ( - shutdownApprovals.length > 0 && - isTeamLead(currentAppState.teamContext) - ) { + if (shutdownApprovals.length > 0 && pollingAsTeamLead) { logForDebugging( `[InboxPoller] Found ${shutdownApprovals.length} shutdown approval(s)`, ) diff --git a/src/utils/swarm/__tests__/teammateModel.test.ts b/src/utils/swarm/__tests__/teammateModel.test.ts new file mode 100644 index 000000000..e69ced12d --- /dev/null +++ b/src/utils/swarm/__tests__/teammateModel.test.ts @@ -0,0 +1,85 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { saveGlobalConfig } from 'src/utils/config.js' +import { + resetSettingsCache, + setSessionSettingsCache, +} from 'src/utils/settings/settingsCache.js' +import { resolveTeammateModel } from '../teammateModel.js' + +const providerEnvKeys = [ + 'CLAUDE_CODE_USE_BEDROCK', + 'CLAUDE_CODE_USE_VERTEX', + 'CLAUDE_CODE_USE_FOUNDRY', + 'CLAUDE_CODE_USE_OPENAI', + 'CLAUDE_CODE_USE_GEMINI', + 'CLAUDE_CODE_USE_GROK', + 'CLAUDE_CODE_USE_COSTRICT', +] as const + +let previousProviderEnv: Record< + (typeof providerEnvKeys)[number], + string | undefined +> + +function clearTeammateDefaultModel(): void { + resetSettingsCache() + setSessionSettingsCache({ settings: {}, errors: [] }) + saveGlobalConfig(config => { + const next = { ...config } + delete next.teammateDefaultModel + return next + }) +} + +beforeEach(() => { + previousProviderEnv = {} as Record< + (typeof providerEnvKeys)[number], + string | undefined + > + for (const key of providerEnvKeys) { + previousProviderEnv[key] = process.env[key] + delete process.env[key] + } + clearTeammateDefaultModel() +}) + +afterEach(() => { + for (const key of providerEnvKeys) { + if (previousProviderEnv[key] === undefined) { + delete process.env[key] + } else { + process.env[key] = previousProviderEnv[key] + } + } + clearTeammateDefaultModel() +}) + +describe('resolveTeammateModel', () => { + test('inherits the leader model by default for CoStrict', () => { + process.env.CLAUDE_CODE_USE_COSTRICT = '1' + + expect(resolveTeammateModel(undefined, 'CoStrict-DeepSeek-V4-Pro')).toBe( + 'CoStrict-DeepSeek-V4-Pro', + ) + }) + + test('keeps the hardcoded Opus fallback by default for first-party', () => { + expect(resolveTeammateModel(undefined, 'claude-sonnet-4-5')).toBe( + 'claude-opus-4-6', + ) + }) + + test('keeps explicit teammate model values', () => { + process.env.CLAUDE_CODE_USE_COSTRICT = '1' + + expect( + resolveTeammateModel('deepseek-v4-flash', 'deepseek-v4-pro[1m]'), + ).toBe('deepseek-v4-flash') + }) + + test('keeps inherit alias behavior', () => { + expect(resolveTeammateModel('inherit', 'claude-sonnet-4-5')).toBe( + 'claude-sonnet-4-5', + ) + }) +}) diff --git a/src/utils/swarm/teammateModel.ts b/src/utils/swarm/teammateModel.ts index 0a49009b2..833063a31 100644 --- a/src/utils/swarm/teammateModel.ts +++ b/src/utils/swarm/teammateModel.ts @@ -1,5 +1,7 @@ import { CLAUDE_OPUS_4_6_CONFIG } from '../model/configs.js' +import { parseUserSpecifiedModel } from '../model/model.js' import { getAPIProvider } from '../model/providers.js' +import { getGlobalConfig } from '../config.js' // @[MODEL LAUNCH]: Update the fallback model below. // When the user has never set teammateDefaultModel in /config, new teammates @@ -8,3 +10,32 @@ import { getAPIProvider } from '../model/providers.js' export function getHardcodedTeammateModelFallback(): string { return CLAUDE_OPUS_4_6_CONFIG[getAPIProvider()] } + +export function getDefaultTeammateModel(leaderModel: string | null): string { + const configured = getGlobalConfig().teammateDefaultModel + if (configured === null) { + // User picked "Default" in the /config picker — follow the leader. + return leaderModel ?? getHardcodedTeammateModelFallback() + } + if (configured !== undefined) { + return parseUserSpecifiedModel(configured) + } + if (getAPIProvider() === 'costrict' && leaderModel) { + return leaderModel + } + return getHardcodedTeammateModelFallback() +} + +/** + * Resolve a teammate model value. Handles the 'inherit' alias (from agent + * frontmatter) by substituting the leader's model. + */ +export function resolveTeammateModel( + inputModel: string | undefined, + leaderModel: string | null, +): string { + if (inputModel === 'inherit') { + return leaderModel ?? getDefaultTeammateModel(leaderModel) + } + return inputModel ?? getDefaultTeammateModel(leaderModel) +}