From 1827d8470997e379daf1459d543f88b224b1cec2 Mon Sep 17 00:00:00 2001 From: Bogdan Date: Tue, 30 Jun 2026 05:23:21 +0200 Subject: [PATCH] feat(agents): add per-agent step limits (#1815) * feat(agents): add per-agent step limits Add maxSteps agent configuration for markdown, JSON, plugin, and SDK agent definitions. Enforce the limit in subagent query execution by blocking over-limit tool calls, preserving a no-tool summary turn, and recording an agent_step_limit terminal reason. Add focused coverage for default behavior, invalid values, multi-turn accumulation, plugin parsing, failure-loop interaction, and summary-tool blocking. * test(agents): isolate agent loader fixtures * test(agents): stabilize agent loader config fixtures * fix(agents): harden step-limit summaries * fix(sdk): harden agent injection follow-up * fix(sdk): report invalid agent step limits --- README.md | 14 + src/components/ContextVisualization.tsx | 6 +- src/components/agents/AgentsList.tsx | 2 +- src/components/agents/AgentsMenu.test.tsx | 115 ++++ src/components/agents/AgentsMenu.tsx | 5 +- src/components/agents/agentFileUtils.ts | 6 + src/components/agents/types.ts | 2 +- src/components/agents/utils.ts | 5 +- src/entrypoints/sdk.d.ts | 23 +- src/entrypoints/sdk/agentDefinitions.test.ts | 196 ++++++ src/entrypoints/sdk/agentDefinitions.ts | 118 ++++ src/entrypoints/sdk/coreSchemas.ts | 14 +- src/entrypoints/sdk/coreTypes.generated.ts | 1 + src/entrypoints/sdk/query.ts | 289 ++++---- src/entrypoints/sdk/v2.ts | 90 ++- src/query.ts | 245 ++++++- src/query/agentStepLimit.test.ts | 690 +++++++++++++++++++ src/query/agentStepLimit.ts | 2 + src/query/toolFailureLoopGuard.test.ts | 37 +- src/query/toolFailureLoopGuard.ts | 20 +- src/query/transitions.ts | 6 + src/services/api/claude.ts | 6 +- src/tools/AgentTool/agentDisplay.ts | 3 +- src/tools/AgentTool/agentToolUtils.ts | 21 +- src/tools/AgentTool/loadAgentsDir.test.ts | 137 +++- src/tools/AgentTool/loadAgentsDir.ts | 27 +- src/tools/AgentTool/runAgent.ts | 118 ++-- src/types/message.ts | 2 + src/utils/analyzeContext.ts | 2 +- src/utils/frontmatterParser.ts | 9 +- src/utils/messages.ts | 3 + src/utils/plugins/loadPluginAgents.test.ts | 143 ++++ src/utils/plugins/loadPluginAgents.ts | 10 + src/utils/settings/constants.ts | 4 +- tests/sdk/package-consumer-types.test.ts | 21 + tests/sdk/query-happy-path.test.ts | 195 +++++- tests/sdk/sdk-v2-lifecycle.test.ts | 379 +++++++++- 37 files changed, 2723 insertions(+), 243 deletions(-) create mode 100644 src/entrypoints/sdk/agentDefinitions.test.ts create mode 100644 src/entrypoints/sdk/agentDefinitions.ts create mode 100644 src/query/agentStepLimit.test.ts create mode 100644 src/query/agentStepLimit.ts create mode 100644 src/utils/plugins/loadPluginAgents.test.ts diff --git a/README.md b/README.md index 3bc01fee7..00754d172 100644 --- a/README.md +++ b/README.md @@ -278,6 +278,20 @@ The `is_async` field reported in the `tengu_agent_tool_selected` event and the a For best results, use models with strong tool/function calling support. +### Agent step limits + +Custom agents can define `maxSteps` as a positive integer to cap how many tool-use steps a sub-agent may execute. When the limit is reached, OpenClaude stops additional tool calls and asks the sub-agent for a concise final summary covering completed work, findings, remaining tasks, and whether another run is needed. Omitting `maxSteps`, or setting it to an invalid value such as `0` or malformed input, preserves the default unlimited behavior. + +```markdown +--- +name: bounded-researcher +description: Use for focused research with bounded tool use +maxSteps: 8 +--- + +You are a focused research agent. +``` + ## Agent Routing OpenClaude can route different agents to different models through settings-based routing. This is useful for cost optimization or splitting work by model strength. diff --git a/src/components/ContextVisualization.tsx b/src/components/ContextVisualization.tsx index e3ef5f62c..ee50805c4 100644 --- a/src/components/ContextVisualization.tsx +++ b/src/components/ContextVisualization.tsx @@ -70,12 +70,12 @@ function CollapseStatus() { return null; } -// Order for displaying source groups: Project > User > Managed > Plugin > Built-in -const SOURCE_DISPLAY_ORDER = ['Project', 'User', 'Managed', 'Plugin', 'Built-in']; +// Order for displaying source groups: Project > User > Managed > Plugin > SDK > Built-in +const SOURCE_DISPLAY_ORDER = ['Project', 'User', 'Managed', 'Plugin', 'SDK', 'Built-in']; /** Group items by source type for display, sorted by tokens descending within each group */ function groupBySource(items: T[]): Map { const groups = new Map(); diff --git a/src/components/agents/AgentsList.tsx b/src/components/agents/AgentsList.tsx index e46815140..33f694068 100644 --- a/src/components/agents/AgentsList.tsx +++ b/src/components/agents/AgentsList.tsx @@ -13,7 +13,7 @@ import { Dialog } from '../design-system/Dialog.js'; import { Divider } from '../design-system/Divider.js'; import { getAgentSourceDisplayName } from './utils.js'; type Props = { - source: SettingSource | 'all' | 'built-in' | 'plugin'; + source: SettingSource | 'all' | 'built-in' | 'plugin' | 'sdk'; agents: ResolvedAgent[]; onBack: () => void; onSelect: (agent: AgentDefinition) => void; diff --git a/src/components/agents/AgentsMenu.test.tsx b/src/components/agents/AgentsMenu.test.tsx index 41a7b247e..e984ba706 100644 --- a/src/components/agents/AgentsMenu.test.tsx +++ b/src/components/agents/AgentsMenu.test.tsx @@ -379,3 +379,118 @@ test('omits set-active action when no session setter is available', async () => stdout.end() } }) + +test('shows SDK agents in all view without file edit actions', async () => { + const AgentsMenu = await importAgentsMenu() + const reviewer = createAgent('reviewer') + const sdkHelper = createAgent('sdk-helper', 'sdk') + const initialState = { + ...getDefaultAppState(), + agentDefinitions: { + activeAgents: [reviewer, sdkHelper], + allAgents: [reviewer, sdkHelper], + }, + } + const listStreams = createTestStreams() + const listRoot = await createRoot({ + stdout: listStreams.stdout as unknown as NodeJS.WriteStream, + stdin: listStreams.stdin as unknown as NodeJS.ReadStream, + patchConsole: false, + }) + + listRoot.render( + + {}} + initialModeState={{ mode: 'list-agents', source: 'all' }} + /> + , + ) + + try { + const listOutput = await waitForOutput( + listStreams.getOutput, + frame => frame.includes('reviewer') && frame.includes('sdk-helper'), + ) + expect(listOutput).toContain('sdk-helper') + } finally { + listRoot.unmount() + listStreams.stdin.end() + listStreams.stdout.end() + } + + const menuStreams = createTestStreams() + const menuRoot = await createRoot({ + stdout: menuStreams.stdout as unknown as NodeJS.WriteStream, + stdin: menuStreams.stdin as unknown as NodeJS.ReadStream, + patchConsole: false, + }) + + menuRoot.render( + + {}} + initialModeState={{ + mode: 'agent-menu', + agent: sdkHelper, + previousMode: { mode: 'list-agents', source: 'all' }, + }} + /> + , + ) + + try { + const menuOutput = await waitForOutput(menuStreams.getOutput, frame => + frame.includes('View agent'), + ) + expect(menuOutput).not.toContain('Edit agent') + expect(menuOutput).not.toContain('Delete agent') + } finally { + menuRoot.unmount() + menuStreams.stdin.end() + menuStreams.stdout.end() + } +}) + +test('shows only SDK agents in the dedicated SDK source list', async () => { + const AgentsMenu = await importAgentsMenu() + const reviewer = createAgent('reviewer') + const sdkHelper = createAgent('sdk-helper', 'sdk') + const initialState = { + ...getDefaultAppState(), + agentDefinitions: { + activeAgents: [reviewer, sdkHelper], + allAgents: [reviewer, sdkHelper], + }, + } + const { stdout, stdin, getOutput } = createTestStreams() + const root = await createRoot({ + stdout: stdout as unknown as NodeJS.WriteStream, + stdin: stdin as unknown as NodeJS.ReadStream, + patchConsole: false, + }) + + root.render( + + {}} + initialModeState={{ mode: 'list-agents', source: 'sdk' }} + /> + , + ) + + try { + const output = await waitForOutput(getOutput, frame => + frame.includes('sdk-helper'), + ) + expect(output).toContain('sdk-helper') + expect(output).not.toContain('reviewer') + } finally { + root.unmount() + stdin.end() + stdout.end() + } +}) diff --git a/src/components/agents/AgentsMenu.tsx b/src/components/agents/AgentsMenu.tsx index 39fba903b..40fc9ebbd 100644 --- a/src/components/agents/AgentsMenu.tsx +++ b/src/components/agents/AgentsMenu.tsx @@ -135,6 +135,7 @@ export function AgentsMenu(t0) { localSettings: t7, flagSettings: t8, plugin: t9, + sdk: allAgents.filter(a => a.source === "sdk"), all: allAgents }; $[16] = allAgents; @@ -202,7 +203,7 @@ export function AgentsMenu(t0) { { let t13; if ($[28] !== agentsBySource || $[29] !== modeState.source) { - t13 = modeState.source === "all" ? [...agentsBySource["built-in"], ...agentsBySource.userSettings, ...agentsBySource.projectSettings, ...agentsBySource.localSettings, ...agentsBySource.policySettings, ...agentsBySource.flagSettings, ...agentsBySource.plugin] : agentsBySource[modeState.source]; + t13 = modeState.source === "all" ? [...agentsBySource["built-in"], ...agentsBySource.userSettings, ...agentsBySource.projectSettings, ...agentsBySource.localSettings, ...agentsBySource.policySettings, ...agentsBySource.flagSettings, ...agentsBySource.plugin, ...agentsBySource.sdk] : agentsBySource[modeState.source]; $[28] = agentsBySource; $[29] = modeState.source; $[30] = t13; @@ -320,7 +321,7 @@ export function AgentsMenu(t0) { } const freshAgent_1 = t13; const agentToUse = freshAgent_1 || modeState.agent; - const isEditable = agentToUse.source !== "built-in" && agentToUse.source !== "plugin" && agentToUse.source !== "flagSettings"; + const isEditable = agentToUse.source !== "built-in" && agentToUse.source !== "plugin" && agentToUse.source !== "flagSettings" && agentToUse.source !== "sdk"; const isActiveAgent = agentToUse.agentType === activeAgentName; const sessionAgentToUse = agents.find(a_10 => a_10.agentType === agentToUse.agentType) ?? agentToUse; const editableItems = isEditable ? [{ diff --git a/src/components/agents/agentFileUtils.ts b/src/components/agents/agentFileUtils.ts index 0287b2944..7c9571ec2 100644 --- a/src/components/agents/agentFileUtils.ts +++ b/src/components/agents/agentFileUtils.ts @@ -110,6 +110,9 @@ export function getActualAgentFilePath(agent: AgentDefinition): string { if (agent.source === 'plugin') { throw new Error('Cannot get file path for plugin agents') } + if (agent.source === 'sdk') { + throw new Error('Cannot get file path for SDK agents') + } const dirPath = agent.baseDir || getAgentDirectoryPath(agent.source) const filename = agent.filename || agent.agentType @@ -144,6 +147,9 @@ export function getActualRelativeAgentFilePath(agent: AgentDefinition): string { if (agent.source === 'flagSettings') { return 'CLI argument' } + if (agent.source === 'sdk') { + return 'SDK agent' + } const dirPath = agent.baseDir && diff --git a/src/components/agents/types.ts b/src/components/agents/types.ts index 1a47a6e35..8a2daec67 100644 --- a/src/components/agents/types.ts +++ b/src/components/agents/types.ts @@ -13,7 +13,7 @@ type WithAgent = { agent: AgentDefinition } // Simplified state type using intersection types export type ModeState = | { mode: 'main-menu' } - | { mode: 'list-agents'; source: SettingSource | 'all' | 'built-in' } + | { mode: 'list-agents'; source: SettingSource | 'all' | 'built-in' | 'plugin' | 'sdk' } | ({ mode: 'agent-menu' } & WithAgent & WithPreviousMode) | ({ mode: 'view-agent' } & WithAgent & WithPreviousMode) | { mode: 'create-agent' } diff --git a/src/components/agents/utils.ts b/src/components/agents/utils.ts index 3497e68a8..0de332df2 100644 --- a/src/components/agents/utils.ts +++ b/src/components/agents/utils.ts @@ -3,7 +3,7 @@ import type { SettingSource } from 'src/utils/settings/constants.js' import { getSettingSourceName } from 'src/utils/settings/constants.js' export function getAgentSourceDisplayName( - source: SettingSource | 'all' | 'built-in' | 'plugin', + source: SettingSource | 'all' | 'built-in' | 'plugin' | 'sdk', ): string { if (source === 'all') { return 'Agents' @@ -14,5 +14,8 @@ export function getAgentSourceDisplayName( if (source === 'plugin') { return 'Plugin agents' } + if (source === 'sdk') { + return 'SDK agents' + } return capitalize(getSettingSourceName(source)) } diff --git a/src/entrypoints/sdk.d.ts b/src/entrypoints/sdk.d.ts index 9586611f1..5dfcbceb8 100644 --- a/src/entrypoints/sdk.d.ts +++ b/src/entrypoints/sdk.d.ts @@ -288,12 +288,18 @@ export type QueryOptions = { | { type: 'custom'; content: string } /** Agent definitions to register with the query engine. */ agents?: Record settingSources?: string[] /** When true, yields stream_event messages for token-by-token streaming. */ @@ -403,6 +409,21 @@ export type SDKSessionOptions = { onPermissionRequest?: (message: SDKPermissionRequestMessage) => void /** Tools to disallow (blanket deny by tool name). */ disallowedTools?: string[] + /** Agent definitions to register with the session engine. */ + agents?: Record } export interface SDKSession { diff --git a/src/entrypoints/sdk/agentDefinitions.test.ts b/src/entrypoints/sdk/agentDefinitions.test.ts new file mode 100644 index 000000000..4bea0eb3d --- /dev/null +++ b/src/entrypoints/sdk/agentDefinitions.test.ts @@ -0,0 +1,196 @@ +import { describe, expect, test } from 'bun:test' + +import { + buildSdkUserAgents, + mergeSdkUserAgents, +} from './agentDefinitions.js' + +describe('buildSdkUserAgents', () => { + test('preserves valid maxSteps, reports invalid limits, and rejects malformed SDK agents safely', () => { + const failures: Array<{ name: string; error: string }> = [] + + const agents = buildSdkUserAgents( + { + valid: { + description: 'Use for valid SDK agent coverage', + prompt: 'valid prompt', + maxTurns: 4, + maxSteps: 2, + }, + zero: { + description: 'Use for invalid SDK agent coverage', + prompt: 'zero prompt', + maxSteps: 0, + }, + malformed: { + description: 'Use for malformed SDK agent coverage', + prompt: 'malformed prompt', + maxSteps: '2' as unknown as number, + }, + negative: { + description: 'Use for negative SDK agent coverage', + prompt: 'negative prompt', + maxSteps: -1, + }, + fractional: { + description: 'Use for fractional SDK agent coverage', + prompt: 'fractional prompt', + maxSteps: 1.5, + }, + invalidTurns: { + description: 'Use for invalid maxTurns SDK coverage', + prompt: 'invalid turns prompt', + maxTurns: 0, + }, + malformedTurns: { + description: 'Use for malformed maxTurns SDK coverage', + prompt: 'malformed turns prompt', + maxTurns: '2' as unknown as number, + }, + missingDescription: { + prompt: 'missing description prompt', + maxSteps: 3, + } as unknown as { + description: string + prompt: string + maxSteps: number + }, + missingPrompt: { + maxSteps: 3, + } as unknown as { + description: string + prompt: string + maxSteps: number + }, + broken: { + description: 'Use for broken SDK agent coverage', + prompt: 2 as unknown as string, + }, + scalar: 'not an object', + array: [], + nullish: null, + }, + (name, error) => failures.push({ name, error }), + ) + + expect(agents.map(agent => agent.agentType)).toEqual([ + 'valid', + 'missingDescription', + ]) + expect(agents.find(agent => agent.agentType === 'valid')?.maxSteps).toBe(2) + expect(agents.find(agent => agent.agentType === 'valid')?.maxTurns).toBe(4) + expect( + agents.find(agent => agent.agentType === 'missingDescription'), + ).toMatchObject({ + whenToUse: 'missingDescription', + maxSteps: 3, + }) + expect( + agents + .find(agent => agent.agentType === 'missingDescription') + ?.getSystemPrompt(), + ).toBe('missing description prompt') + expect(failures.map(failure => failure.name)).toEqual([ + 'zero', + 'malformed', + 'negative', + 'fractional', + 'invalidTurns', + 'malformedTurns', + 'missingPrompt', + 'broken', + 'scalar', + 'array', + 'nullish', + ]) + for (const name of ['zero', 'malformed', 'negative', 'fractional']) { + expect(failures.find(failure => failure.name === name)?.error).toContain( + 'maxSteps', + ) + } + for (const name of ['invalidTurns', 'malformedTurns']) { + expect(failures.find(failure => failure.name === name)?.error).toContain( + 'maxTurns', + ) + } + expect( + failures.find(failure => failure.name === 'missingPrompt')?.error, + ).toContain('prompt') + expect(failures.find(failure => failure.name === 'broken')?.error).toContain( + 'prompt', + ) + expect( + failures + .filter(failure => ['scalar', 'array', 'nullish'].includes(failure.name)) + .map(failure => failure.error), + ).toEqual([ + 'Agent definition must be an object', + 'Agent definition must be an object', + 'Agent definition must be an object', + ]) + }) + + test('mergeSdkUserAgents gives SDK agents deterministic active-agent precedence', () => { + const existingHelper = { + agentType: 'helper', + whenToUse: 'Use filesystem helper', + getSystemPrompt: () => 'filesystem helper prompt', + source: 'projectSettings', + } + const filesystemOnly = { + agentType: 'filesystem', + whenToUse: 'Use filesystem-only agent', + getSystemPrompt: () => 'filesystem-only prompt', + source: 'projectSettings', + } + const sdkHelper = { + agentType: 'helper', + whenToUse: 'Use SDK helper', + getSystemPrompt: () => 'sdk helper prompt', + source: 'sdk' as const, + maxSteps: 2, + } + + const merged = mergeSdkUserAgents( + { + activeAgents: [filesystemOnly, existingHelper], + allAgents: [filesystemOnly, existingHelper], + }, + [sdkHelper], + ) + + expect(merged.activeAgents).toEqual([filesystemOnly, sdkHelper]) + expect(merged.allAgents).toEqual([ + filesystemOnly, + existingHelper, + sdkHelper, + ]) + }) + + test('mergeSdkUserAgents does not let SDK agents override managed policy agents', () => { + const managedHelper = { + agentType: 'helper', + whenToUse: 'Use managed helper', + getSystemPrompt: () => 'managed helper prompt', + source: 'policySettings', + } + const sdkHelper = { + agentType: 'helper', + whenToUse: 'Use SDK helper', + getSystemPrompt: () => 'sdk helper prompt', + source: 'sdk' as const, + maxSteps: 2, + } + + const merged = mergeSdkUserAgents( + { + activeAgents: [managedHelper], + allAgents: [managedHelper], + }, + [sdkHelper], + ) + + expect(merged.activeAgents).toEqual([managedHelper]) + expect(merged.allAgents).toEqual([managedHelper, sdkHelper]) + }) +}) diff --git a/src/entrypoints/sdk/agentDefinitions.ts b/src/entrypoints/sdk/agentDefinitions.ts new file mode 100644 index 000000000..a5da6e440 --- /dev/null +++ b/src/entrypoints/sdk/agentDefinitions.ts @@ -0,0 +1,118 @@ +import { AgentDefinitionSchema } from './coreSchemas.js' + +export type SdkAgentDefinitionInput = { + description?: string + prompt: string + /** + * Tool allowlist for this agent. If omitted or set to ['*'], the agent can use + * all tools available to subagents after disallowedTools is applied. + */ + tools?: string[] + /** + * Tool denylist for this agent. Deny entries always override tools entries. + */ + disallowedTools?: string[] + model?: string + maxTurns?: number + maxSteps?: number +} + +export type SdkInjectedAgentDefinition = { + agentType: string + whenToUse: string + getSystemPrompt: () => string + source: 'sdk' + tools?: string[] + disallowedTools?: string[] + model?: string + maxTurns?: number + maxSteps?: number +} + +export type SdkMergeableAgentDefinition = { + agentType: string + source: string +} + +export type SdkAgentDefinitionSet< + TAgent extends SdkMergeableAgentDefinition = SdkMergeableAgentDefinition, +> = { + activeAgents: TAgent[] + allAgents: TAgent[] +} + +export function buildSdkUserAgents( + userAgents: Record | undefined, + reportInvalidAgent: (name: string, errorMessage: string) => void, +): SdkInjectedAgentDefinition[] { + if (!userAgents || Object.keys(userAgents).length === 0) { + return [] + } + + return Object.entries(userAgents).flatMap(([name, def]) => { + if (def === null || typeof def !== 'object' || Array.isArray(def)) { + reportInvalidAgent(name, 'Agent definition must be an object') + return [] + } + + const candidate = def as Partial + const normalizedDef = { + ...candidate, + description: candidate.description ?? name, + } + + const parsed = AgentDefinitionSchema().safeParse(normalizedDef) + if (!parsed.success) { + reportInvalidAgent(name, parsed.error.message) + return [] + } + + const data = parsed.data + return [ + { + agentType: name, + whenToUse: data.description, + getSystemPrompt: () => data.prompt, + source: 'sdk', + ...(data.tools ? { tools: data.tools } : {}), + ...(data.disallowedTools + ? { disallowedTools: data.disallowedTools } + : {}), + ...(data.model ? { model: data.model } : {}), + ...(data.maxTurns !== undefined ? { maxTurns: data.maxTurns } : {}), + ...(data.maxSteps !== undefined ? { maxSteps: data.maxSteps } : {}), + }, + ] + }) +} + +export function mergeSdkUserAgents( + agentDefs: SdkAgentDefinitionSet, + userAgents: SdkInjectedAgentDefinition[], +): SdkAgentDefinitionSet { + if (userAgents.length === 0) { + return agentDefs + } + + const protectedAgentTypes = new Set( + agentDefs.activeAgents + .filter(agent => agent.source === 'policySettings') + .map(agent => agent.agentType), + ) + const activeUserAgents = userAgents.filter( + agent => !protectedAgentTypes.has(agent.agentType), + ) + const activeUserAgentTypes = new Set( + activeUserAgents.map(agent => agent.agentType), + ) + return { + ...agentDefs, + activeAgents: [ + ...agentDefs.activeAgents.filter( + agent => !activeUserAgentTypes.has(agent.agentType), + ), + ...activeUserAgents, + ], + allAgents: [...agentDefs.allAgents, ...userAgents], + } +} diff --git a/src/entrypoints/sdk/coreSchemas.ts b/src/entrypoints/sdk/coreSchemas.ts index 0a6aa3a8a..d40401b75 100644 --- a/src/entrypoints/sdk/coreSchemas.ts +++ b/src/entrypoints/sdk/coreSchemas.ts @@ -1141,12 +1141,14 @@ export const AgentDefinitionSchema = lazySchema(() => .array(z.string()) .optional() .describe( - 'Array of allowed tool names. If omitted, inherits all tools from parent', + 'Array of allowed tool names. If omitted or set to ["*"], inherits all tools from parent before disallowedTools is applied', ), disallowedTools: z .array(z.string()) .optional() - .describe('Array of tool names to explicitly disallow for this agent'), + .describe( + 'Array of tool names to explicitly disallow for this agent. Deny entries always override tools entries', + ), prompt: z.string().describe("The agent's system prompt"), model: z .string() @@ -1177,6 +1179,14 @@ export const AgentDefinitionSchema = lazySchema(() => .describe( 'Maximum number of agentic turns (API round-trips) before stopping', ), + maxSteps: z + .number() + .int() + .positive() + .optional() + .describe( + 'Maximum number of subagent tool-use steps before forcing a concise final summary', + ), background: z .boolean() .optional() diff --git a/src/entrypoints/sdk/coreTypes.generated.ts b/src/entrypoints/sdk/coreTypes.generated.ts index b24ec1c6d..567f3fa91 100644 --- a/src/entrypoints/sdk/coreTypes.generated.ts +++ b/src/entrypoints/sdk/coreTypes.generated.ts @@ -1532,6 +1532,7 @@ export type AgentDefinition = { skills?: string[] initialPrompt?: string maxTurns?: number + maxSteps?: number background?: boolean memory?: "user" | "project" | "local" effort?: "low" | "medium" | "high" | "xhigh" | "max" | number diff --git a/src/entrypoints/sdk/query.ts b/src/entrypoints/sdk/query.ts index 02e4277ad..0d5d894ed 100644 --- a/src/entrypoints/sdk/query.ts +++ b/src/entrypoints/sdk/query.ts @@ -18,7 +18,7 @@ import { getEmptyToolPermissionContext, type ToolPermissionContext, } from '../../Tool.js' -import { getTools } from '../../tools.js' +import { assembleToolPool, getTools } from '../../tools.js' import { createFileStateCacheWithSizeLimit } from '../../utils/fileStateCache.js' import { init } from '../init.js' import { @@ -35,7 +35,10 @@ import { runWithSdkContext, } from '../../bootstrap/state.js' import type { SessionId } from '../../types/ids.js' -import { getAgentDefinitionsWithOverrides } from '../../tools/AgentTool/loadAgentsDir.js' +import { + getAgentDefinitionsWithOverrides, + type AgentDefinitionsResult, +} from '../../tools/AgentTool/loadAgentsDir.js' import type { RewindFilesResult, McpServerStatus, @@ -81,6 +84,11 @@ import { buildConversationChain, stripExtraFields, } from './transcript.js' +import { + buildSdkUserAgents, + mergeSdkUserAgents, + type SdkAgentDefinitionInput, +} from './agentDefinitions.js' // ============================================================================ // QueryOptions type @@ -150,14 +158,7 @@ export type QueryOptions = { | { type: 'preset'; preset: string; append?: string } | { type: 'custom'; content: string } /** Agent definitions to register with the query engine. */ - agents?: Record + agents?: Record /** Setting sources to load. */ settingSources?: string[] /** When true, yields stream_event messages for token-by-token streaming. */ @@ -395,6 +396,7 @@ class QueryImpl implements Query { private resumeSessionAt?: string private userAgents?: QueryOptions['agents'] private mcpServers?: Record + private sdkMcpTools: Parameters[1] = [] private permissionContext: ToolPermissionContext private timeoutQueue: SDKPermissionTimeoutMessage[] = [] private agentFailureQueue: SDKAgentLoadFailureMessage[] = [] @@ -520,65 +522,74 @@ class QueryImpl implements Query { } // Load agent definitions BEFORE creating engine context - let agentDefs: { activeAgents: any[]; allAgents: any[] } = { activeAgents: [], allAgents: [] } - try { - agentDefs = await getAgentDefinitionsWithOverrides(self.cwd) - } catch (err) { - // Agent loading failed — continue without agents but emit failure event - const errorMessage = err instanceof Error ? err.message : String(err) - console.warn('SDK: agent definitions loading failed:', errorMessage) + let agentDefs: AgentDefinitionsResult = { activeAgents: [], allAgents: [] } + const reportDefinitionFailure = (errorMessage: string) => { + self.pushAgentFailure({ + type: 'agent_load_failure', + stage: 'definitions', + error_message: errorMessage, + }) + } + try { + agentDefs = await getAgentDefinitionsWithOverrides(self.cwd) + for (const failedFile of agentDefs.failedFiles ?? []) { + reportDefinitionFailure( + `Failed to load agent definition '${failedFile.path}': ${failedFile.error}`, + ) + } + } catch (err) { + // Agent loading failed — continue without agents but emit failure event + const errorMessage = err instanceof Error ? err.message : String(err) + console.warn('SDK: agent definitions loading failed:', errorMessage) + reportDefinitionFailure(errorMessage) + } + + // Inject agents into the engine + const userAgents = buildSdkUserAgents( + self.userAgents, + (name, errorMessage) => { self.pushAgentFailure({ type: 'agent_load_failure', - stage: 'definitions', + stage: 'injection', + error_message: `Invalid SDK agent '${name}': ${errorMessage}`, + }) + }, + ) + const mergedAgentDefs = mergeSdkUserAgents(agentDefs, userAgents) + if (mergedAgentDefs.activeAgents.length > 0) { + self.appStateStore.setState(prev => ({ + ...prev, + agentDefinitions: mergedAgentDefs, + })) + try { + self.engine.injectAgents(mergedAgentDefs.activeAgents) + } catch (err) { + // Agent injection failed — continue without agents but emit failure event + const errorMessage = err instanceof Error ? err.message : String(err) + console.warn('SDK: agent injection failed:', errorMessage) + self.appStateStore.setState(prev => ({ + ...prev, + agentDefinitions: { activeAgents: [], allAgents: [] }, + })) + self.pushAgentFailure({ + type: 'agent_load_failure', + stage: 'injection', error_message: errorMessage, }) } - - // Update AppState with agents + } else { self.appStateStore.setState(prev => ({ ...prev, - agentDefinitions: agentDefs, + agentDefinitions: mergedAgentDefs, })) + } - // Inject agents into the engine - if (self.userAgents && Object.keys(self.userAgents).length > 0) { - const userAgents: Array<{ - agentType: string - whenToUse: string - getSystemPrompt: () => string - tools?: string[] - disallowedTools?: string[] - model?: string - maxTurns?: number - }> = Object.entries(self.userAgents).map(([name, def]) => ({ - agentType: name, - whenToUse: def.description ?? name, - getSystemPrompt: () => def.prompt ?? '', - ...(def.tools ? { tools: def.tools } : {}), - ...(def.disallowedTools ? { disallowedTools: def.disallowedTools } : {}), - ...(def.model ? { model: def.model } : {}), - ...(def.maxTurns ? { maxTurns: def.maxTurns } : {}), - })) - agentDefs.activeAgents.push(...userAgents) - } - if (agentDefs.activeAgents.length > 0) { - try { - self.engine.injectAgents(agentDefs.activeAgents) - } catch (err) { - // Agent injection failed — continue without agents but emit failure event - const errorMessage = err instanceof Error ? err.message : String(err) - console.warn('SDK: agent injection failed:', errorMessage) - self.pushAgentFailure({ - type: 'agent_load_failure', - stage: 'injection', - error_message: errorMessage, - }) - } - } - + let envMutexAcquired = false + try { // Apply env overrides AFTER init() with full-duration mutex (SEC-1) if (hasEnvOverrides) { await acquireEnvMutex() + envMutexAcquired = true self.envSnapshot = {} for (const key of Object.keys(self.envOverrides!)) { self.envSnapshot[key] = process.env[key] @@ -592,103 +603,103 @@ class QueryImpl implements Query { } } - try { - // Connect MCP servers if provided - if (self.mcpServers && Object.keys(self.mcpServers).length > 0) { - try { - const { clients: mcpClients, tools: mcpTools } = await connectSdkMcpServers(self.mcpServers) - if (mcpClients.length > 0) { - self.engine.setMcpClients(mcpClients) - } - if (mcpTools.length > 0) { - const allTools = [...getTools(self.permissionContext)] // Mutable copy - for (const mcpTool of mcpTools) { - if (!allTools.some(t => t.name === mcpTool.name)) { - allTools.push(mcpTool) - } - } - self.engine.updateTools(allTools) - } - } catch (err) { - // MCP connection failed — continue without MCP tools - console.warn('SDK: MCP server connection failed:', err instanceof Error ? err.message : String(err)) + // Connect MCP servers if provided + if (self.mcpServers && Object.keys(self.mcpServers).length > 0) { + try { + const { clients: mcpClients, tools: mcpTools } = await connectSdkMcpServers(self.mcpServers) + self.sdkMcpTools = mcpTools + if (mcpClients.length > 0) { + self.engine.setMcpClients(mcpClients) } + if (mcpTools.length > 0) { + self.engine.updateTools(assembleToolPool(self.permissionContext, mcpTools)) + } + } catch (err) { + // MCP connection failed — continue without MCP tools + console.warn('SDK: MCP server connection failed:', err instanceof Error ? err.message : String(err)) } + } - // Handle continue/fork/resume session resolution - let effectiveSessionId: string | undefined = self._sessionId - let resolvedTranscriptDir: string | null = null + // Handle continue/fork/resume session resolution + let effectiveSessionId: string | undefined = self._sessionId + let resolvedTranscriptDir: string | null = null - if (self.continueSession && !self._sessionIdExplicitlyProvided) { - const sessions = await listSessions({ dir: self.cwd, limit: 1 }) - if (sessions.length > 0) { - effectiveSessionId = sessions[0].sessionId - const result = await loadAndInjectSessionMessages(effectiveSessionId, self.cwd, self.engine, self.resumeSessionAt) - if (result.loaded) { - resolvedTranscriptDir = result.transcriptDir - } else { - effectiveSessionId = undefined - } - } else { - // No existing sessions — keep the constructor-created UUID for fresh query - effectiveSessionId = self._sessionId - } - } else if (self.shouldFork && self._sessionId) { - try { - const forkResult = await forkSession(self._sessionId, { dir: self.cwd }) - effectiveSessionId = forkResult.sessionId - const result = await loadAndInjectSessionMessages(effectiveSessionId, self.cwd, self.engine, self.resumeSessionAt) - if (result.loaded) { - resolvedTranscriptDir = result.transcriptDir - } else { - effectiveSessionId = undefined - } - } catch { - effectiveSessionId = undefined - } - } else if (self._sessionId) { - const result = await loadAndInjectSessionMessages(self._sessionId, self.cwd, self.engine, self.resumeSessionAt) + if (self.continueSession && !self._sessionIdExplicitlyProvided) { + const sessions = await listSessions({ dir: self.cwd, limit: 1 }) + if (sessions.length > 0) { + effectiveSessionId = sessions[0].sessionId + const result = await loadAndInjectSessionMessages(effectiveSessionId, self.cwd, self.engine, self.resumeSessionAt) if (result.loaded) { resolvedTranscriptDir = result.transcriptDir } else { - // Session file not found — preserve constructor UUID for fresh session - effectiveSessionId = self._sessionId + effectiveSessionId = undefined } + } else { + // No existing sessions — keep the constructor-created UUID for fresh query + effectiveSessionId = self._sessionId } - - // Switch session for transcript writes using the resolved transcript dir - if (!effectiveSessionId) { - regenerateSessionId() - effectiveSessionId = getSessionId() + } else if (self.shouldFork && self._sessionId) { + try { + const forkResult = await forkSession(self._sessionId, { + dir: self.cwd, + ...(self.resumeSessionAt + ? { upToMessageId: self.resumeSessionAt } + : {}), + }) + effectiveSessionId = forkResult.sessionId + const result = await loadAndInjectSessionMessages(effectiveSessionId, self.cwd, self.engine) + if (result.loaded) { + resolvedTranscriptDir = result.transcriptDir + } else { + effectiveSessionId = undefined + } + } catch { + effectiveSessionId = undefined } - switchSession(effectiveSessionId as SessionId, resolvedTranscriptDir) + } else if (self._sessionId) { + const result = await loadAndInjectSessionMessages(self._sessionId, self.cwd, self.engine, self.resumeSessionAt) + if (result.loaded) { + resolvedTranscriptDir = result.transcriptDir + } else { + // Session file not found — preserve constructor UUID for fresh session + effectiveSessionId = self._sessionId + } + } - // Sync resolved sessionId and transcript dir back to authoritative fields - self._sessionId = effectiveSessionId - sdkContext.sessionId = effectiveSessionId as SessionId - sdkContext.sessionProjectDir = resolvedTranscriptDir + // Switch session for transcript writes using the resolved transcript dir + if (!effectiveSessionId) { + regenerateSessionId() + effectiveSessionId = getSessionId() + } + switchSession(effectiveSessionId as SessionId, resolvedTranscriptDir) - // Submit to engine - if (typeof self.prompt === 'string') { - for await (const engineMsg of self.engine.submitMessage(self.prompt)) { + // Sync resolved sessionId and transcript dir back to authoritative fields + self._sessionId = effectiveSessionId + sdkContext.sessionId = effectiveSessionId as SessionId + sdkContext.sessionProjectDir = resolvedTranscriptDir + yield* self.drainAgentFailureQueue() + + // Submit to engine + if (typeof self.prompt === 'string') { + for await (const engineMsg of self.engine.submitMessage(self.prompt)) { + yield engineMsg + yield* self.drainTimeoutQueue() + yield* self.drainAgentFailureQueue() + } + } else { + for await (const userMessage of self.prompt) { + if (self.abortController.signal.aborted) break + const content = extractPromptFromUserMessage(userMessage) + for await (const engineMsg of self.engine.submitMessage(content, { uuid: userMessage.uuid })) { yield engineMsg yield* self.drainTimeoutQueue() yield* self.drainAgentFailureQueue() } - } else { - for await (const userMessage of self.prompt) { - if (self.abortController.signal.aborted) break - const content = extractPromptFromUserMessage(userMessage) - for await (const engineMsg of self.engine.submitMessage(content, { uuid: userMessage.uuid })) { - yield engineMsg - yield* self.drainTimeoutQueue() - yield* self.drainAgentFailureQueue() - } - } } - // Final drain for timeout/failure messages that fired on the last engine yield - yield* self.drainTimeoutQueue() - yield* self.drainAgentFailureQueue() + } + // Final drain for timeout/failure messages that fired on the last engine yield + yield* self.drainTimeoutQueue() + yield* self.drainAgentFailureQueue() } finally { // Clean up timeout and agent failure queues self.timeoutQueue.length = 0 @@ -705,7 +716,7 @@ class QueryImpl implements Query { } self.envSnapshot = undefined } - if (hasEnvOverrides) { + if (envMutexAcquired) { releaseEnvMutex() } } @@ -740,7 +751,7 @@ class QueryImpl implements Query { toolPermissionContext: newPermissionContext, })) // Refresh the engine's tool list to reflect new permissions - const updatedTools = getTools(newPermissionContext) + const updatedTools = assembleToolPool(newPermissionContext, this.sdkMcpTools) this.engine.updateTools(updatedTools) } diff --git a/src/entrypoints/sdk/v2.ts b/src/entrypoints/sdk/v2.ts index 84f4e9a41..621e693e7 100644 --- a/src/entrypoints/sdk/v2.ts +++ b/src/entrypoints/sdk/v2.ts @@ -17,7 +17,7 @@ import { createStore, type Store } from '../../state/store.js' import { type ToolPermissionContext, } from '../../Tool.js' -import { getTools } from '../../tools.js' +import { assembleToolPool, getTools } from '../../tools.js' import { createFileStateCacheWithSizeLimit } from '../../utils/fileStateCache.js' import { init } from '../init.js' import { @@ -32,7 +32,10 @@ import { runWithSdkContext, } from '../../bootstrap/state.js' import type { SessionId } from '../../types/ids.js' -import { getAgentDefinitionsWithOverrides } from '../../tools/AgentTool/loadAgentsDir.js' +import { + getAgentDefinitionsWithOverrides, + type AgentDefinitionsResult, +} from '../../tools/AgentTool/loadAgentsDir.js' import type { PermissionResult, SDKResultMessage as GeneratedSDKResultMessage, @@ -66,6 +69,11 @@ import { buildConversationChain as buildChain, stripExtraFields as stripChainFields, } from './transcript.js' +import { + buildSdkUserAgents, + mergeSdkUserAgents, + type SdkAgentDefinitionInput, +} from './agentDefinitions.js' // ============================================================================ // V2 API Types @@ -104,6 +112,8 @@ export type SDKSessionOptions = { onPermissionRequest?: (message: SDKPermissionRequestMessage) => void /** Tools to disallow (blanket deny by tool name). */ disallowedTools?: string[] + /** Agent definitions to register with the session engine. */ + agents?: Record } /** @@ -266,25 +276,66 @@ class SDKSessionImpl implements SDKSession { // Load agent definitions once (not on every sendMessage call) if (!self.agentsLoaded) { - try { - const agentDefs = await getAgentDefinitionsWithOverrides(self.options.cwd) - self.appStateStore.setState(prev => ({ - ...prev, - agentDefinitions: agentDefs, - })) - if (agentDefs.activeAgents.length > 0) { - self.engine.injectAgents(agentDefs.activeAgents) - } - } catch (err) { - // Agent loading failed — continue without agents but emit failure event - const errorMessage = err instanceof Error ? err.message : String(err) - console.warn('SDK: agent loading failed:', errorMessage) + let agentDefs: AgentDefinitionsResult = { activeAgents: [], allAgents: [] } + const reportDefinitionFailure = (errorMessage: string) => { self.pushAgentFailure({ type: 'agent_load_failure', stage: 'definitions', error_message: errorMessage, }) } + try { + agentDefs = await getAgentDefinitionsWithOverrides(self.options.cwd) + for (const failedFile of agentDefs.failedFiles ?? []) { + reportDefinitionFailure( + `Failed to load agent definition '${failedFile.path}': ${failedFile.error}`, + ) + } + } catch (err) { + // Agent loading failed — continue without filesystem agents but emit failure event + const errorMessage = err instanceof Error ? err.message : String(err) + console.warn('SDK: agent loading failed:', errorMessage) + reportDefinitionFailure(errorMessage) + } + + const userAgents = buildSdkUserAgents( + self.options.agents, + (name, errorMessage) => { + self.pushAgentFailure({ + type: 'agent_load_failure', + stage: 'injection', + error_message: `Invalid SDK agent '${name}': ${errorMessage}`, + }) + }, + ) + const mergedAgentDefs = mergeSdkUserAgents(agentDefs, userAgents) + if (mergedAgentDefs.activeAgents.length > 0) { + self.appStateStore.setState(prev => ({ + ...prev, + agentDefinitions: mergedAgentDefs, + })) + try { + self.engine.injectAgents(mergedAgentDefs.activeAgents) + } catch (err) { + // Agent injection failed — continue without agents but emit failure event + const errorMessage = err instanceof Error ? err.message : String(err) + console.warn('SDK: agent injection failed:', errorMessage) + self.appStateStore.setState(prev => ({ + ...prev, + agentDefinitions: { activeAgents: [], allAgents: [] }, + })) + self.pushAgentFailure({ + type: 'agent_load_failure', + stage: 'injection', + error_message: errorMessage, + }) + } + } else { + self.appStateStore.setState(prev => ({ + ...prev, + agentDefinitions: mergedAgentDefs, + })) + } self.agentsLoaded = true } @@ -297,13 +348,7 @@ class SDKSessionImpl implements SDKSession { } if (mcpTools.length > 0) { const permissionContext = self.appStateStore.getState().toolPermissionContext - const allTools = [...getTools(permissionContext)] // Mutable copy - for (const mcpTool of mcpTools) { - if (!allTools.some(t => t.name === mcpTool.name)) { - allTools.push(mcpTool) - } - } - self.engine.updateTools(allTools) + self.engine.updateTools(assembleToolPool(permissionContext, mcpTools)) } } catch (err) { // MCP connection failed — continue without MCP tools @@ -317,6 +362,7 @@ class SDKSessionImpl implements SDKSession { try { if (self._abortController?.signal.aborted) return + yield* self.drainAgentFailureQueue() for await (const engineMsg of self.engine.submitMessage(content)) { if (self._abortController?.signal.aborted) break yield engineMsg diff --git a/src/query.ts b/src/query.ts index c5483c73e..c9774dd64 100644 --- a/src/query.ts +++ b/src/query.ts @@ -49,6 +49,7 @@ import { } from './services/api/errors.js' import { logAntError, logForDebugging } from './utils/debug.js' import { + createAssistantMessage, createUserMessage, createUserInterruptionMessage, normalizeMessagesForAPI, @@ -112,6 +113,7 @@ import { createToolFailureLoopGuardState, updateToolFailureLoopGuard, } from './query/toolFailureLoopGuard.js' +import { AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX } from './query/agentStepLimit.js' import { buildQueryConfig } from './query/config.js' import { getGlobalConfig, @@ -180,6 +182,123 @@ function* yieldMissingToolResultBlocks( const MAX_OUTPUT_TOKENS_RECOVERY_LIMIT = 3 const MAX_CONTINUATION_NUDGES = 20 +type AgentStepLimitConfig = { + maxSteps: number + agentType?: string +} + +type AgentStepLimitState = AgentStepLimitConfig & { + stepsUsed: number + summaryRequested: boolean +} + +function normalizeAgentStepLimit( + limit: AgentStepLimitConfig | undefined, +): AgentStepLimitState | undefined { + if ( + !limit || + !Number.isInteger(limit.maxSteps) || + limit.maxSteps <= 0 + ) { + return undefined + } + return { + maxSteps: limit.maxSteps, + agentType: limit.agentType, + stepsUsed: 0, + summaryRequested: false, + } +} + +function findAssistantMessageForToolUse( + assistantMessages: AssistantMessage[], + toolUseId: string, +): AssistantMessage | undefined { + return assistantMessages.find(message => + message.message.content.some( + content => content.type === 'tool_use' && content.id === toolUseId, + ), + ) +} + +function createAgentStepLimitToolResult( + toolUse: ToolUseBlock, + assistantMessage: AssistantMessage | undefined, + limit: AgentStepLimitState, +): UserMessage { + const content = + `${AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX} for '${limit.agentType ?? 'subagent'}' ` + + `(${limit.stepsUsed}/${limit.maxSteps} tool uses). This tool call was not executed. ` + + 'Do not call more tools; provide the final summary requested next.' + + return createUserMessage({ + content: [ + { + type: 'tool_result', + content: `${content}`, + is_error: true, + tool_use_id: toolUse.id, + }, + ], + toolUseResult: content, + isAgentStepLimitToolResult: true, + ...(assistantMessage + ? { sourceToolAssistantUUID: assistantMessage.uuid } + : {}), + }) +} + +function createAgentStepLimitSummaryRequest( + limit: AgentStepLimitState, +): UserMessage { + return createUserMessage({ + content: + `Agent '${limit.agentType ?? 'subagent'}' reached its configured step limit ` + + `after ${limit.stepsUsed}/${limit.maxSteps} tool uses. Stop using tools now. ` + + 'Provide a concise final summary with these sections: completed work, findings, ' + + 'remaining tasks, and whether another run is needed.', + isMeta: true, + }) +} + +function createAgentStepLimitForcedSummary( + limit: AgentStepLimitState, + blockedToolUseCount: number, +): AssistantMessage { + const blockedCallText = + blockedToolUseCount === 1 + ? '1 additional tool call was blocked' + : `${blockedToolUseCount} additional tool calls were blocked` + + return createAssistantMessage({ + content: + `Completed work: Agent '${limit.agentType ?? 'subagent'}' reached its configured step limit ` + + `after ${limit.stepsUsed}/${limit.maxSteps} tool uses. ` + + `Findings: ${blockedCallText} during the forced summary step and no more tools were run. ` + + 'Remaining tasks: continue any unfinished work in another run if more tool access is needed. ' + + 'Another run needed: yes, if the requested task is not complete.', + }) +} + +function hasAssistantSummaryText( + assistantMessage: AssistantMessage | undefined, +): boolean { + const text = (assistantMessage?.message.content ?? []) + .map(part => + part.type === 'text' && typeof part.text === 'string' + ? part.text.toLowerCase() + : '', + ) + .join('\n') + + return ( + text.includes('completed') && + text.includes('findings') && + text.includes('remaining tasks') && + text.includes('another run') + ) +} + function formatAutoCompactRetryDelay(delayMs: number): string { const totalSeconds = Math.max(1, Math.ceil(delayMs / 1000)) if (totalSeconds < 60) { @@ -234,6 +353,7 @@ export type QueryParams = { // budget for the whole agentic turn; `remaining` is computed per iteration // from cumulative API usage. See configureTaskBudgetParams in claude.ts. taskBudget?: { total: number } + agentStepLimit?: AgentStepLimitConfig deps?: QueryDeps } @@ -263,6 +383,7 @@ type State = { // Why the previous iteration continued. Undefined on first iteration. // Lets tests assert recovery paths fired without inspecting message contents. transition: Continue | undefined + agentStepLimit: AgentStepLimitState | undefined } export async function* query( @@ -337,6 +458,7 @@ async function* queryLoop( continuationNudgeCount: 0, pendingToolUseSummary: undefined, transition: undefined, + agentStepLimit: normalizeAgentStepLimit(params.agentStepLimit), } const budgetTracker = feature('TOKEN_BUDGET') ? createBudgetTracker() : null @@ -388,6 +510,7 @@ async function* queryLoop( pendingToolUseSummary, stopHookActive, turnCount, + agentStepLimit, } = state const effectiveMaxOutputTokensOverride = maxOutputTokensOverride === undefined @@ -737,7 +860,8 @@ async function* queryLoop( let needsFollowUp = false queryCheckpoint('query_setup_start') - const useStreamingToolExecution = config.gates.streamingToolExecution + const useStreamingToolExecution = + config.gates.streamingToolExecution && agentStepLimit === undefined let streamingToolExecutor = useStreamingToolExecution ? new StreamingToolExecutor( toolUseContext.options.tools, @@ -877,6 +1001,9 @@ async function* queryLoop( } let attemptWithFallback = true + const toolsForModel = agentStepLimit?.summaryRequested + ? [] + : toolUseContext.options.tools queryCheckpoint('query_api_loop_start') try { @@ -889,7 +1016,7 @@ async function* queryLoop( messages: prependUserContext(messagesForQuery, userContext), systemPrompt: fullSystemPrompt, thinkingConfig: toolUseContext.options.thinkingConfig, - tools: toolUseContext.options.tools, + tools: toolsForModel, signal: toolUseContext.abortController.signal, options: { async getToolPermissionContext() { @@ -927,6 +1054,9 @@ async function* queryLoop( agentId: toolUseContext.agentId, addNotification: toolUseContext.addNotification, providerOverride: toolUseContext.options.providerOverride, + ...(toolsForModel !== toolUseContext.options.tools && { + messageNormalizationTools: toolUseContext.options.tools, + }), ...(params.taskBudget && { taskBudget: { total: params.taskBudget.total, @@ -1359,6 +1489,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'collapse_drain_retry', committed: drained.committed, @@ -1416,6 +1547,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'reactive_compact_retry' }, } state = next @@ -1477,6 +1609,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'provider_max_tokens_retry', cap: providerMaxTokensCap, @@ -1524,6 +1657,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'max_output_tokens_escalate' }, } state = next @@ -1555,6 +1689,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'max_output_tokens_recovery', attempt: maxOutputTokensRecoveryCount + 1, @@ -1634,6 +1769,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'provider_fallback_retry' }, } state = next @@ -1700,6 +1836,7 @@ async function* queryLoop( stopHookActive: stopHookResult.stopHookActive, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'stop_hook_blocking' }, } state = next @@ -1739,6 +1876,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount, + agentStepLimit, transition: { reason: 'token_budget_continuation' }, } continue @@ -1766,6 +1904,7 @@ async function* queryLoop( // when the model keeps matching signals without ever calling tools. if ( assistantMessages.length > 0 && + !agentStepLimit?.summaryRequested && turnCount < (maxTurns ?? Infinity) && state.continuationNudgeCount < MAX_CONTINUATION_NUDGES ) { @@ -1806,6 +1945,7 @@ async function* queryLoop( stopHookActive: undefined, turnCount, continuationNudgeCount: state.continuationNudgeCount + 1, + agentStepLimit, transition: { reason: 'continuation_nudge' }, } state = next @@ -1814,6 +1954,15 @@ async function* queryLoop( } } + if (agentStepLimit?.summaryRequested) { + return { + reason: 'agent_step_limit', + turnCount, + stepsUsed: agentStepLimit.stepsUsed, + maxSteps: agentStepLimit.maxSteps, + } + } + return { reason: 'completed' } } @@ -1822,6 +1971,38 @@ async function* queryLoop( queryCheckpoint('query_tool_execution_start') + let toolUseBlocksToExecute = toolUseBlocks + let blockedToolUseBlocks: ToolUseBlock[] = [] + let nextAgentStepLimit = agentStepLimit + let shouldRequestAgentStepSummary = false + const shouldTerminateAgentStepSummary = + agentStepLimit?.summaryRequested === true && toolUseBlocks.length > 0 + + if (agentStepLimit) { + if (agentStepLimit.summaryRequested) { + toolUseBlocksToExecute = [] + blockedToolUseBlocks = toolUseBlocks + } else { + const remainingSteps = Math.max( + 0, + agentStepLimit.maxSteps - agentStepLimit.stepsUsed, + ) + toolUseBlocksToExecute = toolUseBlocks.slice(0, remainingSteps) + blockedToolUseBlocks = toolUseBlocks.slice(remainingSteps) + const stepsUsed = + agentStepLimit.stepsUsed + toolUseBlocksToExecute.length + const summaryRequested = + stepsUsed >= agentStepLimit.maxSteps || + blockedToolUseBlocks.length > 0 + + nextAgentStepLimit = { + ...agentStepLimit, + stepsUsed, + summaryRequested, + } + shouldRequestAgentStepSummary = summaryRequested + } + } if (streamingToolExecutor) { logEvent('tengu_streaming_tool_execution_used', { @@ -1839,7 +2020,12 @@ async function* queryLoop( const toolUpdates = streamingToolExecutor ? streamingToolExecutor.getRemainingResults() - : runTools(toolUseBlocks, assistantMessages, canUseTool, toolUseContext) + : runTools( + toolUseBlocksToExecute, + assistantMessages, + canUseTool, + toolUseContext, + ) for await (const update of toolUpdates) { if (update.message) { @@ -1866,6 +2052,34 @@ async function* queryLoop( } } } + + if (nextAgentStepLimit && blockedToolUseBlocks.length > 0) { + for (const toolUse of blockedToolUseBlocks) { + const message = createAgentStepLimitToolResult( + toolUse, + findAssistantMessageForToolUse(assistantMessages, toolUse.id), + nextAgentStepLimit, + ) + yield message + toolResults.push(message) + } + } + + if (shouldTerminateAgentStepSummary && nextAgentStepLimit) { + if (!hasAssistantSummaryText(assistantMessages.at(-1))) { + yield createAgentStepLimitForcedSummary( + nextAgentStepLimit, + blockedToolUseBlocks.length, + ) + } + return { + reason: 'agent_step_limit', + turnCount, + stepsUsed: nextAgentStepLimit.stepsUsed, + maxSteps: nextAgentStepLimit.maxSteps, + } + } + queryCheckpoint('query_tool_execution_end') // Track multi-turn context after tool execution @@ -1970,6 +2184,24 @@ async function* queryLoop( return { reason: 'tool_failure_loop' } } + if (shouldRequestAgentStepSummary && nextAgentStepLimit) { + const summaryRequest = + createAgentStepLimitSummaryRequest(nextAgentStepLimit) + yield summaryRequest + toolResults.push(summaryRequest) + logForDebugging( + `[Agent: ${nextAgentStepLimit.agentType ?? 'subagent'}] Reached maxSteps limit (${nextAgentStepLimit.stepsUsed}/${nextAgentStepLimit.maxSteps}); requesting final summary`, + ) + logEvent('tengu_agent_step_limit_reached', { + agent_type: + (nextAgentStepLimit.agentType ?? + 'subagent') as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, + steps_used: nextAgentStepLimit.stepsUsed, + max_steps: nextAgentStepLimit.maxSteps, + blocked_tool_uses: blockedToolUseBlocks.length, + }) + } + // Generate tool use summary after tool batch completes — passed to next recursive call let nextPendingToolUseSummary: | Promise @@ -2229,7 +2461,11 @@ async function* queryLoop( } // Check if we've reached the max turns limit - if (maxTurns && nextTurnCount > maxTurns) { + if ( + maxTurns && + nextTurnCount > maxTurns && + !nextAgentStepLimit?.summaryRequested + ) { yield createAttachmentMessage({ type: 'max_turns_reached', maxTurns, @@ -2253,6 +2489,7 @@ async function* queryLoop( maxOutputTokensOverride: undefined, providerMaxOutputTokensCap, stopHookActive, + agentStepLimit: nextAgentStepLimit, transition: { reason: 'next_turn' }, } state = next diff --git a/src/query/agentStepLimit.test.ts b/src/query/agentStepLimit.test.ts new file mode 100644 index 000000000..ee9a8f639 --- /dev/null +++ b/src/query/agentStepLimit.test.ts @@ -0,0 +1,690 @@ +import { describe, expect, test } from 'bun:test' +import { z } from 'zod/v4' + +import { query, type QueryParams } from '../query.js' +import { buildTool, type Tools } from '../Tool.js' +import type { QueryDeps } from './deps.js' +import { + createAssistantMessage, + createUserMessage, + normalizeMessagesForAPI, +} from '../utils/messages.js' +import { asSystemPrompt } from '../utils/systemPromptType.js' +import { countToolUses } from '../tools/AgentTool/agentToolUtils.js' +import { AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX } from './agentStepLimit.js' + +const echoCalls: string[] = [] + +const echoTool = buildTool({ + name: 'Echo', + inputSchema: z.object({ text: z.string() }), + maxResultSizeChars: Infinity, + async description() { + return 'Echo input text' + }, + async prompt() { + return '' + }, + async call(input) { + echoCalls.push(input.text) + return { data: `echo:${input.text}` } + }, + mapToolResultToToolResultBlockParam(content, toolUseID) { + return { + type: 'tool_result', + tool_use_id: toolUseID, + content: String(content), + } + }, + renderToolUseMessage() { + return null + }, + renderToolResultMessage() { + return null + }, +}) + +function makeToolUseContext(tools: Tools = []): QueryParams['toolUseContext'] { + const abortController = new AbortController() + let inProgressToolUseIDs = new Set() + + return { + abortController, + agentId: 'agent-test', + getAppState: () => ({ + fastMode: false, + mcp: { tools: [], clients: [] }, + toolPermissionContext: { + mode: 'default', + additionalWorkingDirectories: new Map(), + alwaysAllowRules: {}, + alwaysDenyRules: {}, + alwaysAskRules: {}, + isBypassPermissionsModeAvailable: false, + }, + sessionHooks: new Map(), + mainLoopModel: 'gpt-4o', + effortValue: undefined, + advisorModel: undefined, + }), + options: { + commands: [], + debug: false, + thinkingConfig: { type: 'disabled' }, + tools, + verbose: false, + mcpClients: [], + mcpResources: {}, + isNonInteractiveSession: false, + agentDefinitions: { activeAgents: [], allAgents: [] }, + appendSystemPrompt: undefined, + providerOverride: undefined, + mainLoopModel: 'gpt-4o', + }, + addNotification: () => {}, + messages: [], + setInProgressToolUseIDs: updater => { + inProgressToolUseIDs = updater(inProgressToolUseIDs) + }, + setResponseLength: () => {}, + updateFileHistoryState: () => {}, + updateAttributionState: () => {}, + } as unknown as QueryParams['toolUseContext'] +} + +function makeParams( + callModel: QueryDeps['callModel'], + tools: Tools = [], + agentStepLimit?: { maxSteps: number; agentType: string }, +): QueryParams { + return { + messages: [createUserMessage({ content: 'inspect' })], + systemPrompt: asSystemPrompt([]), + userContext: {}, + systemContext: {}, + canUseTool: async () => ({ behavior: 'allow' }), + toolUseContext: makeToolUseContext(tools), + querySource: 'agent:builtin:general-purpose', + ...(agentStepLimit ? { agentStepLimit } : {}), + deps: { + callModel, + microcompact: async messages => ({ messages }), + autocompact: async () => ({ + compactionResult: null, + consecutiveFailures: undefined, + }), + uuid: () => '00000000-0000-4000-8000-000000000000', + } as unknown as QueryDeps, + } +} + +async function drain(params: QueryParams): Promise<{ + yielded: any[] + returned: any +}> { + const yielded: any[] = [] + const generator = query(params) + while (true) { + const next = await generator.next() + if (next.done) return { yielded, returned: next.value } + yielded.push(next.value) + } +} + +describe('agent step limits', () => { + test('without a configured limit, tool use behavior is unchanged', async () => { + echoCalls.length = 0 + let modelCalls = 0 + + const { yielded, returned } = await drain( + makeParams( + async function* () { + modelCalls++ + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_1', + name: 'Echo', + input: { text: 'first' }, + }, + { + type: 'tool_use', + id: 'toolu_echo_2', + name: 'Echo', + input: { text: 'second' }, + }, + ], + }) + return + } + yield createAssistantMessage({ content: 'done' }) + }, + [echoTool], + ), + ) + + expect(returned.reason).toBe('completed') + expect(modelCalls).toBe(2) + expect(echoCalls).toEqual(['first', 'second']) + expect( + yielded.some( + message => + message?.type === 'user' && + message?.isMeta && + typeof message.message.content === 'string' && + message.message.content.includes('configured step limit'), + ), + ).toBe(false) + }) + + test('invalid configured limit is ignored safely', async () => { + echoCalls.length = 0 + let modelCalls = 0 + + const { returned } = await drain( + makeParams( + async function* () { + modelCalls++ + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_1', + name: 'Echo', + input: { text: 'first' }, + }, + { + type: 'tool_use', + id: 'toolu_echo_2', + name: 'Echo', + input: { text: 'second' }, + }, + ], + }) + return + } + yield createAssistantMessage({ content: 'done' }) + }, + [echoTool], + { maxSteps: 0, agentType: 'general-purpose' }, + ), + ) + + expect(returned.reason).toBe('completed') + expect(modelCalls).toBe(2) + expect(echoCalls).toEqual(['first', 'second']) + }) + + test('configured limit stops further tool calls and requests a no-tools summary', async () => { + echoCalls.length = 0 + let modelCalls = 0 + const requestToolCounts: number[] = [] + const requestMessageNormalizationToolCounts: number[] = [] + const requestMessages: any[][] = [] + + const { yielded, returned } = await drain( + makeParams( + async function* ({ messages, options, tools }) { + modelCalls++ + requestToolCounts.push(tools.length) + requestMessageNormalizationToolCounts.push( + options.messageNormalizationTools?.length ?? 0, + ) + requestMessages.push(messages) + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_1', + name: 'Echo', + input: { text: 'allowed' }, + }, + { + type: 'tool_use', + id: 'toolu_echo_2', + name: 'Echo', + input: { text: 'blocked' }, + }, + ], + }) + return + } + yield createAssistantMessage({ + content: + 'Completed: checked the allowed step. Findings: limit reached. Remaining tasks: continue later. Another run needed: yes.', + }) + }, + [echoTool], + { maxSteps: 1, agentType: 'general-purpose' }, + ), + ) + + expect(returned).toMatchObject({ + reason: 'agent_step_limit', + turnCount: 2, + stepsUsed: 1, + maxSteps: 1, + }) + expect(modelCalls).toBe(2) + expect(requestToolCounts).toEqual([1, 0]) + expect(requestMessageNormalizationToolCounts).toEqual([0, 1]) + expect(echoCalls).toEqual(['allowed']) + expect(countToolUses(yielded)).toBe(1) + expect( + yielded.some( + message => + message?.type === 'assistant' && + message.message.content.some( + part => + part.type === 'text' && + part.text.includes('Completed: checked the allowed step') && + part.text.includes('Another run needed: yes'), + ), + ), + ).toBe(true) + expect( + yielded.some( + message => + message?.type === 'user' && + Array.isArray(message.message.content) && + message.message.content.some( + (part: any) => + part.type === 'tool_result' && + part.tool_use_id === 'toolu_echo_2' && + part.is_error === true && + String(part.content).includes('Agent step limit reached'), + ), + ), + ).toBe(true) + expect( + yielded.some( + message => + message?.type === 'user' && + message?.isMeta && + typeof message.message.content === 'string' && + message.message.content.includes('completed work') && + message.message.content.includes('remaining tasks') && + message.message.content.includes('another run is needed'), + ), + ).toBe(true) + expect( + requestMessages[1]?.some( + message => + message.type === 'user' && + message.isMeta && + typeof message.message.content === 'string' && + message.message.content.includes('configured step limit'), + ), + ).toBe(true) + + const normalizedSecondRequest = normalizeMessagesForAPI( + requestMessages[1] as any, + [echoTool], + ) + const summaryUser = normalizedSecondRequest.find( + message => + message.type === 'user' && + Array.isArray(message.message.content) && + message.message.content.some( + part => + part.type === 'text' && + part.text.includes('completed work') && + part.text.includes('another run is needed'), + ), + ) + expect(summaryUser).toBeDefined() + if ( + summaryUser?.type === 'user' && + Array.isArray(summaryUser.message.content) + ) { + const toolResultText = summaryUser.message.content + .filter(part => part.type === 'tool_result') + .map(part => String(part.content)) + .join('\n') + expect(toolResultText).not.toContain('completed work') + } + }) + + test('step count accumulates across turns before later tool calls are blocked', async () => { + echoCalls.length = 0 + let modelCalls = 0 + const requestToolCounts: number[] = [] + const requestMessageNormalizationToolCounts: number[] = [] + + const { yielded, returned } = await drain( + makeParams( + async function* ({ options, tools }) { + modelCalls++ + requestToolCounts.push(tools.length) + requestMessageNormalizationToolCounts.push( + options.messageNormalizationTools?.length ?? 0, + ) + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_turn_1', + name: 'Echo', + input: { text: 'first' }, + }, + ], + }) + return + } + if (modelCalls === 2) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_turn_2_allowed', + name: 'Echo', + input: { text: 'second' }, + }, + { + type: 'tool_use', + id: 'toolu_turn_2_blocked', + name: 'Echo', + input: { text: 'third' }, + }, + ], + }) + return + } + yield createAssistantMessage({ + content: + 'Completed work: handled two allowed steps. Findings: a later step was blocked. Remaining tasks: continue later. Another run needed: yes.', + }) + }, + [echoTool], + { maxSteps: 2, agentType: 'general-purpose' }, + ), + ) + + expect(returned).toMatchObject({ + reason: 'agent_step_limit', + turnCount: 3, + stepsUsed: 2, + maxSteps: 2, + }) + expect(modelCalls).toBe(3) + expect(requestToolCounts).toEqual([1, 1, 0]) + expect(requestMessageNormalizationToolCounts).toEqual([0, 0, 1]) + expect(echoCalls).toEqual(['first', 'second']) + expect(countToolUses(yielded)).toBe(2) + expect( + yielded.some( + message => + message?.type === 'user' && + Array.isArray(message.message.content) && + message.message.content.some( + (part: any) => + part.type === 'tool_result' && + part.tool_use_id === 'toolu_turn_2_blocked' && + part.is_error === true && + String(part.content).includes('Agent step limit reached'), + ), + ), + ).toBe(true) + expect( + yielded.some( + message => + message?.type === 'assistant' && + message.message.content.some( + part => + part.type === 'text' && + part.text.includes('Completed work: handled two allowed steps') && + part.text.includes('Another run needed: yes'), + ), + ), + ).toBe(true) + }) + + test('multiple over-limit tool calls do not trip the failure-loop guard', async () => { + echoCalls.length = 0 + let modelCalls = 0 + + const { returned } = await drain( + makeParams( + async function* () { + modelCalls++ + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_allowed', + name: 'Echo', + input: { text: 'allowed' }, + }, + ...[1, 2, 3, 4].map(i => ({ + type: 'tool_use' as const, + id: `toolu_echo_blocked_${i}`, + name: 'Echo', + input: { text: `blocked-${i}` }, + })), + ], + }) + return + } + yield createAssistantMessage({ + content: + 'Completed work: one allowed step. Findings: extra calls were blocked. Remaining tasks: continue later. Another run needed: yes.', + }) + }, + [echoTool], + { maxSteps: 1, agentType: 'general-purpose' }, + ), + ) + + expect(returned).toMatchObject({ + reason: 'agent_step_limit', + turnCount: 2, + stepsUsed: 1, + maxSteps: 1, + }) + expect(modelCalls).toBe(2) + expect(echoCalls).toEqual(['allowed']) + }) + + test('forced summary turn cannot execute more tools', async () => { + echoCalls.length = 0 + let modelCalls = 0 + + const { yielded, returned } = await drain( + makeParams( + async function* () { + modelCalls++ + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_1', + name: 'Echo', + input: { text: 'allowed' }, + }, + ], + }) + return + } + yield createAssistantMessage({ + content: [ + { + type: 'text', + text: 'I should inspect one more thing first.', + citations: null, + }, + { + type: 'tool_use', + id: 'toolu_echo_summary', + name: 'Echo', + input: { text: 'must-not-run' }, + }, + ], + }) + }, + [echoTool], + { maxSteps: 1, agentType: 'general-purpose' }, + ), + ) + + expect(returned).toMatchObject({ + reason: 'agent_step_limit', + turnCount: 2, + stepsUsed: 1, + maxSteps: 1, + }) + expect(modelCalls).toBe(2) + expect(echoCalls).toEqual(['allowed']) + expect(countToolUses(yielded)).toBe(1) + expect( + yielded.some( + message => + message?.type === 'user' && + Array.isArray(message.message.content) && + message.message.content.some( + (part: any) => + part.type === 'tool_result' && + part.tool_use_id === 'toolu_echo_summary' && + part.is_error === true && + String(part.content).includes('Agent step limit reached'), + ), + ), + ).toBe(true) + + const finalAssistantMessage = yielded + .filter(message => message?.type === 'assistant') + .at(-1) + expect( + finalAssistantMessage?.message.content.some( + part => + part.type === 'text' && + part.text.includes('Completed work: Agent') && + part.text.includes( + 'Findings: 1 additional tool call was blocked', + ) && + part.text.includes('Another run needed: yes'), + ), + ).toBe(true) + }) + + test('forced summary turn does not add a duplicate synthetic summary after a valid model summary', async () => { + echoCalls.length = 0 + let modelCalls = 0 + + const { yielded, returned } = await drain( + makeParams( + async function* () { + modelCalls++ + if (modelCalls === 1) { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_echo_1', + name: 'Echo', + input: { text: 'allowed' }, + }, + ], + }) + return + } + yield createAssistantMessage({ + content: [ + { + type: 'text', + text: 'Completed work: one step.', + citations: null, + }, + { + type: 'text', + text: 'Findings: the limit was reached.', + citations: null, + }, + { + type: 'text', + text: 'Remaining tasks: continue later.', + citations: null, + }, + { + type: 'text', + text: 'Another run needed: yes.', + citations: null, + }, + { + type: 'tool_use', + id: 'toolu_echo_summary', + name: 'Echo', + input: { text: 'must-not-run' }, + }, + ], + }) + }, + [echoTool], + { maxSteps: 1, agentType: 'general-purpose' }, + ), + ) + + expect(returned).toMatchObject({ + reason: 'agent_step_limit', + turnCount: 2, + stepsUsed: 1, + maxSteps: 1, + }) + expect(modelCalls).toBe(2) + expect(echoCalls).toEqual(['allowed']) + expect(countToolUses(yielded)).toBe(1) + + const assistantMessages = yielded.filter( + message => message?.type === 'assistant', + ) + expect(assistantMessages).toHaveLength(2) + const finalAssistantText = assistantMessages + .at(-1) + ?.message.content.filter(part => part.type === 'text') + .map(part => part.text) + .join('\n') + expect(finalAssistantText).toContain('Completed work: one step') + expect(finalAssistantText).toContain('Another run needed: yes') + expect( + assistantMessages.some(message => + message.message.content.some( + part => + part.type === 'text' && + part.text.includes("Agent 'general-purpose' reached"), + ), + ), + ).toBe(false) + }) + + test('real tool output with the readable limit prefix still counts as a tool use', () => { + const messages = [ + createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_real_output', + name: 'Echo', + input: { text: 'prefix-collision' }, + }, + ], + }), + createUserMessage({ + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_real_output', + content: `${AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX}: this is real tool output, not synthetic`, + }, + ], + }), + ] + + expect(countToolUses(messages)).toBe(1) + }) +}) diff --git a/src/query/agentStepLimit.ts b/src/query/agentStepLimit.ts new file mode 100644 index 000000000..e1ad81757 --- /dev/null +++ b/src/query/agentStepLimit.ts @@ -0,0 +1,2 @@ +export const AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX = + 'Agent step limit reached' diff --git a/src/query/toolFailureLoopGuard.test.ts b/src/query/toolFailureLoopGuard.test.ts index 81ca15c31..49e705389 100644 --- a/src/query/toolFailureLoopGuard.test.ts +++ b/src/query/toolFailureLoopGuard.test.ts @@ -26,9 +26,15 @@ function toolResult( toolUseId: string, content: string, isError = true, -): { type: 'user'; message: { content: unknown[] } } { + isAgentStepLimitToolResult = false, +): { + type: 'user' + isAgentStepLimitToolResult?: boolean + message: { content: unknown[] } +} { return { type: 'user', + ...(isAgentStepLimitToolResult ? { isAgentStepLimitToolResult } : {}), message: { content: [ { @@ -274,6 +280,35 @@ test('real tool errors that merely mention ignored phrases are still counted', ( expect(decision.tripped).toBe(true) }) +test('agent step-limit text is ignored only with the structured message flag', () => { + const spoofedState = createToolFailureLoopGuardState() + + update(spoofedState, [toolUse('a', 'Bash')], [ + toolResult('a', 'Agent step limit reached while parsing logs'), + ]) + const spoofedDecision = update( + spoofedState, + [toolUse('b', 'Bash')], + [toolResult('b', 'Agent step limit reached while parsing logs')], + 2, + ) + + expect(spoofedDecision.tripped).toBe(true) + + const syntheticState = createToolFailureLoopGuardState() + update(syntheticState, [toolUse('c', 'Bash')], [ + toolResult('c', 'Agent step limit reached for subagent', true, true), + ]) + const syntheticDecision = update( + syntheticState, + [toolUse('d', 'Bash')], + [toolResult('d', 'Agent step limit reached for subagent', true, true)], + 2, + ) + + expect(syntheticDecision.tripped).toBe(false) +}) + test('same failing file_path across repeated failures trips the guard', () => { const state = createToolFailureLoopGuardState() diff --git a/src/query/toolFailureLoopGuard.ts b/src/query/toolFailureLoopGuard.ts index 1fbec945a..32537d634 100644 --- a/src/query/toolFailureLoopGuard.ts +++ b/src/query/toolFailureLoopGuard.ts @@ -88,7 +88,10 @@ export function updateToolFailureLoopGuard(params: { continue } - if (isIgnoredSyntheticToolResult(content)) { + if ( + block.isAgentStepLimitToolResult || + isIgnoredSyntheticToolResult(content) + ) { continue } @@ -202,6 +205,7 @@ type ToolResultBlockLike = { tool_use_id?: unknown content?: unknown is_error?: unknown + isAgentStepLimitToolResult?: boolean } type FailureInfo = { @@ -264,7 +268,10 @@ function getToolResultBlocks( for (const block of message.message.content) { if (isToolResultBlock(block)) { - blocks.push(block) + blocks.push({ + ...block, + isAgentStepLimitToolResult: message.isAgentStepLimitToolResult, + }) } } } @@ -290,9 +297,12 @@ function toolResultContentToString(content: unknown): string { } if (typeof content === 'object' && content !== null) { - const text = (content as { text?: unknown }).text - if (typeof text === 'string') { - return text + const block = content as { text?: unknown; content?: unknown } + if (block.text !== undefined) { + return toolResultContentToString(block.text) + } + if (block.content !== undefined) { + return toolResultContentToString(block.content) } } diff --git a/src/query/transitions.ts b/src/query/transitions.ts index 95f2b0edf..40e052305 100644 --- a/src/query/transitions.ts +++ b/src/query/transitions.ts @@ -9,6 +9,12 @@ export type Terminal = | { reason: 'aborted_tools' } | { reason: 'hook_stopped' } | { reason: 'max_turns'; turnCount: number } + | { + reason: 'agent_step_limit' + turnCount: number + stepsUsed: number + maxSteps: number + } | { reason: 'tool_failure_loop' } export type Continue = diff --git a/src/services/api/claude.ts b/src/services/api/claude.ts index ad3435195..3633957a0 100644 --- a/src/services/api/claude.ts +++ b/src/services/api/claude.ts @@ -718,6 +718,7 @@ export type Options = { taskBudget?: { total: number; remaining?: number } providerOverride?: { model: string; baseURL: string; apiKey: string } queryLifecycle?: QueryLifecycleOperationTracker + messageNormalizationTools?: Tools } export async function queryModelWithoutStreaming({ @@ -1294,7 +1295,10 @@ async function* queryModel( }) queryCheckpoint('query_message_normalization_start') - let messagesForAPI = normalizeMessagesForAPI(messages, filteredTools) + let messagesForAPI = normalizeMessagesForAPI( + messages, + options.messageNormalizationTools ?? filteredTools, + ) queryCheckpoint('query_message_normalization_end') // Apply hybrid context strategy for optimal cache/fresh balance diff --git a/src/tools/AgentTool/agentDisplay.ts b/src/tools/AgentTool/agentDisplay.ts index 2090ae999..4d64e0007 100644 --- a/src/tools/AgentTool/agentDisplay.ts +++ b/src/tools/AgentTool/agentDisplay.ts @@ -10,7 +10,7 @@ import { } from '../../utils/settings/constants.js' import type { AgentDefinition } from './loadAgentsDir.js' -type AgentSource = SettingSource | 'built-in' | 'plugin' +type AgentSource = SettingSource | 'built-in' | 'plugin' | 'sdk' export type AgentSourceGroup = { label: string @@ -28,6 +28,7 @@ export const AGENT_SOURCE_GROUPS: AgentSourceGroup[] = [ { label: 'Managed agents', source: 'policySettings' }, { label: 'Plugin agents', source: 'plugin' }, { label: 'CLI arg agents', source: 'flagSettings' }, + { label: 'SDK agents', source: 'sdk' }, { label: 'Built-in agents', source: 'built-in' }, ] diff --git a/src/tools/AgentTool/agentToolUtils.ts b/src/tools/AgentTool/agentToolUtils.ts index a566ac317..64e4a473d 100644 --- a/src/tools/AgentTool/agentToolUtils.ts +++ b/src/tools/AgentTool/agentToolUtils.ts @@ -260,11 +260,30 @@ export const agentToolResultSchema = lazySchema(() => export type AgentToolResult = z.input> export function countToolUses(messages: MessageType[]): number { + const blockedStepLimitToolUseIds = new Set() + for (const m of messages) { + if ( + m.type !== 'user' || + !m.isAgentStepLimitToolResult || + !Array.isArray(m.message.content) + ) { + continue + } + for (const block of m.message.content) { + if (block.type === 'tool_result') { + blockedStepLimitToolUseIds.add(String(block.tool_use_id)) + } + } + } + let count = 0 for (const m of messages) { if (m.type === 'assistant') { for (const block of m.message.content) { - if (block.type === 'tool_use') { + if ( + block.type === 'tool_use' && + !blockedStepLimitToolUseIds.has(block.id) + ) { count++ } } diff --git a/src/tools/AgentTool/loadAgentsDir.test.ts b/src/tools/AgentTool/loadAgentsDir.test.ts index 2c0787939..944f63e1a 100644 --- a/src/tools/AgentTool/loadAgentsDir.test.ts +++ b/src/tools/AgentTool/loadAgentsDir.test.ts @@ -2,17 +2,28 @@ import { afterEach, beforeEach, describe, expect, test } from 'bun:test' import { mkdtemp, mkdir, rm, writeFile } from 'fs/promises' import { tmpdir } from 'os' import { dirname, join } from 'path' +import { setAllowedSettingSources } from '../../bootstrap/state.js' +import { + getClaudeConfigHomeDir, + getClaudeConfigHomeDirOverrideForTesting, + setClaudeConfigHomeDirForTesting, +} from '../../utils/envUtils.js' import { clearAgentDefinitionsCache, getAgentDefinitionsWithOverrides, + parseAgentFromJson, } from './loadAgentsDir.js' import { loadMarkdownFilesForSubdir } from '../../utils/markdownConfigLoader.js' +import { SETTING_SOURCES } from '../../utils/settings/constants.js' +import { resetSettingsCache } from '../../utils/settings/settingsCache.js' import { acquireSharedMutationLock, releaseSharedMutationLock, } from '../../test/sharedMutationLock.js' const originalEnv = { + HOME: process.env.HOME, + OPENCLAUDE_CONFIG_DIR: process.env.OPENCLAUDE_CONFIG_DIR, CLAUDE_CONFIG_DIR: process.env.CLAUDE_CONFIG_DIR, CLAUDE_CODE_SIMPLE: process.env.CLAUDE_CODE_SIMPLE, CLAUDE_CODE_USE_NATIVE_FILE_SEARCH: @@ -21,24 +32,52 @@ const originalEnv = { } let tempDir: string +let projectRootDir: string +let userConfigDir: string +let previousConfigHomeOverride: string | undefined beforeEach(async () => { await acquireSharedMutationLock('loadAgentsDir.test.ts') tempDir = await mkdtemp(join(tmpdir(), 'openclaude-agents-test-')) - process.env.CLAUDE_CONFIG_DIR = join(tempDir, '.openclaude') + projectRootDir = await mkdtemp(join(tmpdir(), 'openclaude-agents-project-')) + const configDir = join(tempDir, '.openclaude') + previousConfigHomeOverride = getClaudeConfigHomeDirOverrideForTesting() + setClaudeConfigHomeDirForTesting(configDir) + process.env.HOME = tempDir + process.env.OPENCLAUDE_CONFIG_DIR = configDir + process.env.CLAUDE_CONFIG_DIR = configDir process.env.CLAUDE_CODE_USE_NATIVE_FILE_SEARCH = '1' delete process.env.CLAUDE_CODE_SIMPLE + setAllowedSettingSources([...SETTING_SOURCES]) + getClaudeConfigHomeDir.cache?.clear?.() + const resolvedConfigDir = getClaudeConfigHomeDir() + userConfigDir = resolvedConfigDir.startsWith(join(tmpdir(), '')) + ? resolvedConfigDir + : configDir + resetSettingsCache() clearAgentDefinitionsCache() loadMarkdownFilesForSubdir.cache.clear?.() }) afterEach(async () => { try { + await rm(join(userConfigDir, 'agents', 'user-agent.md'), { force: true }) + await rm(join(userConfigDir, 'agents', 'shared-limited.md'), { + force: true, + }) await rm(tempDir, { recursive: true, force: true }) + await rm(projectRootDir, { recursive: true, force: true }) + restoreEnv('HOME') + restoreEnv('OPENCLAUDE_CONFIG_DIR') restoreEnv('CLAUDE_CONFIG_DIR') restoreEnv('CLAUDE_CODE_SIMPLE') restoreEnv('CLAUDE_CODE_USE_NATIVE_FILE_SEARCH') restoreEnv('USER_TYPE') + setAllowedSettingSources([...SETTING_SOURCES]) + setClaudeConfigHomeDirForTesting(previousConfigHomeOverride) + previousConfigHomeOverride = undefined + getClaudeConfigHomeDir.cache?.clear?.() + resetSettingsCache() clearAgentDefinitionsCache() loadMarkdownFilesForSubdir.cache.clear?.() } finally { @@ -78,7 +117,7 @@ ${prompt} describe('agent definition loading', () => { test('loads user agents from the OpenClaude config dir in simple mode', async () => { await writeAgent( - join(process.env.CLAUDE_CONFIG_DIR!, 'agents', 'user-agent.md'), + join(userConfigDir, 'agents', 'user-agent.md'), 'user-agent', ) @@ -86,7 +125,9 @@ describe('agent definition loading', () => { clearAgentDefinitionsCache() loadMarkdownFilesForSubdir.cache.clear?.() - const { activeAgents } = await getAgentDefinitionsWithOverrides(tempDir) + const { activeAgents } = await getAgentDefinitionsWithOverrides( + projectRootDir, + ) expect(activeAgents.some(agent => agent.agentType === 'user-agent')).toBe( true, @@ -94,7 +135,7 @@ describe('agent definition loading', () => { }) test('loads project agents from .openclaude/agents', async () => { - const projectDir = join(tempDir, 'project') + const projectDir = join(projectRootDir, 'project') await writeAgent( join(projectDir, '.openclaude', 'agents', 'project-agent.md'), 'project-agent', @@ -108,7 +149,7 @@ describe('agent definition loading', () => { }) test('prefers .openclaude project agents over legacy .claude agents', async () => { - const projectDir = join(tempDir, 'project') + const projectDir = join(projectRootDir, 'project') await writeAgent( join(projectDir, '.claude', 'agents', 'shared-agent.md'), 'shared-agent', @@ -127,7 +168,7 @@ describe('agent definition loading', () => { }) test('accepts worktree isolation in markdown agent frontmatter', async () => { - const projectDir = join(tempDir, 'project') + const projectDir = join(projectRootDir, 'project') await writeAgent( join(projectDir, '.openclaude', 'agents', 'worktree-agent.md'), 'worktree-agent', @@ -143,7 +184,7 @@ describe('agent definition loading', () => { test('rejects removed remote isolation in markdown agent frontmatter', async () => { process.env.USER_TYPE = 'ant' - const projectDir = join(tempDir, 'project') + const projectDir = join(projectRootDir, 'project') await writeAgent( join(projectDir, '.openclaude', 'agents', 'remote-agent.md'), 'remote-agent', @@ -157,4 +198,86 @@ describe('agent definition loading', () => { expect(agent).toBeDefined() expect(agent?.isolation).toBeUndefined() }) + + test('loads maxSteps from markdown agent frontmatter', async () => { + const projectDir = join(projectRootDir, 'project') + await writeAgent( + join(projectDir, '.openclaude', 'agents', 'limited-agent.md'), + 'limited-agent', + 'limited prompt', + 'maxSteps: 3\n', + ) + + const { activeAgents } = await getAgentDefinitionsWithOverrides(projectDir) + const agent = activeAgents.find(agent => agent.agentType === 'limited-agent') + + expect(agent?.maxSteps).toBe(3) + }) + + test('ignores invalid maxSteps in markdown agent frontmatter', async () => { + const projectDir = join(projectRootDir, 'project') + await writeAgent( + join(projectDir, '.openclaude', 'agents', 'invalid-steps-agent.md'), + 'invalid-steps-agent', + 'invalid steps prompt', + 'maxSteps: 0\n', + ) + await writeAgent( + join(projectDir, '.openclaude', 'agents', 'malformed-steps-agent.md'), + 'malformed-steps-agent', + 'malformed steps prompt', + 'maxSteps: 2abc\n', + ) + + const { activeAgents } = await getAgentDefinitionsWithOverrides(projectDir) + const agent = activeAgents.find( + agent => agent.agentType === 'invalid-steps-agent', + ) + const malformed = activeAgents.find( + agent => agent.agentType === 'malformed-steps-agent', + ) + + expect(agent).toBeDefined() + expect(agent?.maxSteps).toBeUndefined() + expect(malformed).toBeDefined() + expect(malformed?.maxSteps).toBeUndefined() + }) + + test('project agent maxSteps overrides user agent maxSteps for the same name', async () => { + const projectDir = join(projectRootDir, 'project') + await writeAgent( + join(userConfigDir, 'agents', 'shared-limited.md'), + 'shared-limited', + 'user prompt', + 'maxSteps: 1\n', + ) + await writeAgent( + join(projectDir, '.openclaude', 'agents', 'shared-limited.md'), + 'shared-limited', + 'project prompt', + 'maxSteps: 5\n', + ) + + const { activeAgents } = await getAgentDefinitionsWithOverrides(projectDir) + const agent = activeAgents.find(agent => agent.agentType === 'shared-limited') + + expect(agent?.source).toBe('projectSettings') + expect(agent?.maxSteps).toBe(5) + }) + + test('loads maxSteps from JSON agent definitions and rejects invalid values safely', () => { + const valid = parseAgentFromJson('json-limited', { + description: 'Use for JSON maxSteps coverage', + prompt: 'JSON prompt', + maxSteps: 2, + }) + const invalid = parseAgentFromJson('json-invalid', { + description: 'Use for invalid JSON maxSteps coverage', + prompt: 'JSON prompt', + maxSteps: 0, + }) + + expect(valid?.maxSteps).toBe(2) + expect(invalid).toBeNull() + }) }) diff --git a/src/tools/AgentTool/loadAgentsDir.ts b/src/tools/AgentTool/loadAgentsDir.ts index 0ef5b2ff3..98da42ad0 100644 --- a/src/tools/AgentTool/loadAgentsDir.ts +++ b/src/tools/AgentTool/loadAgentsDir.ts @@ -86,6 +86,7 @@ const AgentJsonSchema = lazySchema(() => mcpServers: z.array(AgentMcpServerSpecSchema()).optional(), hooks: HooksSchema().optional(), maxTurns: z.number().int().positive().optional(), + maxSteps: z.number().int().positive().optional(), skills: z.array(z.string()).optional(), initialPrompt: z.string().optional(), memory: z.enum(['user', 'project', 'local']).optional(), @@ -112,6 +113,7 @@ export type BaseAgentDefinition = { effort?: EffortValue permissionMode?: PermissionMode maxTurns?: number // Maximum number of agentic turns before stopping + maxSteps?: number // Maximum number of tool-use steps before forcing a final summary filename?: string // Original filename without .md extension (for user/project/managed agents) baseDir?: string criticalSystemReminder_EXPERIMENTAL?: string // Short message re-injected at every user turn @@ -146,6 +148,11 @@ export type CustomAgentDefinition = BaseAgentDefinition & { baseDir?: string } +export type SdkAgentDefinition = BaseAgentDefinition & { + getSystemPrompt: () => string + source: 'sdk' +} + // Plugin agents - similar to custom but with plugin metadata, prompt stored via closure export type PluginAgentDefinition = BaseAgentDefinition & { getSystemPrompt: () => string @@ -158,6 +165,7 @@ export type PluginAgentDefinition = BaseAgentDefinition & { export type AgentDefinition = | BuiltInAgentDefinition | CustomAgentDefinition + | SdkAgentDefinition | PluginAgentDefinition // Type guards for runtime type checking @@ -170,7 +178,11 @@ export function isBuiltInAgent( export function isCustomAgent( agent: AgentDefinition, ): agent is CustomAgentDefinition { - return agent.source !== 'built-in' && agent.source !== 'plugin' + return ( + agent.source !== 'built-in' && + agent.source !== 'plugin' && + agent.source !== 'sdk' + ) } export function isPluginAgent( @@ -195,6 +207,7 @@ export function getActiveAgentsFromList( const projectAgents = allAgents.filter(a => a.source === 'projectSettings') const managedAgents = allAgents.filter(a => a.source === 'policySettings') const flagAgents = allAgents.filter(a => a.source === 'flagSettings') + const sdkAgents = allAgents.filter(a => a.source === 'sdk') const agentGroups = [ builtInAgents, @@ -202,6 +215,7 @@ export function getActiveAgentsFromList( userAgents, projectAgents, flagAgents, + sdkAgents, managedAgents, ] @@ -484,6 +498,7 @@ export function parseAgentFromJson( : {}), ...(parsed.hooks ? { hooks: parsed.hooks } : {}), ...(parsed.maxTurns !== undefined ? { maxTurns: parsed.maxTurns } : {}), + ...(parsed.maxSteps !== undefined ? { maxSteps: parsed.maxSteps } : {}), ...(parsed.skills && parsed.skills.length > 0 ? { skills: parsed.skills } : {}), @@ -639,6 +654,15 @@ export function parseAgentFromMarkdown( ) } + // Parse maxSteps from frontmatter + const maxStepsRaw = frontmatter['maxSteps'] + const maxSteps = parsePositiveIntFromFrontmatter(maxStepsRaw) + if (maxStepsRaw !== undefined && maxSteps === undefined) { + logForDebugging( + `Agent file ${filePath} has invalid maxSteps '${maxStepsRaw}'. Must be a positive integer.`, + ) + } + // Extract filename without extension const filename = basename(filePath, '.md') @@ -727,6 +751,7 @@ export function parseAgentFromMarkdown( ? { permissionMode: permissionModeRaw as PermissionMode } : {}), ...(maxTurns !== undefined ? { maxTurns } : {}), + ...(maxSteps !== undefined ? { maxSteps } : {}), ...(background ? { background } : {}), ...(memory ? { memory } : {}), ...(isolation ? { isolation } : {}), diff --git a/src/tools/AgentTool/runAgent.ts b/src/tools/AgentTool/runAgent.ts index 9bc7ad745..341840cb9 100644 --- a/src/tools/AgentTool/runAgent.ts +++ b/src/tools/AgentTool/runAgent.ts @@ -13,6 +13,7 @@ import type { QuerySource } from '../../constants/querySource.js' import { getSystemContext, getUserContext } from '../../context.js' import type { CanUseToolFn } from '../../hooks/useCanUseTool.js' import { query } from '../../query.js' +import type { Terminal } from '../../query/transitions.js' import { getFeatureValue_CACHED_MAY_BE_STALE } from '../../services/analytics/growthbook.js' import { getDumpPromptsPath } from '../../services/api/dumpPrompts.js' import { cleanupAgentTracking } from '../../services/api/promptCacheBreakDetection.js' @@ -254,6 +255,7 @@ export async function* runAgent({ override, model, maxTurns, + maxSteps, preserveToolUseResults, availableTools, allowedTools, @@ -286,6 +288,7 @@ export async function* runAgent({ } model?: string maxTurns?: number + maxSteps?: number /** Preserve toolUseResult on messages for subagents with viewable transcripts */ preserveToolUseResults?: boolean /** Precomputed tool pool for the worker agent. Computed by the caller @@ -786,7 +789,15 @@ export async function* runAgent({ let lastRecordedUuid: UUID | null = initialMessages.at(-1)?.uuid ?? null try { - for await (const message of query({ + let queryTerminal: Terminal | undefined + const configuredMaxSteps = + Number.isSafeInteger(maxSteps) && maxSteps! > 0 + ? maxSteps + : Number.isSafeInteger(agentDefinition.maxSteps) && + agentDefinition.maxSteps! > 0 + ? agentDefinition.maxSteps + : undefined + const queryIterator = query({ messages: initialMessages, systemPrompt: agentSystemPrompt, userContext: resolvedUserContext, @@ -795,55 +806,74 @@ export async function* runAgent({ toolUseContext: agentToolUseContext, querySource, maxTurns: maxTurns ?? agentDefinition.maxTurns, - })) { - onQueryProgress?.() - // Forward subagent API request starts to parent's metrics display - // so TTFT/OTPS update during subagent execution. - if ( - message.type === 'stream_event' && - message.event.type === 'message_start' && - message.ttftMs != null - ) { - toolUseContext.pushApiMetricsEntry?.(message.ttftMs) - continue - } + agentStepLimit: + configuredMaxSteps !== undefined + ? { + maxSteps: configuredMaxSteps, + agentType: agentDefinition.agentType, + } + : undefined, + })[Symbol.asyncIterator]() - // Yield attachment messages (e.g., structured_output) without recording them - if (message.type === 'attachment') { - // Handle max turns reached signal from query.ts - if (message.attachment.type === 'max_turns_reached') { - logForDebugging( - `[Agent -: $ -{ - agentDefinition.agentType -} -] Reached max turns limit ($ -{ - message.attachment.maxTurns -} -)`, - ) + try { + while (true) { + const next = await queryIterator.next() + if (next.done) { + queryTerminal = next.value break } - yield message - continue - } - if (isRecordableMessage(message)) { - // Record only the new message with correct parent (O(1) per message) - await recordSidechainTranscript( - [message], - agentId, - lastRecordedUuid, - ).catch(err => - logForDebugging(`Failed to record sidechain transcript: ${err}`), - ) - if (message.type !== 'progress') { - lastRecordedUuid = message.uuid + const message = next.value + onQueryProgress?.() + // Forward subagent API request starts to parent's metrics display + // so TTFT/OTPS update during subagent execution. + if ( + message.type === 'stream_event' && + message.event.type === 'message_start' && + message.ttftMs != null + ) { + toolUseContext.pushApiMetricsEntry?.(message.ttftMs) + continue + } + + // Yield attachment messages (e.g., structured_output) without recording them + if (message.type === 'attachment') { + // Handle max turns reached signal from query.ts + if (message.attachment.type === 'max_turns_reached') { + logForDebugging( + `[Agent: ${agentDefinition.agentType}] Reached max turns limit (${message.attachment.maxTurns})`, + ) + break + } + yield message + continue + } + + if (isRecordableMessage(message)) { + // Record only the new message with correct parent (O(1) per message) + await recordSidechainTranscript( + [message], + agentId, + lastRecordedUuid, + ).catch(err => + logForDebugging(`Failed to record sidechain transcript: ${err}`), + ) + if (message.type !== 'progress') { + lastRecordedUuid = message.uuid + } + yield message } - yield message } + } finally { + if (queryTerminal === undefined) { + await queryIterator.return?.(undefined as never) + } + } + + if (queryTerminal?.reason === 'agent_step_limit') { + logForDebugging( + `[Agent: ${agentDefinition.agentType}] Stopped after reaching maxSteps (${queryTerminal.stepsUsed}/${queryTerminal.maxSteps})`, + ) } if (agentAbortController.signal.aborted) { diff --git a/src/types/message.ts b/src/types/message.ts index 7b7f7eb72..ff7639089 100644 --- a/src/types/message.ts +++ b/src/types/message.ts @@ -116,6 +116,8 @@ export interface UserMessage { } /** Matches the tool's `Output` type for tool_result messages. */ toolUseResult?: unknown + /** Internal marker for synthetic tool_result messages created by agent step limits. */ + isAgentStepLimitToolResult?: boolean /** MCP protocol metadata passed through to SDK consumers. */ mcpMeta?: { _meta?: Record diff --git a/src/utils/analyzeContext.ts b/src/utils/analyzeContext.ts index e54a7834f..e1073758e 100644 --- a/src/utils/analyzeContext.ts +++ b/src/utils/analyzeContext.ts @@ -176,7 +176,7 @@ export interface SystemPromptSectionDetail { interface Agent { agentType: string - source: SettingSource | 'built-in' | 'plugin' + source: SettingSource | 'built-in' | 'plugin' | 'sdk' tokens: number } diff --git a/src/utils/frontmatterParser.ts b/src/utils/frontmatterParser.ts index 543dd3e7e..c8bbf7bd6 100644 --- a/src/utils/frontmatterParser.ts +++ b/src/utils/frontmatterParser.ts @@ -335,9 +335,14 @@ export function parsePositiveIntFromFrontmatter( return undefined } - const parsed = typeof value === 'number' ? value : parseInt(String(value), 10) + const parsed = + typeof value === 'number' + ? value + : /^\d+$/.test(String(value).trim()) + ? Number(String(value).trim()) + : NaN - if (Number.isInteger(parsed) && parsed > 0) { + if (Number.isSafeInteger(parsed) && parsed > 0) { return parsed } diff --git a/src/utils/messages.ts b/src/utils/messages.ts index fd3652da0..6f8255842 100644 --- a/src/utils/messages.ts +++ b/src/utils/messages.ts @@ -471,6 +471,7 @@ export function createUserMessage({ isCollapseSummary, summarizeMetadata, toolUseResult, + isAgentStepLimitToolResult, mcpMeta, uuid, timestamp, @@ -486,6 +487,7 @@ export function createUserMessage({ isCompactSummary?: boolean isCollapseSummary?: boolean toolUseResult?: unknown // Matches tool's `Output` type + isAgentStepLimitToolResult?: boolean /** MCP protocol metadata to pass through to SDK consumers (never sent to model) */ mcpMeta?: { _meta?: Record @@ -526,6 +528,7 @@ export function createUserMessage({ uuid: (uuid as UUID | undefined) || randomUUID(), timestamp: timestamp ?? new Date().toISOString(), toolUseResult, + isAgentStepLimitToolResult, mcpMeta, imagePasteIds, sourceToolAssistantUUID, diff --git a/src/utils/plugins/loadPluginAgents.test.ts b/src/utils/plugins/loadPluginAgents.test.ts new file mode 100644 index 000000000..e8b1e69fd --- /dev/null +++ b/src/utils/plugins/loadPluginAgents.test.ts @@ -0,0 +1,143 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { mkdtemp, mkdir, rm, writeFile } from 'fs/promises' +import { tmpdir } from 'os' +import { join } from 'path' + +import { setInlinePlugins } from '../../bootstrap/state.js' +import { + acquireSharedMutationLock, + releaseSharedMutationLock, +} from '../../test/sharedMutationLock.js' +import { clearPluginCache } from './pluginLoader.js' +import { clearPluginAgentCache, loadPluginAgents } from './loadPluginAgents.js' + +let tempDir: string + +beforeEach(async () => { + await acquireSharedMutationLock('utils/plugins/loadPluginAgents.test.ts') + tempDir = await mkdtemp(join(tmpdir(), 'openclaude-plugin-agents-test-')) + setInlinePlugins([]) + clearPluginCache('loadPluginAgents.test setup') + clearPluginAgentCache() +}) + +afterEach(async () => { + try { + setInlinePlugins([]) + clearPluginCache('loadPluginAgents.test cleanup') + clearPluginAgentCache() + await rm(tempDir, { recursive: true, force: true }) + } finally { + releaseSharedMutationLock() + } +}) + +async function writePluginAgent( + pluginRoot: string, + filename: string, + frontmatter: string, +): Promise { + await mkdir(join(pluginRoot, '.claude-plugin'), { recursive: true }) + await mkdir(join(pluginRoot, 'agents'), { recursive: true }) + await writeFile( + join(pluginRoot, '.claude-plugin', 'plugin.json'), + JSON.stringify({ name: 'step-limit-plugin', version: '0.0.0' }), + ) + await writeFile( + join(pluginRoot, 'agents', filename), + `--- +name: ${filename.replace(/\.md$/, '')} +description: "Use for plugin maxSteps coverage" +${frontmatter} +--- + +Plugin agent prompt. +`, + ) +} + +describe('loadPluginAgents', () => { + test('loads valid maxSteps from plugin agent frontmatter and ignores invalid values safely', async () => { + const pluginRoot = join(tempDir, 'plugin') + await writePluginAgent(pluginRoot, 'valid.md', 'maxSteps: 7\n') + await writePluginAgent(pluginRoot, 'invalid.md', 'maxSteps: 0\n') + await writePluginAgent(pluginRoot, 'malformed.md', 'maxSteps: 2abc\n') + + setInlinePlugins([pluginRoot]) + clearPluginCache('loadPluginAgents.test inline plugin') + clearPluginAgentCache() + + const agents = await loadPluginAgents() + + const valid = agents.find( + agent => agent.agentType === 'step-limit-plugin:valid', + ) + const invalid = agents.find( + agent => agent.agentType === 'step-limit-plugin:invalid', + ) + const malformed = agents.find( + agent => agent.agentType === 'step-limit-plugin:malformed', + ) + expect(valid?.maxSteps).toBe(7) + expect(invalid).toBeDefined() + expect(invalid?.maxSteps).toBeUndefined() + expect(malformed).toBeDefined() + expect(malformed?.maxSteps).toBeUndefined() + }) + + test('loads maxSteps from plugin manifest agent file paths', async () => { + const pluginRoot = join(tempDir, 'manifest-plugin') + await mkdir(join(pluginRoot, '.claude-plugin'), { recursive: true }) + await mkdir(join(pluginRoot, 'custom-agents'), { recursive: true }) + await writeFile( + join(pluginRoot, '.claude-plugin', 'plugin.json'), + JSON.stringify({ + name: 'manifest-step-limit-plugin', + version: '0.0.0', + agents: [ + './custom-agents/valid.md', + './custom-agents/invalid.md', + './custom-agents/malformed.md', + ], + }), + ) + for (const [filename, maxSteps] of [ + ['valid.md', '9'], + ['invalid.md', '0'], + ['malformed.md', '2abc'], + ] as const) { + await writeFile( + join(pluginRoot, 'custom-agents', filename), + `--- +name: ${filename.replace(/\.md$/, '')} +description: "Use for plugin manifest maxSteps coverage" +maxSteps: ${maxSteps} +--- + +Plugin manifest agent prompt. +`, + ) + } + + setInlinePlugins([pluginRoot]) + clearPluginCache('loadPluginAgents.test manifest agents') + clearPluginAgentCache() + + const agents = await loadPluginAgents() + + const valid = agents.find( + agent => agent.agentType === 'manifest-step-limit-plugin:valid', + ) + const invalid = agents.find( + agent => agent.agentType === 'manifest-step-limit-plugin:invalid', + ) + const malformed = agents.find( + agent => agent.agentType === 'manifest-step-limit-plugin:malformed', + ) + expect(valid?.maxSteps).toBe(9) + expect(invalid).toBeDefined() + expect(invalid?.maxSteps).toBeUndefined() + expect(malformed).toBeDefined() + expect(malformed?.maxSteps).toBeUndefined() + }) +}) diff --git a/src/utils/plugins/loadPluginAgents.ts b/src/utils/plugins/loadPluginAgents.ts index d335a1234..93020b3ce 100644 --- a/src/utils/plugins/loadPluginAgents.ts +++ b/src/utils/plugins/loadPluginAgents.ts @@ -176,6 +176,15 @@ async function loadAgentFromFile( ) } + // Parse maxSteps + const maxStepsRaw = frontmatter.maxSteps + const maxSteps = parsePositiveIntFromFrontmatter(maxStepsRaw) + if (maxStepsRaw !== undefined && maxSteps === undefined) { + logForDebugging( + `Plugin agent file ${filePath} has invalid maxSteps '${maxStepsRaw}'. Must be a positive integer.`, + ) + } + // Parse disallowedTools const disallowedTools = frontmatter.disallowedTools !== undefined @@ -219,6 +228,7 @@ async function loadAgentFromFile( ...(isolation ? { isolation } : {}), ...(effort !== undefined ? { effort } : {}), ...(maxTurns !== undefined ? { maxTurns } : {}), + ...(maxSteps !== undefined ? { maxSteps } : {}), } as AgentDefinition } catch (error) { logForDebugging(`Failed to load agent from ${filePath}: ${error}`, { diff --git a/src/utils/settings/constants.ts b/src/utils/settings/constants.ts index 4a8aedad6..5e78db3da 100644 --- a/src/utils/settings/constants.ts +++ b/src/utils/settings/constants.ts @@ -44,7 +44,7 @@ export function getSettingSourceName(source: SettingSource): string { * @returns Short capitalized display name like 'User', 'Project', 'Plugin' */ export function getSourceDisplayName( - source: SettingSource | 'plugin' | 'built-in', + source: SettingSource | 'plugin' | 'built-in' | 'sdk', ): string { switch (source) { case 'userSettings': @@ -59,6 +59,8 @@ export function getSourceDisplayName( return 'Managed' case 'plugin': return 'Plugin' + case 'sdk': + return 'SDK' case 'built-in': return 'Built-in' } diff --git a/tests/sdk/package-consumer-types.test.ts b/tests/sdk/package-consumer-types.test.ts index f0680c185..1c54c4f71 100644 --- a/tests/sdk/package-consumer-types.test.ts +++ b/tests/sdk/package-consumer-types.test.ts @@ -138,6 +138,7 @@ describe('package consumer types', () => { ` SDKRateLimitError,`, ` QueryOptions,`, ` SDKSession,`, + ` SDKSessionOptions,`, `} from '@gitlawb/openclaude/sdk'`, ``, `// Use the types so they're not unused-imports-eliminated`, @@ -153,6 +154,26 @@ describe('package consumer types', () => { `// Verify session types`, `declare const session: SDKSession`, `const _messages: SDKMessage[] = session.getMessages()`, + ``, + `// SDK-provided agents can rely on the runtime name fallback for description.`, + `const _queryOptions: QueryOptions = {`, + ` cwd: '/tmp/project',`, + ` agents: {`, + ` helper: {`, + ` prompt: 'Help with package consumer type coverage',`, + ` maxSteps: 2,`, + ` },`, + ` },`, + `}`, + `const _sessionOptions: SDKSessionOptions = {`, + ` cwd: '/tmp/project',`, + ` agents: {`, + ` helper: {`, + ` prompt: 'Help with persistent SDK session type coverage',`, + ` maxSteps: 2,`, + ` },`, + ` },`, + `}`, ].join('\n'), ) diff --git a/tests/sdk/query-happy-path.test.ts b/tests/sdk/query-happy-path.test.ts index 377eb6fff..9817ee962 100644 --- a/tests/sdk/query-happy-path.test.ts +++ b/tests/sdk/query-happy-path.test.ts @@ -1,6 +1,10 @@ import { describe, test, expect, beforeAll, afterAll } from 'bun:test' import { MockQueryEngine } from './helpers/mock-engine.js' -import { query } from '../../src/entrypoints/sdk/index.js' +import { + createSdkMcpServer, + query, + tool, +} from '../../src/entrypoints/sdk/index.js' import { acquireSharedMutationLock, releaseSharedMutationLock, @@ -130,6 +134,195 @@ describe('Query happy-path — full lifecycle', () => { ) expect(textContent?.text).toContain(prompt) }) + + test('invalid SDK agent definitions are emitted before engine output', async () => { + const mockEngine = new MockQueryEngine() + const q = query({ + prompt: 'agent failure visibility', + options: { + cwd: process.cwd(), + agents: { + broken: { + description: 'Use for broken SDK agent coverage', + prompt: 2 as unknown as string, + }, + badLimit: { + description: 'Use for invalid SDK agent step limit coverage', + prompt: 'bad limit prompt', + maxSteps: 0, + }, + }, + }, + }) + ;(q as any).setEngine(mockEngine) + + const messages: any[] = [] + for await (const msg of q) { + messages.push(msg) + } + + expect(messages[0]).toMatchObject({ + type: 'agent_load_failure', + stage: 'injection', + }) + expect(messages[0].error_message).toContain("Invalid SDK agent 'broken'") + const loadFailures = messages.filter( + message => message?.type === 'agent_load_failure', + ) + expect(loadFailures).toHaveLength(2) + expect(loadFailures[1].error_message).toContain( + "Invalid SDK agent 'badLimit'", + ) + expect(loadFailures[1].error_message).toContain('maxSteps') + const assistantIndex = messages.findIndex( + message => message?.type === 'assistant', + ) + expect(assistantIndex).toBeGreaterThan(1) + expect( + loadFailures.every(failure => messages.indexOf(failure) < assistantIndex), + ).toBe(true) + expect(messages.some(message => message?.type === 'assistant')).toBe(true) + expect( + mockEngine.config.agents.some( + (agent: any) => agent?.agentType === 'badLimit', + ), + ).toBe(false) + }) + + test('valid SDK agents are exposed after successful engine injection', async () => { + const mockEngine = new MockQueryEngine() + const q = query({ + prompt: 'agent injection success', + options: { + cwd: process.cwd(), + agents: { + helper: { + description: 'Use for successful SDK agent injection coverage', + prompt: 'Help with SDK agent injection coverage', + maxSteps: 2, + }, + }, + }, + }) + ;(q as any).setEngine(mockEngine) + + const messages: unknown[] = [] + for await (const msg of q) { + messages.push(msg) + } + + expect(messages.some((message: any) => message?.type === 'assistant')).toBe( + true, + ) + expect( + mockEngine.config.agents.some( + (agent: any) => agent?.agentType === 'helper' && agent?.maxSteps === 2, + ), + ).toBe(true) + expect( + (q as any).appStateStore + .getState() + .agentDefinitions.allAgents.some( + (agent: any) => agent?.agentType === 'helper', + ), + ).toBe(true) + expect(q.supportedAgents()).toContain('helper') + }) + + test('failed SDK agent injection does not expose uninjected user agents', async () => { + const mockEngine = new MockQueryEngine() + mockEngine.injectAgents = () => { + throw new Error('injection rejected') + } + const q = query({ + prompt: 'agent injection failure', + options: { + cwd: process.cwd(), + agents: { + leaky: { + description: 'Use for failed SDK agent injection coverage', + prompt: 'Help with SDK agent injection failure coverage', + maxSteps: 2, + }, + }, + }, + }) + ;(q as any).setEngine(mockEngine) + + const messages: any[] = [] + for await (const msg of q) { + messages.push(msg) + } + + const failureIndex = messages.findIndex( + message => + message?.type === 'agent_load_failure' && + message?.stage === 'injection' && + message?.error_message === 'injection rejected', + ) + const assistantIndex = messages.findIndex( + message => message?.type === 'assistant', + ) + expect(failureIndex).toBeGreaterThanOrEqual(0) + expect(failureIndex).toBeLessThan(assistantIndex) + expect(q.supportedAgents()).not.toContain('leaky') + expect( + (q as any).appStateStore + .getState() + .agentDefinitions.allAgents.some( + (agent: any) => agent?.agentType === 'leaky', + ), + ).toBe(false) + }) + + test('denied SDK MCP tools are filtered before engine updateTools', async () => { + const mockEngine = new MockQueryEngine() + const deniedBash = tool( + 'Bash', + 'Denied MCP duplicate of Bash', + { type: 'object', properties: {} }, + async () => ({ content: [{ type: 'text', text: 'denied' }] }), + ) + const allowedSdkTool = tool( + 'sdkAllowed', + 'Allowed SDK MCP tool', + { type: 'object', properties: {} }, + async () => ({ content: [{ type: 'text', text: 'allowed' }] }), + ) + const q = query({ + prompt: 'mcp deny filtering', + options: { + cwd: process.cwd(), + disallowedTools: ['Bash'], + mcpServers: { + 'sdk-tools': createSdkMcpServer({ + type: 'sdk', + name: 'sdk-tools', + tools: [deniedBash, allowedSdkTool], + }), + }, + }, + }) + const initialToolNames = (q as any).engine.config.tools.map( + (entry: any) => entry?.name, + ) + expect(initialToolNames.length).toBeGreaterThan(0) + mockEngine.config.tools = [...(q as any).engine.config.tools] + ;(q as any).setEngine(mockEngine) + + const messages: any[] = [] + for await (const msg of q) { + messages.push(msg) + } + + const toolNames = mockEngine.config.tools.map((entry: any) => entry?.name) + for (const initialToolName of initialToolNames) { + expect(toolNames).toContain(initialToolName) + } + expect(toolNames).toContain('sdkAllowed') + expect(toolNames).not.toContain('Bash') + expect(messages.some(message => message?.type === 'assistant')).toBe(true) + }) }) describe('mcpServerStatus() reads from engine.config.mcpClients', () => { diff --git a/tests/sdk/sdk-v2-lifecycle.test.ts b/tests/sdk/sdk-v2-lifecycle.test.ts index e952cbe05..662f97310 100644 --- a/tests/sdk/sdk-v2-lifecycle.test.ts +++ b/tests/sdk/sdk-v2-lifecycle.test.ts @@ -1,7 +1,10 @@ import { describe, test, expect, beforeEach, afterEach, beforeAll, afterAll } from 'bun:test' import { randomUUID } from 'crypto' -import { rmSync } from 'fs' +import { mkdirSync, rmSync, writeFileSync } from 'fs' +import { join } from 'path' import { + createSdkMcpServer, + tool, unstable_v2_createSession, unstable_v2_resumeSession, unstable_v2_prompt, @@ -30,6 +33,7 @@ import { isExpectedDrainAbort, UUID_REGEX, } from './helpers/query-test-doubles.js' +import { MockQueryEngine } from './helpers/mock-engine.js' // sendMessage drains trigger init(), which checks auth. Stub it for CI. const AUTH_KEY = 'ANTHROPIC_API_KEY' @@ -46,6 +50,10 @@ let originalOriginalCwd: string // Collect temp dirs for cleanup const tempDirs: string[] = [] +function attachMockEngine(session: unknown, mockEngine: MockQueryEngine): void { + ;(session as { setEngine(engine: MockQueryEngine): void }).setEngine(mockEngine) +} + beforeAll(async () => { await acquireSharedMutationLock('sdk-v2-lifecycle') savedApiKey = process.env[AUTH_KEY] @@ -261,6 +269,375 @@ describe('V2: permission handling', () => { }) }) +describe('V2: SDK agents', () => { + test('createSession() injects SDK agents with maxSteps on first message', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const mockEngine = new MockQueryEngine() + const session = unstable_v2_createSession({ + cwd: dir, + agents: { + helper: { + prompt: 'Help with persistent SDK agent injection coverage', + maxSteps: 2, + }, + }, + }) + attachMockEngine(session, mockEngine) + + const messages: unknown[] = [] + for await (const msg of session.sendMessage('agent injection success')) { + messages.push(msg) + } + + expect( + messages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + expect( + mockEngine.config.agents.some( + (agent: any) => + agent?.agentType === 'helper' && + agent?.whenToUse === 'helper' && + agent?.maxSteps === 2, + ), + ).toBe(true) + + const firstTurnAgents = [...mockEngine.config.agents] + const secondTurnMessages: unknown[] = [] + for await (const msg of session.sendMessage('agent injection second turn')) { + secondTurnMessages.push(msg) + } + expect( + secondTurnMessages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + expect(mockEngine.config.agents.map((agent: any) => agent?.agentType)).toEqual( + firstTurnAgents.map((agent: any) => agent?.agentType), + ) + expect( + mockEngine.config.agents.filter( + (agent: any) => agent?.agentType === 'helper', + ), + ).toHaveLength(1) + }) + }) + + test('createSession() filters denied SDK MCP tools on every turn', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const mockEngine = new MockQueryEngine() + const deniedBash = tool( + 'Bash', + 'Denied persistent SDK MCP Bash duplicate', + { type: 'object', properties: {} }, + async () => ({ content: [{ type: 'text', text: 'denied' }] }), + ) + const allowedSdkTool = tool( + 'sdkAllowed', + 'Allowed persistent SDK MCP tool', + { type: 'object', properties: {} }, + async () => ({ content: [{ type: 'text', text: 'allowed' }] }), + ) + const session = unstable_v2_createSession({ + cwd: dir, + disallowedTools: ['Bash'], + mcpServers: { + 'sdk-tools': createSdkMcpServer({ + type: 'sdk', + name: 'sdk-tools', + tools: [deniedBash, allowedSdkTool], + }), + }, + }) + const initialTools = + (session as unknown as { _engine: { config: { tools: unknown[] } } }) + ._engine.config.tools + const initialToolNames = initialTools.map((entry: any) => entry?.name) + expect(initialToolNames.length).toBeGreaterThan(0) + expect(initialToolNames).not.toContain('Bash') + mockEngine.config.tools = [...initialTools] + attachMockEngine(session, mockEngine) + + const firstTurnMessages: unknown[] = [] + for await (const msg of session.sendMessage('mcp deny filtering')) { + firstTurnMessages.push(msg) + } + const firstTurnToolNames = mockEngine.config.tools.map( + (entry: any) => entry?.name, + ) + for (const initialToolName of initialToolNames) { + expect(firstTurnToolNames).toContain(initialToolName) + } + expect(firstTurnToolNames).toContain('sdkAllowed') + expect(firstTurnToolNames).not.toContain('Bash') + expect( + firstTurnMessages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + + const secondTurnMessages: unknown[] = [] + for await (const msg of session.sendMessage('mcp deny filtering second turn')) { + secondTurnMessages.push(msg) + } + const secondTurnToolNames = mockEngine.config.tools.map( + (entry: any) => entry?.name, + ) + expect(secondTurnToolNames).toEqual(firstTurnToolNames) + expect(secondTurnToolNames).toContain('sdkAllowed') + expect(secondTurnToolNames).not.toContain('Bash') + expect( + secondTurnMessages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + }) + }) + + test('createSession() merges filesystem and SDK agents before injection', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const agentsDir = join(dir, '.openclaude', 'agents') + mkdirSync(agentsDir, { recursive: true }) + writeFileSync( + join(agentsDir, 'filesystem.md'), + [ + '---', + 'name: filesystem', + 'description: Use for filesystem agent merge coverage', + '---', + 'Filesystem agent prompt', + ].join('\n'), + ) + + const mockEngine = new MockQueryEngine() + const session = unstable_v2_createSession({ + cwd: dir, + agents: { + helper: { + description: 'Use for persistent SDK agent merge coverage', + prompt: 'Help with persistent SDK agent merge coverage', + maxSteps: 2, + }, + }, + }) + attachMockEngine(session, mockEngine) + + const messages: unknown[] = [] + for await (const msg of session.sendMessage('agent merge success')) { + messages.push(msg) + } + + expect(messages.some((message: any) => message?.type === 'assistant')).toBe( + true, + ) + const agentTypes = mockEngine.config.agents.map( + (agent: any) => agent?.agentType, + ) + expect(agentTypes).toContain('filesystem') + expect(agentTypes).toContain('helper') + + const secondTurnMessages: unknown[] = [] + for await (const msg of session.sendMessage('agent merge second turn')) { + secondTurnMessages.push(msg) + } + const secondTurnAgentTypes = mockEngine.config.agents.map( + (agent: any) => agent?.agentType, + ) + expect( + secondTurnMessages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + expect(secondTurnAgentTypes).toEqual(agentTypes) + expect( + secondTurnAgentTypes.filter(agentType => agentType === 'filesystem'), + ).toHaveLength(1) + expect(secondTurnAgentTypes.filter(agentType => agentType === 'helper')).toHaveLength(1) + }) + }) + + test('createSession() lets SDK agents override filesystem agents with the same name', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const agentsDir = join(dir, '.openclaude', 'agents') + mkdirSync(agentsDir, { recursive: true }) + writeFileSync( + join(agentsDir, 'helper.md'), + [ + '---', + 'name: helper', + 'description: Use for filesystem collision coverage', + '---', + 'Filesystem helper prompt', + ].join('\n'), + ) + + const mockEngine = new MockQueryEngine() + const session = unstable_v2_createSession({ + cwd: dir, + agents: { + helper: { + description: 'Use for SDK collision coverage', + prompt: 'SDK helper prompt', + maxSteps: 2, + }, + }, + }) + attachMockEngine(session, mockEngine) + + const messages: unknown[] = [] + for await (const msg of session.sendMessage('agent collision success')) { + messages.push(msg) + } + + const helperAgents = mockEngine.config.agents.filter( + (agent: any) => agent?.agentType === 'helper', + ) + expect(messages.some((message: any) => message?.type === 'assistant')).toBe( + true, + ) + expect(helperAgents).toHaveLength(1) + expect((helperAgents[0] as any).getSystemPrompt()).toBe( + 'SDK helper prompt', + ) + expect((helperAgents[0] as any).maxSteps).toBe(2) + + const secondTurnMessages: unknown[] = [] + for await (const msg of session.sendMessage('agent collision second turn')) { + secondTurnMessages.push(msg) + } + const secondTurnHelpers = mockEngine.config.agents.filter( + (agent: any) => agent?.agentType === 'helper', + ) + expect( + secondTurnMessages.some((message: any) => message?.type === 'assistant'), + ).toBe(true) + expect(secondTurnHelpers).toHaveLength(1) + expect((secondTurnHelpers[0] as any).getSystemPrompt()).toBe( + 'SDK helper prompt', + ) + expect((secondTurnHelpers[0] as any).maxSteps).toBe(2) + }) + }) + + test('createSession() emits invalid SDK agent failures before engine output', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const mockEngine = new MockQueryEngine() + const session = unstable_v2_createSession({ + cwd: dir, + agents: { + broken: { + description: 'Use for broken persistent SDK agent coverage', + prompt: 2 as unknown as string, + }, + badLimit: { + description: 'Use for invalid persistent SDK agent step limit coverage', + prompt: 'bad limit prompt', + maxSteps: 0, + }, + }, + }) + attachMockEngine(session, mockEngine) + + const messages: any[] = [] + for await (const msg of session.sendMessage('agent failure visibility')) { + messages.push(msg) + } + + expect(messages[0]).toMatchObject({ + type: 'agent_load_failure', + stage: 'injection', + }) + expect(messages[0].error_message).toContain("Invalid SDK agent 'broken'") + const loadFailures = messages.filter( + message => message?.type === 'agent_load_failure', + ) + expect(loadFailures).toHaveLength(2) + expect(loadFailures[1].error_message).toContain( + "Invalid SDK agent 'badLimit'", + ) + expect(loadFailures[1].error_message).toContain('maxSteps') + const assistantIndex = messages.findIndex( + message => message?.type === 'assistant', + ) + expect(assistantIndex).toBeGreaterThan(1) + expect( + loadFailures.every(failure => messages.indexOf(failure) < assistantIndex), + ).toBe(true) + expect(messages.some(message => message?.type === 'assistant')).toBe(true) + expect( + mockEngine.config.agents.some( + (agent: any) => agent?.agentType === 'broken', + ), + ).toBe(false) + expect( + mockEngine.config.agents.some( + (agent: any) => agent?.agentType === 'badLimit', + ), + ).toBe(false) + + const secondTurnMessages: any[] = [] + for await (const msg of session.sendMessage('agent failure second turn')) { + secondTurnMessages.push(msg) + } + + expect( + secondTurnMessages.some( + message => message?.type === 'agent_load_failure', + ), + ).toBe(false) + expect(secondTurnMessages.some(message => message?.type === 'assistant')).toBe( + true, + ) + }) + }) + + test('createSession() emits filesystem agent parse failures before engine output', async () => { + await withTempDir(async (dir) => { + tempDirs.push(dir) + const agentsDir = join(dir, '.openclaude', 'agents') + mkdirSync(agentsDir, { recursive: true }) + writeFileSync( + join(agentsDir, 'broken.md'), + [ + '---', + 'name: broken', + '---', + 'Broken filesystem agent prompt', + ].join('\n'), + ) + + const mockEngine = new MockQueryEngine() + const session = unstable_v2_createSession({ cwd: dir }) + attachMockEngine(session, mockEngine) + + const messages: any[] = [] + for await (const msg of session.sendMessage('agent parse failure visibility')) { + messages.push(msg) + } + + expect(messages[0]).toMatchObject({ + type: 'agent_load_failure', + stage: 'definitions', + }) + expect(messages[0].error_message).toContain('broken.md') + expect(messages[0].error_message).toContain( + 'Missing required "description" field', + ) + expect(messages.some(message => message?.type === 'assistant')).toBe(true) + + const secondTurnMessages: any[] = [] + for await (const msg of session.sendMessage('agent parse second turn')) { + secondTurnMessages.push(msg) + } + + expect( + secondTurnMessages.some( + message => message?.type === 'agent_load_failure', + ), + ).toBe(false) + expect(secondTurnMessages.some(message => message?.type === 'assistant')).toBe( + true, + ) + }) + }) +}) + describe('V2: unstable_v2_prompt', () => { test('throws when query completes without a result message (aborted)', async () => { const ac = new AbortController()