From 323ddcff3674fb6e4ebafcfe2859de13e22bee17 Mon Sep 17 00:00:00 2001 From: David Matejka Date: Mon, 14 Sep 2026 17:31:20 +0200 Subject: [PATCH 1/2] fix(sdk): count recovered tool calls once in the limits guard --- .../sdk/src/plugins/limits-guard/plugin.ts | 7 ++ .../src/plugins/limits-guard/recovery.test.ts | 96 +++++++++++++++++++ 2 files changed, 103 insertions(+) create mode 100644 packages/sdk/src/plugins/limits-guard/recovery.test.ts diff --git a/packages/sdk/src/plugins/limits-guard/plugin.ts b/packages/sdk/src/plugins/limits-guard/plugin.ts index 83dce86d..f248f143 100644 --- a/packages/sdk/src/plugins/limits-guard/plugin.ts +++ b/packages/sdk/src/plugins/limits-guard/plugin.ts @@ -5,6 +5,7 @@ import { llmEvents } from '~/core/llm/state.js' import { definePlugin } from '~/core/plugins/plugin-builder.js' import { selectPluginState } from '~/core/sessions/reducer.js' import type { SessionState } from '~/core/sessions/state.js' +import type { ToolCallId } from '~/core/tools/schema.js' import { toolEvents } from '~/core/tools/state.js' import { responseFingerprint, toolCallFingerprint } from '~/lib/utils/hash.js' import { mailboxEvents } from '~/plugins/mailbox/state.js' @@ -25,6 +26,8 @@ import { export interface AgentCounters { inferenceCount: number toolCallCount: number + /** Agents execute tools serially; recovery can start the same unfinished call again. */ + countedToolCallId?: ToolCallId spawnedAgentCount: number messagesSentCount: number /** Consecutive outbound-only turns with no newly consumed inbound work. */ @@ -259,12 +262,14 @@ export const limitsGuardPlugin = definePlugin('limits-guard') case 'tool_started': { const counters = limits.get(event.agentId) if (!counters) return limits + if (counters.countedToolCallId === event.toolCallId) return limits const fingerprint = toolCallFingerprint(event.toolName, event.input) const newLimits = new Map(limits) newLimits.set(event.agentId, { ...counters, toolCallCount: counters.toolCallCount + 1, + countedToolCallId: event.toolCallId, recentToolCallHashes: [...counters.recentToolCallHashes, fingerprint].slice(-20), }) return newLimits @@ -282,6 +287,7 @@ export const limitsGuardPlugin = definePlugin('limits-guard') newLimits.set(event.agentId, { ...counters, consecutiveToolFailures: restFailures, + countedToolCallId: undefined, }) return newLimits } @@ -301,6 +307,7 @@ export const limitsGuardPlugin = definePlugin('limits-guard') ...counters.consecutiveToolFailures, [toolName]: { count: (currentEntry?.count ?? 0) + 1, lastError: event.error }, }, + countedToolCallId: undefined, }) return newLimits } diff --git a/packages/sdk/src/plugins/limits-guard/recovery.test.ts b/packages/sdk/src/plugins/limits-guard/recovery.test.ts new file mode 100644 index 00000000..a74df382 --- /dev/null +++ b/packages/sdk/src/plugins/limits-guard/recovery.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, it } from 'bun:test' +import { AgentId } from '~/core/agents/schema.js' +import { agentEvents } from '~/core/agents/state.js' +import { withSessionId } from '~/core/events/test-helpers.js' +import type { DomainEvent } from '~/core/events/types.js' +import { createApplyEvent } from '~/core/sessions/apply-event.js' +import { SessionId } from '~/core/sessions/schema.js' +import { createSessionState, sessionEvents } from '~/core/sessions/state.js' +import { ToolCallId } from '~/core/tools/schema.js' +import { toolEvents } from '~/core/tools/state.js' +import { limitsGuardPlugin, selectAgentCounters } from './plugin.js' +import { checkLimits, resolveAgentLimits } from './limit-guard.js' + +const sessionId = SessionId('limits-recovery') +const agentId = AgentId('worker') +const applyEvent = createApplyEvent([limitsGuardPlugin.create()]) +const spawn = withSessionId( + sessionId, + agentEvents.create('agent_spawned', { + agentId, + definitionName: 'worker', + parentId: null, + }), +) +const start = (id: string) => + withSessionId( + sessionId, + toolEvents.create('tool_started', { + agentId, + toolCallId: ToolCallId(id), + toolName: 'run_command', + input: { command: 'sleep 60' }, + }), + ) +const restart = withSessionId( + sessionId, + sessionEvents.create('session_restarted', { + resetAgentIds: [], + clearedToolAgentIds: [agentId], + }), +) +const counters = (events: DomainEvent[]) => selectAgentCounters(events.reduce(applyEvent, createSessionState(sessionId, 'test', 0)), agentId) + +describe('limits during tool recovery', () => { + it('counts one in-flight call once across repeated restart events and replay', () => { + const events = [spawn, start('call-1'), restart, start('call-1'), restart, start('call-1'), restart, start('call-1')] + const result = counters(events) + expect(result.toolCallCount).toBe(1) + expect(result.recentToolCallHashes).toHaveLength(1) + expect(checkLimits(result, resolveAgentLimits()).status).toBe('ok') + expect(counters(events)).toEqual(result) + }) + + it('still counts completed identical calls as separate work', () => { + const events: DomainEvent[] = [spawn] + for (let i = 0; i < 4; i++) { + const id = ToolCallId(`call-${i}`) + events.push( + start(id), + withSessionId( + sessionId, + toolEvents.create('tool_completed', { + agentId, + toolCallId: id, + result: 'done', + }), + ), + ) + } + const result = counters(events) + expect(result.toolCallCount).toBe(4) + expect(result.recentToolCallHashes).toHaveLength(4) + expect(new Set(result.recentToolCallHashes).size).toBe(1) + expect(checkLimits(result, resolveAgentLimits())).toMatchObject({ status: 'hard_limit', limitName: 'maxRepeatedToolCalls' }) + }) + + it('clears the unfinished-call marker on failure without losing failure accounting', () => { + const result = counters([ + spawn, + start('call-1'), + restart, + start('call-1'), + withSessionId( + sessionId, + toolEvents.create('tool_failed', { + agentId, + toolCallId: ToolCallId('call-1'), + error: 'command failed', + }), + ), + ]) + expect(result.toolCallCount).toBe(1) + expect(result.countedToolCallId).toBeUndefined() + expect(Object.values(result.consecutiveToolFailures)).toEqual([{ count: 1, lastError: 'command failed' }]) + }) +}) From b97be3c679676953aff7745d5250131187534724 Mon Sep 17 00:00:00 2001 From: David Matejka Date: Mon, 14 Sep 2026 17:33:18 +0200 Subject: [PATCH 2/2] test(sdk): pass limits configuration in the recovery fixture --- packages/sdk/src/plugins/limits-guard/recovery.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/sdk/src/plugins/limits-guard/recovery.test.ts b/packages/sdk/src/plugins/limits-guard/recovery.test.ts index a74df382..82889bec 100644 --- a/packages/sdk/src/plugins/limits-guard/recovery.test.ts +++ b/packages/sdk/src/plugins/limits-guard/recovery.test.ts @@ -13,7 +13,7 @@ import { checkLimits, resolveAgentLimits } from './limit-guard.js' const sessionId = SessionId('limits-recovery') const agentId = AgentId('worker') -const applyEvent = createApplyEvent([limitsGuardPlugin.create()]) +const applyEvent = createApplyEvent([limitsGuardPlugin.create({})]) const spawn = withSessionId( sessionId, agentEvents.create('agent_spawned', {