diff --git a/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server-subagent-watchdog.test.ts b/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server-subagent-watchdog.test.ts index 4b01cb8f2..abc04449d 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server-subagent-watchdog.test.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server-subagent-watchdog.test.ts @@ -443,6 +443,79 @@ describe('OpenCode subagent run tracking', () => { } }); + it.each(['cancel', 'dispose'])( + 'cleans up spawn safety deadlines on %s', + async (action) => { + vi.useFakeTimers(); + const { client, harness } = createHarness(); + try { + await connectHarness(harness, client); + await armSpawn(client, harness); + if (action === 'dispose') harness.dispose(); + else + harness.sendCommand({ + commandName: TaskCommandName.CancelTask, + data: {}, + }); + await vi.advanceTimersByTimeAsync(0); + const aborts = client.abort.mock.calls.length; + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs + 1_000); + expect(client.abort).toHaveBeenCalledTimes(aborts); + } finally { + harness.dispose(); + } + }, + ); + + it('starts a fresh safety pass when a new task call resumes a settled child', async () => { + const { client, harness } = createHarness(); + try { + await connectHarness(harness, client); + await armSpawn(client, harness); + const completeTool = (id: number) => + client.emit({ + type: 'message.part.updated', + properties: { + part: { + ...createChildToolPart({ + callId: `grep-${id}`, + tool: 'grep', + status: 'completed', + }), + state: { status: 'completed', input: {}, output: '' }, + }, + }, + }); + for (let id = 0; id < 4; id++) await completeTool(id); + await client.emit({ + type: 'message.part.updated', + properties: { + part: createTaskToolPart({ + status: 'completed', + metadata: { sessionId: 'ses_child_1' }, + }), + }, + }); + await client.emit({ + type: 'message.part.updated', + properties: { + part: createTaskToolPart({ + callId: 'resumed-call', + metadata: { sessionId: 'ses_child_1' }, + }), + }, + }); + for (let id = 0; id < 4; id++) await completeTool(id); + expect(client.abort).not.toHaveBeenCalled(); + await completeTool(4); + expect(client.abort).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: 'ses_child_1' }), + ); + } finally { + harness.dispose(); + } + }); + it('does not extend an older pass deadline when a sibling starts or child progress arrives', async () => { vi.useFakeTimers(); const { client, harness } = createHarness(); diff --git a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts index 2b0e7a783..4e782351f 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts @@ -8,8 +8,7 @@ import { spawnSync } from 'node:child_process'; import { asBoolean, AGENT_SAFETY_STOP_TEXT, - AgentSafetyPass, - AGENT_SAFETY, + AgentSafetyLifecycle, isAgentSafetyStop, type AgentSafetyReason, type AgentSafetyStopDetails, @@ -1650,13 +1649,12 @@ export class OpenCodeServerHarness implements Harness { private readonly client: OpenCodeServerClient; - private readonly safetyPasses = new Map(); + private readonly safety: AgentSafetyLifecycle; private readonly onSafetyStop: OpenCodeServerHarnessOptions['onSafetyStop']; - private readonly safetyTimer: ReturnType; - private safetyStopped = false; + private get safetyStopped(): boolean { + return this.safety.stopped; + } private readonly safetyStoppedSessionIds = new Set(); - private readonly settledSafetySessions = new Set(); - private readonly settledSubagentWatchdogKeys = new Set(); private readonly workspacePath: string; private readonly normalizedWorkspacePath: string; private readonly logger: OpenCodeServerHarnessOptions['logger']; @@ -1854,24 +1852,14 @@ export class OpenCodeServerHarness this.fallbackRole = options.fallbackRole; this.agentFallbacks = options.agentFallbacks ?? {}; this.onSafetyStop = options.onSafetyStop; - this.safetyTimer = setInterval(() => { - for (const watchdog of this.activeSubagentWatchdogs.values()) { - if (Date.now() - watchdog.startedAtMs >= AGENT_SAFETY.passTimeoutMs) { - void this.stopForSafety( - 'pass_time_limit', - watchdog.childSessionId ?? watchdog.sessionId, - ); - return; - } - } - for (const [id, pass] of this.safetyPasses) { - if (id !== this.sessionId && pass.checkTime()) { - void this.stopForSafety('pass_time_limit', id); - break; - } - } - }, 1_000); - this.safetyTimer.unref?.(); + this.safety = new AgentSafetyLifecycle({ + digest: (value) => createHash('sha256').update(value).digest('hex'), + isBounded: (id) => id !== this.sessionId, + includeFailedTools: true, + onDeadline: (id) => { + void this.stopForSafety('pass_time_limit', id); + }, + }); this.commandEnv = options.commandEnv ? { ...options.commandEnv } : undefined; @@ -2098,8 +2086,7 @@ export class OpenCodeServerHarness } dispose(): void { - clearInterval(this.safetyTimer); - this.safetyPasses.clear(); + this.safety.dispose(); this.disposed = true; this.connected = false; this.clearReplayAbortErrorSuppression(); @@ -2758,10 +2745,10 @@ export class OpenCodeServerHarness updatePayload: Record; }, ): void { - if (this.settledSubagentWatchdogKeys.has(eventKey)) return; - if (input.childSessionId) - this.settledSafetySessions.delete(input.childSessionId); - if (input.childSessionId) this.safetyPass(input.childSessionId); + if ( + !this.safety.startSpawn(eventKey, input.sessionId, input.childSessionId) + ) + return; const existing = this.activeSubagentWatchdogs.get(eventKey); if (existing) { @@ -2895,7 +2882,7 @@ export class OpenCodeServerHarness } private stopSubagentWatchdog(eventKey: string): void { - this.settledSubagentWatchdogKeys.add(eventKey); + this.safety.settleSpawn(eventKey); const watchdog = this.activeSubagentWatchdogs.get(eventKey); if (!watchdog) { @@ -2907,8 +2894,6 @@ export class OpenCodeServerHarness } this.clearSubagentActivityFlush(watchdog); if (watchdog.childSessionId) { - this.settledSafetySessions.add(watchdog.childSessionId); - this.safetyPasses.delete(watchdog.childSessionId); this.childSessionWatchdogKeys.delete(watchdog.childSessionId); this.linkedChildSessions.delete(watchdog.childSessionId); } @@ -2937,18 +2922,7 @@ export class OpenCodeServerHarness continue; } - if (watchdog.settlementTimer) { - clearTimeout(watchdog.settlementTimer); - } - this.clearSubagentActivityFlush(watchdog); - this.settledSubagentWatchdogKeys.add(eventKey); - this.activeSubagentWatchdogs.delete(eventKey); - if (watchdog.childSessionId) { - this.settledSafetySessions.add(watchdog.childSessionId); - this.safetyPasses.delete(watchdog.childSessionId); - this.childSessionWatchdogKeys.delete(watchdog.childSessionId); - this.linkedChildSessions.delete(watchdog.childSessionId); - } + this.stopSubagentWatchdog(eventKey); } } @@ -4144,11 +4118,7 @@ export class OpenCodeServerHarness parentSessionId: parentId, agentType: this.linkedChildSessions.get(childId)?.agentType ?? null, }); - if ( - payload.type === 'session.created' && - !this.settledSafetySessions.has(childId) - ) - this.safetyPass(childId); + if (payload.type === 'session.created') this.safety.start(childId); } } @@ -4173,46 +4143,7 @@ export class OpenCodeServerHarness await this.stopForSafety('task_request_limit', sessionId); return; } - const pass = - this.safetyPasses.get(sessionId) ?? - (!this.settledSafetySessions.has(sessionId) && - (['message.part.updated', 'message.updated'].includes(payload.type) || - ['busy', 'retry'].includes( - asString(asRecord(properties?.status)?.type) ?? '', - )) - ? this.safetyPass(sessionId) - : undefined); - const part = asRecord(properties?.part); - let reason: AgentSafetyReason | undefined; - if (payload.type === 'message.part.updated' && part?.type === 'tool') { - const state = asRecord(part.state); - if ( - ['completed', 'error', 'failed'].includes( - asString(state?.status) ?? '', - ) && - asString(part.callID) - ) { - reason = pass?.completed({ - id: String(part.callID), - name: asString(part.tool) ?? '', - args: state?.input, - result: { - status: state?.status, - output: state?.output, - error: state?.error, - }, - }); - } - } - // A model step is an assistant message, not every streaming update. - const info = asRecord(properties?.info); - if ( - sessionId !== this.sessionId && - payload.type === 'message.updated' && - info?.role === 'assistant' && - asString(info.id) - ) - reason ??= pass?.step(String(info.id)); + const reason = this.safety.observe(sessionId, payload); if (reason) { await this.stopForSafety(reason, sessionId); return; @@ -4222,8 +4153,7 @@ export class OpenCodeServerHarness payload.type === 'session.idle' && !this.childSessionWatchdogKeys.has(sessionId) ) { - this.safetyPasses.delete(sessionId); - this.settledSafetySessions.add(sessionId); + this.safety.settle(sessionId); } } @@ -4320,25 +4250,14 @@ export class OpenCodeServerHarness } } - private safetyPass(sessionId: string): AgentSafetyPass { - let pass = this.safetyPasses.get(sessionId); - if (!pass) { - pass = new AgentSafetyPass((value) => - createHash('sha256').update(value).digest('hex'), - ); - this.safetyPasses.set(sessionId, pass); - } - return pass; - } - /** Stop the entire tree: a parent cannot retry or replace a safety-stopped pass. */ private async stopForSafety( reason: AgentSafetyReason, stoppedSessionId: string, ): Promise { - if (this.safetyStopped || this.disposed) return; - this.safetyStopped = true; - clearInterval(this.safetyTimer); + if (this.disposed) return; + const details = this.safety.claimStop(stoppedSessionId); + if (!details) return; const sessionId = this.sessionId; const ids = new Set([ stoppedSessionId, @@ -4346,10 +4265,6 @@ export class OpenCodeServerHarness ...(sessionId ? [sessionId] : []), ]); for (const id of ids) this.safetyStoppedSessionIds.add(id); - const details = { - sessionId: stoppedSessionId, - ...this.safetyPasses.get(stoppedSessionId)?.snapshot(), - }; this.prompts.clear(); this.clearQueuedPromptRetryTimer(); this.clearProviderErrorRecoveryState(); @@ -5981,7 +5896,7 @@ export class OpenCodeServerHarness } if (this.safetyStopped) return; - this.safetyPasses.delete(sessionId); + this.safety.finishRoot(sessionId); this.runtimeEvents.taskCompleted(sessionId, finalized?.tokenUsage, { ...(missingChatCloseoutReminderCount !== null ? { diff --git a/packages/cloud-agents/src/server/__tests__/non-task-provider-usage.test.ts b/packages/cloud-agents/src/server/__tests__/non-task-provider-usage.test.ts index fa76a3207..cc2c4ec7c 100644 --- a/packages/cloud-agents/src/server/__tests__/non-task-provider-usage.test.ts +++ b/packages/cloud-agents/src/server/__tests__/non-task-provider-usage.test.ts @@ -215,82 +215,106 @@ describe('resolveOpenCodeSmallModel', () => { expect(resolveOpenCodeSmallModel()).toBe('openrouter/openai/gpt-4.1-mini'); }); - it('stops identical empty child searches, cancels the tree, persists the task stop, and does not retry', async () => { - const { stopTaskForAgentSafety } = await import('@roomote/db/server'); - mockResolveEffectiveModelRuntimeEnv.mockResolvedValue({ - R_MODEL: 'openrouter/openai/gpt-5.4', - OPENCODE_CONFIG_CONTENT: '', - }); - sessionPromptMock.mockImplementation(() => new Promise(() => {})); - eventSubscribeMock.mockImplementation(async () => ({ - stream: (async function* () { - yield { - type: 'session.created', - properties: { - sessionID: 'child', - info: { id: 'child', parentID: 'session-1' }, - }, - }; - yield { - type: 'session.created', - properties: { - sessionID: 'grandchild', - info: { id: 'grandchild', parentID: 'child' }, - }, - }; - for (let index = 0; index < 1_000; index++) { + it.each(['ordinary', 'resumed child', 'failed persistence and abort'])( + 'stops identical empty child searches and cancels the tree without retry (%s)', + async (scenario) => { + const { stopTaskForAgentSafety } = await import('@roomote/db/server'); + if (scenario === 'failed persistence and abort') { + vi.mocked(stopTaskForAgentSafety).mockRejectedValueOnce( + new Error('persistence unavailable'), + ); + sessionAbortMock.mockRejectedValue(new Error('abort unavailable')); + } + mockResolveEffectiveModelRuntimeEnv.mockResolvedValue({ + R_MODEL: 'openrouter/openai/gpt-5.4', + OPENCODE_CONFIG_CONTENT: '', + }); + sessionPromptMock.mockImplementation(() => new Promise(() => {})); + eventSubscribeMock.mockImplementation(async () => ({ + stream: (async function* () { yield { - type: 'message.part.updated', + type: 'session.created', properties: { - part: { - type: 'tool', - id: `part-${index}`, - callID: `call-${index}`, - sessionID: 'child', - messageID: 'message', - tool: 'grep', - state: { - status: 'completed', - input: { pattern: 'missing' }, - output: '', - }, - }, + sessionID: 'child', + info: { id: 'child', parentID: 'session-1' }, }, }; - } - })(), - })); - const { generateTrackedNonTaskText } = - await import('../non-task-provider-usage'); - await expect( - generateTrackedNonTaskText({ - surface: 'test', - taskId: 'task-1', - prompt: 'Do work', - }), - ).rejects.toThrow('Roomote stopped this task at a safety limit'); - expect(stopTaskForAgentSafety).toHaveBeenCalledExactlyOnceWith( - expect.objectContaining({ - taskId: 'task-1', - reason: 'identical_tool_loop', - details: expect.objectContaining({ - sessionId: 'child', - completedTools: 5, + yield { + type: 'session.created', + properties: { + sessionID: 'grandchild', + info: { id: 'grandchild', parentID: 'child' }, + }, + }; + for (let index = 0; index < 1_000; index++) { + if (scenario === 'resumed child' && index === 4) { + yield { + type: 'session.idle', + properties: { sessionID: 'child' }, + }; + yield { + type: 'session.updated', + properties: { + sessionID: 'child', + info: { id: 'child', parentID: 'session-1' }, + }, + }; + } + const event = { + type: 'message.part.updated', + properties: { + part: { + type: 'tool', + id: `part-${index}`, + callID: `call-${index}`, + sessionID: 'child', + messageID: 'message', + tool: 'grep', + state: { + status: 'completed', + input: { pattern: 'missing' }, + output: '', + }, + }, + }, + }; + yield event; + yield event; // Replay must not count as another completion. + } + })(), + })); + const { generateTrackedNonTaskText } = + await import('../non-task-provider-usage'); + await expect( + generateTrackedNonTaskText({ + surface: 'test', + taskId: 'task-1', + prompt: 'Do work', }), - }), - ); - expect( - new Set(sessionAbortMock.mock.calls.map(([input]) => input.sessionID)), - ).toEqual(new Set(['session-1', 'child', 'grandchild'])); - expect(sessionPromptMock).toHaveBeenCalledTimes(1); - const config = JSON.parse( - spawnMock.mock.calls.at(-1)?.[2]?.env.OPENCODE_CONFIG_CONTENT, - ); - expect(config.provider.openrouter.options.apiKey).toBe('test-run-token'); - expect(config.provider.openrouter.options.baseURL).toContain( - '/api/inference/openrouter', - ); - }); + ).rejects.toThrow('Roomote stopped this task at a safety limit'); + expect(stopTaskForAgentSafety).toHaveBeenCalledExactlyOnceWith( + expect.objectContaining({ + taskId: 'task-1', + reason: 'identical_tool_loop', + details: expect.objectContaining({ + sessionId: 'child', + completedTools: 5, + }), + }), + ); + expect( + new Set(sessionAbortMock.mock.calls.map(([input]) => input.sessionID)), + ).toEqual(new Set(['session-1', 'child', 'grandchild'])); + expect(sessionPromptMock).toHaveBeenCalledTimes(1); + const config = JSON.parse( + spawnMock.mock.calls.at(-1)?.[2]?.env.OPENCODE_CONFIG_CONTENT, + ); + expect(config.provider.openrouter.options.apiKey).toBe('test-run-token'); + expect(config.provider.openrouter.options.baseURL).toContain( + '/api/inference/openrouter', + ); + }, + ); it('falls back to the deployment task model', () => { process.env = { diff --git a/packages/cloud-agents/src/server/non-task-provider-usage.ts b/packages/cloud-agents/src/server/non-task-provider-usage.ts index c3b32881c..b9105303d 100644 --- a/packages/cloud-agents/src/server/non-task-provider-usage.ts +++ b/packages/cloud-agents/src/server/non-task-provider-usage.ts @@ -24,7 +24,7 @@ import { formatErrorForLog, AGENT_SAFETY, AGENT_SAFETY_STOP_TEXT, - AgentSafetyPass, + AgentSafetyLifecycle, isAgentSafetyStop, buildInferenceGatewayUrl, type AgentSafetyReason, @@ -1878,33 +1878,27 @@ async function runNonTaskSdkPrompt( await options.onSessionReady?.(sessionId); const trackedSessionIds = new Set([sessionId]); - const safetyPasses = new Map(); - const safetyPass = (id: string) => { - let pass = safetyPasses.get(id); - if (!pass) { - pass = new AgentSafetyPass((value) => - createHash('sha256').update(value).digest('hex'), - ); - safetyPasses.set(id, pass); - } - return pass; - }; - safetyPass(sessionId); - let safetyStopped = false; + const safety = new AgentSafetyLifecycle({ + digest: (value) => createHash('sha256').update(value).digest('hex'), + isBounded: () => true, + // Native sessions can continue a child without a worker task-tool spawn. + resumeOnActivity: true, + onDeadline: (id) => { + void safetyStop('pass_time_limit', id).catch(() => undefined); + }, + }); + safety.start(sessionId); const safetyStop = async ( reason: AgentSafetyReason, stoppedSessionId = sessionId, ) => { - if (safetyStopped) return; - safetyStopped = true; + const details = safety.claimStop(stoppedSessionId); + if (!details) return; safetyPersistence = params.taskId ? stopTaskForAgentSafety({ taskId: params.taskId, reason, - details: { - sessionId: stoppedSessionId, - ...safetyPasses.get(stoppedSessionId)?.snapshot(), - }, + details, }) : Promise.resolve(); // Reject and abort locally immediately, then persist the durable latch. @@ -1913,15 +1907,6 @@ async function runNonTaskSdkPrompt( abortController.abort(error); await safetyPersistence; }; - const safetyTimer = setInterval(() => { - for (const [id, pass] of safetyPasses) { - if (pass.checkTime()) { - void safetyStop('pass_time_limit', id).catch(() => undefined); - break; - } - } - }, 1_000); - safetyTimer.unref(); const sessionsCreatedThisTurn = new Set(); const usageRecordings = new Map>(); const observedUsageEventKeys = new Set(); @@ -2067,7 +2052,7 @@ async function runNonTaskSdkPrompt( const childId = event.properties.info.id ?? event.properties.sessionID; trackedSessionIds.add(childId); - safetyPass(childId); + safety.start(childId, true); if (event.type === 'session.created') { sessionsCreatedThisTurn.add(event.properties.sessionID); } @@ -2083,7 +2068,7 @@ async function runNonTaskSdkPrompt( event.type === 'session.idle' && event.properties.sessionID !== sessionId ) { - safetyPasses.delete(event.properties.sessionID); + safety.settle(event.properties.sessionID); } else if (event.type === 'message.part.updated') { const part = event.properties.part; if ( @@ -2091,22 +2076,7 @@ async function runNonTaskSdkPrompt( trackedSessionIds.has(part.sessionID) && ['completed', 'error'].includes(part.state.status) ) { - const state = part.state as { - status: string; - input?: unknown; - output?: unknown; - error?: unknown; - }; - const reason = safetyPass(part.sessionID).completed({ - id: part.callID, - name: part.tool, - args: state.input, - result: { - status: state.status, - output: state.output, - error: state.error, - }, - }); + const reason = safety.observe(part.sessionID, event); if (reason) { await safetyStop(reason, part.sessionID); return; @@ -2160,7 +2130,7 @@ async function runNonTaskSdkPrompt( await safetyStop('task_request_limit'); return; } - const reason = safetyPass(info.sessionID).step(info.id); + const reason = safety.observe(info.sessionID, event); if (reason) { await safetyStop(reason, info.sessionID); return; @@ -2530,7 +2500,7 @@ async function runNonTaskSdkPrompt( } throw error; } finally { - clearInterval(safetyTimer); + safety.dispose(); abortController.signal.removeEventListener('abort', abortEventMonitor); eventAbortController.abort(); void eventMonitor?.catch(() => undefined); diff --git a/packages/types/src/__tests__/agent-safety-lifecycle.test.ts b/packages/types/src/__tests__/agent-safety-lifecycle.test.ts new file mode 100644 index 000000000..cc40a372d --- /dev/null +++ b/packages/types/src/__tests__/agent-safety-lifecycle.test.ts @@ -0,0 +1,158 @@ +import { createHash } from 'node:crypto'; +import { AGENT_SAFETY, AgentSafetyLifecycle } from '../agent-safety'; + +const toolEvent = (id: number, status = 'completed') => ({ + type: 'message.part.updated', + properties: { + part: { + type: 'tool', + callID: String(id), + tool: 'grep', + state: { status, input: { pattern: 'missing' }, output: '' }, + }, + }, +}); +const stepEvent = (id: number) => ({ + type: 'message.updated', + properties: { info: { role: 'assistant', id: String(id) } }, +}); + +describe('agent safety lifecycle', () => { + const owners: AgentSafetyLifecycle[] = []; + function create( + options: Partial< + ConstructorParameters[0] + > = {}, + ) { + const onDeadline = vi.fn(); + const owner = new AgentSafetyLifecycle({ + digest: (text) => createHash('sha256').update(text).digest('hex'), + isBounded: (id) => id !== 'root', + onDeadline, + ...options, + }); + owners.push(owner); + return { owner, onDeadline }; + } + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + for (const owner of owners.splice(0)) owner.dispose(); + vi.useRealTimers(); + }); + + it('deduplicates completions per session and claims a tree stop only once', () => { + const { owner } = create(); + for (let id = 0; id < 4; id++) { + for (let duplicate = 0; duplicate < 10; duplicate++) { + expect(owner.observe('child', toolEvent(id))).toBeUndefined(); + expect(owner.observe('sibling', toolEvent(id))).toBeUndefined(); + } + } + expect(owner.observe('child', toolEvent(4))).toBe('identical_tool_loop'); + expect(owner.claimStop('child')).toMatchObject({ + sessionId: 'child', + completedTools: 5, + }); + expect(owner.claimStop('sibling')).toBeUndefined(); + expect(owner.startSpawn('new', 'root', 'new-child')).toBe(false); + expect(owner.observe('sibling', toolEvent(4))).toBeUndefined(); + expect(owner.stopped).toBe(true); + }); + + it('keeps the worker root exempt from step and wall limits but counts its tools', async () => { + const { owner, onDeadline } = create(); + for (let id = 0; id < AGENT_SAFETY.passSteps + 1; id++) + expect(owner.observe('root', stepEvent(id))).toBeUndefined(); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(onDeadline).not.toHaveBeenCalled(); + for (let id = 0; id < 5; id++) + expect(owner.observe('root', toolEvent(id))).toBe( + id === 4 ? 'identical_tool_loop' : undefined, + ); + }); + + it('does not reset an older pending spawn deadline when child metadata or siblings arrive', async () => { + const { owner, onDeadline } = create(); + owner.startSpawn('spawn', 'root', null); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs / 2); + owner.startSpawn('spawn', 'root', 'child'); + owner.startSpawn('sibling-spawn', 'root', 'sibling'); + owner.observe('child', toolEvent(0)); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs / 2); + expect(onDeadline).toHaveBeenCalledExactlyOnceWith('child'); + }); + + it('bounds pending spawns even before a child session exists', async () => { + const { owner, onDeadline } = create(); + owner.startSpawn('pending', 'root', null); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(onDeadline).toHaveBeenCalledExactlyOnceWith('root'); + }); + + it('keeps a session-only deadline when a sibling or more message updates arrive', async () => { + const { owner, onDeadline } = create({ isBounded: () => true }); + owner.start('root'); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs / 2); + owner.start('child'); + owner.observe('root', stepEvent(0)); + owner.observe('child', toolEvent(0)); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs / 2); + expect(onDeadline).toHaveBeenCalledExactlyOnceWith('root'); + }); + + it('ignores settled child events and duplicate spawns but permits an explicit resumed pass', async () => { + const { owner, onDeadline } = create(); + owner.startSpawn('first', 'root', 'child'); + for (let id = 0; id < 4; id++) owner.observe('child', toolEvent(id)); + owner.settleSpawn('first'); + expect(owner.startSpawn('first', 'root', 'child')).toBe(false); + owner.start('child'); // late session.created + for (let id = 4; id < 20; id++) + expect(owner.observe('child', toolEvent(id))).toBeUndefined(); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(onDeadline).not.toHaveBeenCalled(); + expect(owner.startSpawn('resumed', 'root', 'child')).toBe(true); + for (let id = 0; id < 5; id++) + expect(owner.observe('child', toolEvent(id))).toBe( + id === 4 ? 'identical_tool_loop' : undefined, + ); + }); + + it('allows native activity to resume a settled pass and bounds the native root', () => { + const { owner } = create({ isBounded: () => true, resumeOnActivity: true }); + for (let id = 0; id < 4; id++) owner.observe('child', toolEvent(id)); + owner.settle('child'); + for (let id = 0; id < 5; id++) + expect(owner.observe('child', toolEvent(id))).toBe( + id === 4 ? 'identical_tool_loop' : undefined, + ); + for (let id = 0; id < AGENT_SAFETY.passSteps; id++) { + const reason = + id === AGENT_SAFETY.passSteps - 1 ? 'pass_step_limit' : undefined; + expect(owner.observe('root', stepEvent(id))).toBe(reason); + expect(owner.observe('root', stepEvent(id))).toBe(reason); + } + }); + + it('cleans up deadlines on root completion, child settlement, stop and disposal', async () => { + const { owner, onDeadline } = create({ isBounded: () => true }); + owner.start('root'); + owner.startSpawn('child-spawn', 'root', 'child'); + owner.settleSpawn('child-spawn'); + owner.finishRoot('root'); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(onDeadline).not.toHaveBeenCalled(); + owner.start('root'); + owner.claimStop('root'); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(onDeadline).not.toHaveBeenCalled(); + const disposed = create(); + disposed.owner.start('child'); + disposed.owner.dispose(); + expect(disposed.owner.observe('child', toolEvent(0))).toBeUndefined(); + await vi.advanceTimersByTimeAsync(AGENT_SAFETY.passTimeoutMs); + expect(disposed.onDeadline).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/types/src/agent-safety.ts b/packages/types/src/agent-safety.ts index 3dc41c815..7ff9db5c5 100644 --- a/packages/types/src/agent-safety.ts +++ b/packages/types/src/agent-safety.ts @@ -127,3 +127,175 @@ export class AgentSafetyPass { return this.stopped; } } + +/** In-memory safety state for one runtime's session tree. Transport, persistence + * and transcript effects stay with the caller. A settled child can only resume + * through an explicit new spawn unless the runtime opts into activity resumes. */ +export class AgentSafetyLifecycle { + private readonly passes = new Map(); + private readonly settled = new Set(); + private readonly spawns = new Map< + string, + { + startedAtMs: number; + sessionId: string; + childSessionId?: string; + } + >(); + private readonly settledSpawns = new Set(); + private readonly timer: ReturnType; + private disposed = false; + private stopClaimed = false; + + constructor( + private readonly options: { + digest: (value: string) => string; + isBounded: (sessionId: string) => boolean; + onDeadline: (sessionId: string) => void; + resumeOnActivity?: boolean; + includeFailedTools?: boolean; + }, + ) { + this.timer = setInterval(() => { + if (this.disposed || this.stopped) return; + for (const spawn of this.spawns.values()) { + if (Date.now() - spawn.startedAtMs >= AGENT_SAFETY.passTimeoutMs) { + this.options.onDeadline(spawn.childSessionId ?? spawn.sessionId); + return; + } + } + for (const [id, pass] of this.passes) { + if (this.options.isBounded(id) && pass.checkTime()) { + this.options.onDeadline(id); + return; + } + } + }, 1_000); + this.timer.unref?.(); + } + + get stopped(): boolean { + return this.stopClaimed; + } + + start(sessionId: string, resume = false): AgentSafetyPass | undefined { + if (this.disposed || this.stopped) return; + if (resume) this.settled.delete(sessionId); + if (this.settled.has(sessionId)) return; + let pass = this.passes.get(sessionId); + if (!pass) { + pass = new AgentSafetyPass(this.options.digest); + this.passes.set(sessionId, pass); + } + return pass; + } + + settle(sessionId: string): void { + if (this.disposed) return; + this.passes.delete(sessionId); + this.settled.add(sessionId); + } + + /** Root turns may restart on activity without an explicit child spawn. */ + finishRoot(sessionId: string): void { + this.passes.delete(sessionId); + } + + startSpawn( + key: string, + sessionId: string, + childSessionId: string | null, + ): boolean { + if (this.disposed || this.stopped || this.settledSpawns.has(key)) + return false; + const existing = this.spawns.get(key); + this.spawns.set(key, { + startedAtMs: existing?.startedAtMs ?? Date.now(), + sessionId, + childSessionId: childSessionId ?? existing?.childSessionId, + }); + if (childSessionId) this.start(childSessionId, true); + return true; + } + + settleSpawn(key: string): void { + if (this.disposed) return; + this.settledSpawns.add(key); + const spawn = this.spawns.get(key); + if (spawn?.childSessionId) this.settle(spawn.childSessionId); + this.spawns.delete(key); + } + + /** The caller has already established membership in its session tree. */ + observe( + sessionId: string, + event: { type: string; properties?: unknown }, + ): AgentSafetyReason | undefined { + if (this.disposed || this.stopped) return; + const properties = safetyRecord(event.properties); + const status = safetyRecord(properties?.status); + const activity = + event.type === 'message.part.updated' || + event.type === 'message.updated' || + (event.type === 'session.status' && + ['busy', 'retry'].includes(String(status?.type))); + const pass = + this.passes.get(sessionId) ?? + (activity + ? this.start(sessionId, this.options.resumeOnActivity) + : undefined); + const part = safetyRecord(properties?.part); + if (event.type === 'message.part.updated' && part?.type === 'tool') { + const state = safetyRecord(part.state); + const terminal = + state?.status === 'completed' || + state?.status === 'error' || + (this.options.includeFailedTools && state?.status === 'failed'); + if (terminal && typeof part.callID === 'string' && part.callID) { + return pass?.completed({ + id: part.callID, + name: typeof part.tool === 'string' ? part.tool : '', + args: state?.input, + result: { + status: state?.status, + output: state?.output, + error: state?.error, + }, + }); + } + } + const info = safetyRecord(properties?.info); + if ( + event.type === 'message.updated' && + info?.role === 'assistant' && + typeof info.id === 'string' && + info.id && + this.options.isBounded(sessionId) + ) { + return pass?.step(info.id); + } + } + + /** Claim before any async effects so concurrent limits cannot stop twice. */ + claimStop(sessionId: string): AgentSafetyStopDetails | undefined { + if (this.disposed || this.stopped) return; + this.stopClaimed = true; + clearInterval(this.timer); + return { sessionId, ...this.passes.get(sessionId)?.snapshot() }; + } + + dispose(): void { + this.disposed = true; + clearInterval(this.timer); + this.passes.clear(); + this.settled.clear(); + this.spawns.clear(); + this.settledSpawns.clear(); + } +} + +function safetyRecord(value: unknown): Record | undefined { + return value && typeof value === 'object' && !Array.isArray(value) + ? (value as Record) + : undefined; +}