diff --git a/src/agents/tool-restrictions.test.ts b/src/agents/tool-restrictions.test.ts index f4e881b04..9f80c1617 100644 --- a/src/agents/tool-restrictions.test.ts +++ b/src/agents/tool-restrictions.test.ts @@ -54,6 +54,20 @@ describe("read-only agent tool restrictions", () => { } }) + test("allows team tools for team member prompt restrictions", () => { + // given + const teamMemberAgentName = "sisyphus-junior" + + // when + const restrictions = getAgentToolRestrictions(teamMemberAgentName, { includeTeamToolDenylist: false }) + + // then + for (const toolName of TEAM_TOOL_NAMES) { + expect(restrictions[toolName]).toBeUndefined() + } + expect(restrictions.task).toBe(false) + }) + describe("Oracle", () => { test("denies all file-writing tools", () => { // given diff --git a/src/features/background-agent/fallback-retry-handler.test.ts b/src/features/background-agent/fallback-retry-handler.test.ts index a9c4f0cd1..f51ce2059 100644 --- a/src/features/background-agent/fallback-retry-handler.test.ts +++ b/src/features/background-agent/fallback-retry-handler.test.ts @@ -2,7 +2,7 @@ import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test" const sharedLogMock = mock(() => {}) const readConnectedProvidersCacheMock = mock(() => null) -const readProviderModelsCacheMock = mock(() => null) +const readProviderModelsCacheMock = mock((): { connected: string[] } | null => null) const shouldRetryErrorMock = mock(() => true) const getNextFallbackMock = mock((chain: Array<{ model: string }>, attempt: number) => chain[attempt]) const hasMoreFallbacksMock = mock((chain: Array<{ model: string }>, attempt: number) => attempt < chain.length) @@ -88,7 +88,7 @@ function createMockConcurrencyManager(): ConcurrencyManager { acquire: mock(async () => {}), getQueueLength: mock(() => 0), getActiveCount: mock(() => 0), - } as unknown as ConcurrencyManager + } as never } function createMockClient(): { @@ -101,7 +101,7 @@ function createMockClient(): { session: { abort: abortMock, }, - } as unknown as OpencodeClient, + } as never, abortMock, } } @@ -133,9 +133,9 @@ describe("tryFallbackRetry", () => { }) beforeEach(() => { - ;(shouldRetryError as any).mockImplementation(() => true) - ;(selectFallbackProvider as any).mockImplementation((providers: string[]) => providers[0]) - ;(readProviderModelsCache as any).mockReturnValue(null) + shouldRetryError.mockImplementation(() => true) + selectFallbackProvider.mockImplementation((providers: string[]) => providers[0]) + readProviderModelsCache.mockReturnValue(null) }) describe("#given retryable error with fallback chain", () => { @@ -260,6 +260,21 @@ describe("tryFallbackRetry", () => { expect(args.processKey).toHaveBeenCalledWith(key) }) + test("preserves team identity and session callback in retry input", async () => { + const onSessionCreated = mock(async () => {}) + const args = createDefaultArgs({ + teamRunId: "team-run-1", + onSessionCreated, + }) + + await tryFallbackRetry(args) + + const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` + const retryInput = args.queuesByKey.get(key)?.[0]?.input + expect(retryInput?.teamRunId).toBe("team-run-1") + expect(retryInput?.onSessionCreated).toBe(onSessionCreated) + }) + test("finalizes the failed attempt, creates a new pending attempt, and enqueues its explicit attemptID", async () => { const args = createDefaultArgs({ status: "running", @@ -308,13 +323,16 @@ describe("tryFallbackRetry", () => { const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` const queue = args.queuesByKey.get(key) expect(queue).toBeDefined() - expect((queue?.[0] as QueueItem & { attemptID?: string })?.attemptID).toBe(nextAttempt?.attemptId) + const queuedAttemptID = queue?.[0]?.attemptID + expect(queuedAttemptID).toBeDefined() + expect(nextAttempt?.attemptId).toBeDefined() + expect(queuedAttemptID).toBe(nextAttempt?.attemptId ?? "") }) }) describe("#given non-retryable error", () => { test("returns false when shouldRetryError returns false", async () => { - ;(shouldRetryError as any).mockImplementation(() => false) + shouldRetryError.mockImplementation(() => false) const args = createDefaultArgs() const result = await tryFallbackRetry(args) @@ -415,8 +433,8 @@ describe("tryFallbackRetry", () => { describe("#given disconnected fallback providers with connected preferred provider", () => { test("keeps fallback entry and selects connected preferred provider", async () => { - ;(readProviderModelsCache as any).mockReturnValueOnce({ connected: ["provider-a"] }) - ;(selectFallbackProvider as any).mockImplementationOnce( + readProviderModelsCache.mockReturnValueOnce({ connected: ["provider-a"] }) + selectFallbackProvider.mockImplementationOnce( (_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b", ) diff --git a/src/features/background-agent/fallback-retry-handler.ts b/src/features/background-agent/fallback-retry-handler.ts index f5f31a6ad..d61a839b1 100644 --- a/src/features/background-agent/fallback-retry-handler.ts +++ b/src/features/background-agent/fallback-retry-handler.ts @@ -170,10 +170,12 @@ export async function tryFallbackRetry(args: { parentModel: task.parentModel, parentAgent: task.parentAgent, parentTools: task.parentTools, + teamRunId: task.teamRunId, model: nextModel, fallbackChain: task.fallbackChain, category: task.category, isUnstableAgent: task.isUnstableAgent, + onSessionCreated: task.onSessionCreated, } if (previousSessionID) { diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index 10dce905d..1b0d4acdf 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -424,6 +424,7 @@ export class BackgroundManager { fallbackChain: input.fallbackChain, attemptCount: 0, category: input.category, + onSessionCreated: input.onSessionCreated, } const firstAttempt = startAttempt(task, input.model) @@ -722,7 +723,9 @@ The fallback retry session is now created and can be inspected directly. task: false, call_omo_agent: true, question: false, - ...getAgentToolRestrictions(input.agent), + ...getAgentToolRestrictions(input.agent, { + includeTeamToolDenylist: input.teamRunId === undefined, + }), } setSessionTools(sessionID, tools) return tools @@ -742,7 +745,9 @@ The fallback retry session is now created and can be inspected directly. taskId: task.id, }) try { - const fallbackBody = buildFallbackBody(promptBody, FALLBACK_AGENT) + const fallbackBody = buildFallbackBody(promptBody, FALLBACK_AGENT, { + includeTeamToolDenylist: input.teamRunId === undefined, + }) setSessionTools(sessionID, fallbackBody.tools as Record) await promptWithModelSuggestionRetry(this.client, { path: { id: sessionID }, @@ -1103,7 +1108,9 @@ The fallback retry session is now created and can be inspected directly. task: false, call_omo_agent: true, question: false, - ...getAgentToolRestrictions(existingTask.agent), + ...getAgentToolRestrictions(existingTask.agent, { + includeTeamToolDenylist: existingTask.teamRunId === undefined, + }), } setSessionTools(existingTask.sessionId!, tools) return tools @@ -1584,7 +1591,7 @@ The fallback retry session is now created and can be inspected directly. }) } - private tryFallbackRetry( + private async tryFallbackRetry( task: BackgroundTask, errorInfo: { name?: string; message?: string }, source: string, @@ -1620,14 +1627,13 @@ The task was re-queued on a fallback model after a retryable failure. ) }, }) - return result.then((retried) => { - if (retried && previousSessionID) { - this.clearSessionOutputObserved(previousSessionID) - this.clearSessionTodoObservation(previousSessionID) - subagentSessions.delete(previousSessionID) - } - return retried - }) + const retried = await result + if (retried && previousSessionID) { + this.clearSessionOutputObserved(previousSessionID) + this.clearSessionTodoObservation(previousSessionID) + subagentSessions.delete(previousSessionID) + } + return retried } markForNotification(task: BackgroundTask): void { diff --git a/src/features/background-agent/spawner.test.ts b/src/features/background-agent/spawner.test.ts index 8a228866e..dc0ed17d0 100644 --- a/src/features/background-agent/spawner.test.ts +++ b/src/features/background-agent/spawner.test.ts @@ -29,7 +29,7 @@ describe("background-agent spawner agent-not-found fallback", () => { return { data: {} } }, }, - } as any + } as never const onTaskError = mock(() => {}) @@ -64,7 +64,7 @@ describe("background-agent spawner agent-not-found fallback", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) // Wait for the fire-and-forget prompt chain to settle await new Promise(resolve => setTimeout(resolve, 50)) @@ -76,11 +76,23 @@ describe("background-agent spawner agent-not-found fallback", () => { expect(promptCalls[1].body.agent).toBe("general") // Original prompt content preserved in fallback expect(promptCalls[1].body.parts).toEqual(promptCalls[0].body.parts) - // Tool restrictions recomputed for fallback agent (general has no restrictions) + // Tool restrictions recomputed for fallback agent while preserving delegated-subagent team tool denial expect(promptCalls[1].body.tools).toEqual({ task: false, call_omo_agent: true, question: false, + team_create: false, + team_delete: false, + team_shutdown_request: false, + team_approve_shutdown: false, + team_reject_shutdown: false, + team_send_message: false, + team_task_create: false, + team_task_list: false, + team_task_update: false, + team_task_get: false, + team_status: false, + team_list: false, }) // Task agent identity updated to reflect fallback expect(task.agent).toBe("general") @@ -101,7 +113,7 @@ describe("background-agent spawner agent-not-found fallback", () => { throw new Error("Connection timeout") }, }, - } as any + } as never const onTaskError = mock(() => {}) @@ -133,7 +145,7 @@ describe("background-agent spawner agent-not-found fallback", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise(resolve => setTimeout(resolve, 50)) //#then @@ -154,7 +166,7 @@ describe("background-agent spawner agent-not-found fallback", () => { throw new Error('Agent not found: "Sisyphus-Junior". Available agents: build, explore, general, plan') }, }, - } as any + } as never const onTaskError = mock(() => {}) @@ -186,7 +198,7 @@ describe("background-agent spawner agent-not-found fallback", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise(resolve => setTimeout(resolve, 50)) //#then @@ -213,7 +225,7 @@ describe("background-agent spawner agent-not-found fallback", () => { return { data: {} } }, }, - } as any + } as never const onTaskError = mock(() => {}) @@ -248,7 +260,7 @@ describe("background-agent spawner agent-not-found fallback", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise(resolve => setTimeout(resolve, 50)) //#then @@ -276,7 +288,7 @@ describe("background-agent spawner agent-not-found fallback", () => { return { data: {} } }, }, - } as any + } as never const onTaskError = mock(() => {}) @@ -311,7 +323,7 @@ describe("background-agent spawner agent-not-found fallback", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise(resolve => setTimeout(resolve, 50)) //#then @@ -338,11 +350,11 @@ describe("background-agent spawner fallback model promotion", () => { return { data: {} } }), }, - } as any + } as never const concurrencyManager = { release: mock(() => {}), - } as any + } as never const onTaskError = mock(() => {}) @@ -455,7 +467,7 @@ describe("background-agent spawner fallback model promotion", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) //#then expect(promptCalls).toHaveLength(1) @@ -569,7 +581,7 @@ describe("background-agent spawner fallback model promotion", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise((resolve) => setTimeout(resolve, 0)) //#then @@ -623,7 +635,7 @@ describe("background-agent spawner fallback model promotion", () => { } //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise((resolve) => setTimeout(resolve, 0)) //#then @@ -653,7 +665,7 @@ describe("background-agent spawner tmux callback ordering", () => { return { data: {} } }, }, - } as any + } as never const onSubagentSessionCreated = mock(async () => { events.push("tmux.callback.start") @@ -694,7 +706,7 @@ describe("background-agent spawner tmux callback ordering", () => { try { //#when - await startTask(item as any, ctx as any) + await startTask(item as never, ctx as never) await new Promise((resolve) => setTimeout(resolve, 20)) //#then diff --git a/src/features/background-agent/spawner.ts b/src/features/background-agent/spawner.ts index aefe15829..bfbe675da 100644 --- a/src/features/background-agent/spawner.ts +++ b/src/features/background-agent/spawner.ts @@ -28,6 +28,7 @@ export function isAgentNotFoundError(error: unknown): boolean { export function buildFallbackBody( originalBody: Record, fallbackAgent: string, + options: { includeTeamToolDenylist?: boolean } = {}, ): Record { return { ...originalBody, @@ -36,7 +37,7 @@ export function buildFallbackBody( task: false, call_omo_agent: true, question: false, - ...getAgentToolRestrictions(fallbackAgent), + ...getAgentToolRestrictions(fallbackAgent, options), }, } } @@ -60,9 +61,11 @@ export function createTask(input: LaunchInput): BackgroundTask { agent: input.agent, parentSessionId: input.parentSessionId, parentMessageId: input.parentMessageId, + teamRunId: input.teamRunId, parentModel: input.parentModel, parentAgent: input.parentAgent, model: input.model, + onSessionCreated: input.onSessionCreated, } } @@ -160,7 +163,9 @@ export async function startTask( task: false, call_omo_agent: true, question: false, - ...getAgentToolRestrictions(normalizedAgent), + ...getAgentToolRestrictions(normalizedAgent, { + includeTeamToolDenylist: input.teamRunId === undefined, + }), }, parts: [createInternalAgentTextPart(input.prompt)], } @@ -179,7 +184,9 @@ export async function startTask( try { await promptWithModelSuggestionRetry(client, { path: { id: sessionID }, - body: buildFallbackBody(promptBody, FALLBACK_AGENT), + body: buildFallbackBody(promptBody, FALLBACK_AGENT, { + includeTeamToolDenylist: input.teamRunId === undefined, + }), }) task.agent = FALLBACK_AGENT return @@ -294,7 +301,9 @@ export async function resumeTask( task: false, call_omo_agent: true, question: false, - ...getAgentToolRestrictions(task.agent), + ...getAgentToolRestrictions(task.agent, { + includeTeamToolDenylist: task.teamRunId === undefined, + }), }, parts: [createInternalAgentTextPart(input.prompt)], } @@ -312,7 +321,9 @@ export async function resumeTask( try { await promptWithModelSuggestionRetry(client, { path: { id: task.sessionId! }, - body: buildFallbackBody(resumeBody, FALLBACK_AGENT), + body: buildFallbackBody(resumeBody, FALLBACK_AGENT, { + includeTeamToolDenylist: task.teamRunId === undefined, + }), }) task.agent = FALLBACK_AGENT return diff --git a/src/features/background-agent/types.ts b/src/features/background-agent/types.ts index d72ab4a5e..76c242a41 100644 --- a/src/features/background-agent/types.ts +++ b/src/features/background-agent/types.ts @@ -77,6 +77,7 @@ export interface BackgroundTask { isUnstableAgent?: boolean /** Category used for this task (e.g., 'quick', 'visual-engineering') */ category?: string + onSessionCreated?: (sessionId: string) => void | Promise /** Pending retry notification details for the next spawned retry session */ retryNotification?: { previousSessionID?: string diff --git a/src/features/team-mode/team-runtime/create.ts b/src/features/team-mode/team-runtime/create.ts index 5bfdf9e01..7671b03ed 100644 --- a/src/features/team-mode/team-runtime/create.ts +++ b/src/features/team-mode/team-runtime/create.ts @@ -199,12 +199,18 @@ export async function createTeamRun( skillContent: resolvedMember.systemContent, category: member.kind === "category" ? member.category : undefined, sessionPermission: QUESTION_DENIED_SESSION_PERMISSION, - onSessionCreated: (sessionId) => { + onSessionCreated: async (sessionId) => { registerTeamSession(sessionId, { teamRunId: runtimeState.teamRunId, memberName: member.name, role: member.name === spec.leadAgentId ? "lead" : "member", }) + runtimeState = await transitionRuntimeState(runtimeState.teamRunId, (currentState) => ({ + ...currentState, + members: currentState.members.map((currentMember, currentIndex) => currentIndex === memberIndex + ? { ...currentMember, sessionId, status: "running" } + : currentMember), + }), config) }, }) resource.taskId = task.id diff --git a/src/shared/agent-tool-restrictions.ts b/src/shared/agent-tool-restrictions.ts index d84a2dfbe..21e481c5c 100644 --- a/src/shared/agent-tool-restrictions.ts +++ b/src/shared/agent-tool-restrictions.ts @@ -59,14 +59,18 @@ const AGENT_RESTRICTIONS: Record> = { }, } -export function getAgentToolRestrictions(agentName: string): Record { +type AgentToolRestrictionsOptions = { + includeTeamToolDenylist?: boolean +} + +export function getAgentToolRestrictions(agentName: string, options: AgentToolRestrictionsOptions = {}): Record { const stripped = stripInvisibleAgentCharacters(agentName) const agentRestrictions = AGENT_RESTRICTIONS[stripped] ?? Object.entries(AGENT_RESTRICTIONS).find(([key]) => key.toLowerCase() === stripped.toLowerCase())?.[1] ?? {} return { - ...TEAM_TOOL_DENYLIST, + ...(options.includeTeamToolDenylist === false ? {} : TEAM_TOOL_DENYLIST), ...agentRestrictions, } }