diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index de1647618a..7dbbe7484d 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -9316,6 +9316,11 @@ describe('AiSdkBackend RunTrace', () => { request: { messages: [] }, continuation: 'none', }), + toolCallSafety: Promise.resolve({ + hadRawArgumentEvidence: false, + proofs: new Map(), + atomicProofs: new Map(), + }), }; }; @@ -12007,6 +12012,11 @@ describe('AiSdkBackend thinking persistence', () => { request: { messages: [] }, continuation: 'none', }), + toolCallSafety: Promise.resolve({ + hadRawArgumentEvidence: false, + proofs: new Map(), + atomicProofs: new Map(), + }), }); for await (const event of backend.send({ turnId: 'turn-1', text: 'hi', context: [] })) { diff --git a/packages/runtime/src/__tests__/deferred-guard.test.ts b/packages/runtime/src/__tests__/deferred-guard.test.ts index 8669f2a930..6f059d9f42 100644 --- a/packages/runtime/src/__tests__/deferred-guard.test.ts +++ b/packages/runtime/src/__tests__/deferred-guard.test.ts @@ -138,7 +138,31 @@ describe('tool-availability execute-boundary guard', () => { test('keeps WriteStdin args exact across canonical ledgers and projects telemetry', async () => { const h = makeHarness(); const implCalls: string[] = []; - const t = tool('WriteStdin', implCalls); + // The real WriteStdin tool (shell-tools.ts, buildWriteStdinTool) declares + // a heavyweight z.preprocess/refine schema enforcing PTY-specific business + // rules (ref format, input byte length, well-formed Unicode, ...) that + // this test has nothing to do with. `tool()`'s shared z.object({}) is + // wrong here for a different reason: ToolRuntime now executes the + // schema's own parsed value (see tool-runtime.ts), and an empty schema + // strips every key, which used to be masked only because the pre-fix + // ToolRuntime discarded that parsed value and executed the raw input + // instead. A schema matching WriteStdin's real field names, with no + // business-rule refinements, is what this test actually needs: proof + // that already-valid args pass through Runtime's plumbing unchanged, not + // a reproduction of WriteStdin's own validation. + const t: MakaTool = { + name: 'WriteStdin', + description: 'WriteStdin', + parameters: z.object({ + ref: z.string(), + input: z.string(), + size: z.object({ cols: z.number(), rows: z.number() }), + }), + impl: () => { + implCalls.push('WriteStdin'); + return { ok: true }; + }, + }; const args = { ref: 'maka://runtime/background-tasks/pty-1', input: 'password=ordinary-audited-input\r', diff --git a/packages/runtime/src/__tests__/length-cutoff-tool-execution-repro.test.ts b/packages/runtime/src/__tests__/length-cutoff-tool-execution-repro.test.ts new file mode 100644 index 0000000000..1a167af2da --- /dev/null +++ b/packages/runtime/src/__tests__/length-cutoff-tool-execution-repro.test.ts @@ -0,0 +1,987 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import { describe, test } from 'node:test'; +import type { LanguageModelV4StreamPart } from '@ai-sdk/provider'; +import type { BackendSendInput } from '@maka/core/backend-types'; +import type { SessionEvent } from '@maka/core/events'; +import type { LlmConnection } from '@maka/core/llm-connections'; +import type { RuntimeEvent } from '@maka/core/runtime-event'; +import type { SessionHeader } from '@maka/core/session'; +import { MockLanguageModelV4, simulateReadableStream } from 'ai/test'; +import { z } from 'zod'; + +import { createSessionEventMapMemory, mapSessionEventToRuntimeEvent } from '../ai-sdk-flow.js'; +import type { InvocationContext } from '../invocation-context.js'; +import type { MakaTool } from '../tool-runtime.js'; +import { createTestAiSdkBackend } from './execution-boundary-test-helpers.js'; + +/** + * Production-path regression coverage for the execution boundary introduced by + * tool-call-execution-guard.ts. Every assertion reaches the real + * AiSdkBackend -> ModelAdapter -> ToolRuntime settlement path and checks the + * irreversible result (zero executions or exactly one), rather than restating + * the guard implementation. + */ + +function header(): SessionHeader { + return { + id: 'session-1', + workspaceRoot: '/tmp/maka-repro', + cwd: '/tmp/maka-repro', + createdAt: 1, + name: 'Repro', + titleIsManual: true, + isFlagged: false, + labels: [], + isArchived: false, + status: 'active', + statusUpdatedAt: 1, + hasUnread: false, + backend: 'ai-sdk', + llmConnectionSlug: 'anthropic-main', + connectionLocked: true, + model: 'claude-sonnet-4-5-20250929', + permissionMode: 'ask', + schemaVersion: 1, + }; +} + +function connection(): LlmConnection { + return { + slug: 'anthropic-main', + name: 'Anthropic', + providerType: 'anthropic', + defaultModel: 'claude-sonnet-4-5-20250929', + enabled: true, + createdAt: 1, + updatedAt: 1, + }; +} + +function idGenerator(): () => string { + let index = 0; + return () => `id-${++index}`; +} + +function monotonicClock(): () => number { + let value = 1_000; + return () => ++value; +} + +function durableTurnHarness(turnId: string, text: string) { + const runId = 'run-1'; + const invocationId = 'invocation-1'; + const anchor: RuntimeEvent = { + id: `runtime-user-${turnId}`, + invocationId, + runId, + sessionId: 'session-1', + turnId, + ts: 1, + partial: false, + role: 'user', + author: 'user', + content: { kind: 'text', text }, + }; + const ledger: RuntimeEvent[] = [anchor]; + const memory = createSessionEventMapMemory(); + const ctx: InvocationContext = { + sessionId: 'session-1', + invocationId, + runId, + turnId, + source: 'desktop', + startedAt: 1, + request: { + sessionId: 'session-1', + turnId, + text, + source: 'desktop', + initialRuntimeEvent: anchor, + }, + newId: idGenerator(), + now: monotonicClock(), + }; + return { + anchor, + ledger, + loadTurnRuntimeEvents: async (requestedTurnId: string) => + ledger.filter((event) => event.turnId === requestedTurnId), + input: (overrides: Partial = {}): BackendSendInput => ({ + turnId, + text, + context: [], + headAnchorRuntimeEvent: anchor, + ...overrides, + }), + record: (event: SessionEvent): void => { + const mapped = mapSessionEventToRuntimeEvent(event, ctx, memory); + if (mapped.partial !== true && mapped.content?.kind !== 'error') ledger.push(mapped); + }, + }; +} + +async function drainDurably( + iterable: AsyncIterable, + durable: ReturnType, +): Promise { + const events: SessionEvent[] = []; + for await (const event of iterable) { + durable.record(event); + events.push(event); + } + return events; +} + +const ZERO_USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 1, text: 1, reasoning: 0 }, +}; + +function makeGate(): { promise: Promise; release: () => void } { + let release!: () => void; + const promise = new Promise((resolve) => { + release = resolve; + }); + return { promise, release }; +} + +function hangingProviderStream( + chunks: readonly LanguageModelV4StreamPart[], + signal: AbortSignal | undefined, +): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(chunk); + const abort = () => controller.error(signal?.reason ?? new Error('aborted')); + if (signal?.aborted) abort(); + else signal?.addEventListener('abort', abort, { once: true }); + }, + }); +} + +type UnifiedFinishReason = 'length' | 'stop' | 'tool-calls' | 'content-filter' | 'error' | 'other'; + +type FinishReason = { + unified: UnifiedFinishReason; + raw: string | undefined; +}; + +function doneChunks(): LanguageModelV4StreamPart[] { + return [ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 'text-final' }, + { type: 'text-delta', id: 'text-final', delta: 'done' }, + { type: 'text-end', id: 'text-final' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]; +} + +function twoStepModel(firstChunks: LanguageModelV4StreamPart[]): MockLanguageModelV4 { + let calls = 0; + return new MockLanguageModelV4({ + doStream: async () => { + calls += 1; + return { + stream: simulateReadableStream({ + chunks: calls === 1 ? firstChunks : doneChunks(), + initialDelayInMs: null, + chunkDelayInMs: null, + }), + }; + }, + }); +} + +function toolCallChunks( + delivery: 'incremental' | 'atomic', + finishReason: FinishReason, + options: { + rawId?: string; + resolvedId?: string; + rawToolName?: string; + resolvedToolName?: string; + rawInput?: unknown; + projectedInput?: unknown; + } = {}, +): LanguageModelV4StreamPart[] { + const rawId = options.rawId ?? 'call-1'; + const resolvedId = options.resolvedId ?? rawId; + const rawToolName = options.rawToolName ?? 'Write'; + const resolvedToolName = options.resolvedToolName ?? rawToolName; + const rawInput = options.rawInput ?? { path: 'notes.md', content: 'hello from the model' }; + const projectedInput = options.projectedInput ?? rawInput; + const chunks: LanguageModelV4StreamPart[] = [ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: rawId, toolName: rawToolName }, + ]; + if (delivery === 'incremental') { + chunks.push({ type: 'tool-input-delta', id: rawId, delta: JSON.stringify(rawInput) }); + } + chunks.push( + { type: 'tool-input-end', id: rawId }, + { + type: 'tool-call', + toolCallId: resolvedId, + toolName: resolvedToolName, + input: JSON.stringify(projectedInput), + }, + { type: 'finish', finishReason, usage: ZERO_USAGE }, + ); + return chunks; +} + +function writeTool(onExecute: (input: unknown) => void): MakaTool { + return { + name: 'Write', + description: 'Write file contents', + parameters: z.object({ path: z.string(), content: z.string() }), + impl: async (input) => { + onExecute(input); + return { ok: true }; + }, + }; +} + +function notifyTool(onExecute: (input: unknown) => void): MakaTool { + return { + name: 'Notify', + description: 'Send a notification', + parameters: z.object({ message: z.string() }), + impl: async (input) => { + onExecute(input); + return { ok: true }; + }, + }; +} + +function shellTool(onExecute: (input: unknown) => void): MakaTool { + return { + name: 'Shell', + description: 'Run a shell command', + parameters: z.object({ command: z.string() }), + impl: async (input) => { + onExecute(input); + return { ok: true }; + }, + }; +} + +/** A genuinely zero-argument tool -- the schema itself takes nothing. */ +function pingTool(onExecute: (input: unknown) => void): MakaTool { + return { + name: 'Ping', + description: 'Zero-argument health check', + parameters: z.object({}), + impl: async (input) => { + onExecute(input); + return { ok: true }; + }, + }; +} + +/** Zero required arguments, but a declared default -- composes with P1. */ +function listTodosTool(onExecute: (input: unknown) => void): MakaTool { + return { + name: 'ListTodos', + description: 'List todos', + parameters: z.object({ limit: z.number().default(10) }), + impl: async (input) => { + onExecute(input); + return { ok: true }; + }, + }; +} + +async function runModel( + model: MockLanguageModelV4, + tools: MakaTool[], + turnId = 'turn-1', +): Promise { + const durable = durableTurnHarness(turnId, 'write it'); + const backend = createTestAiSdkBackend({ + sessionId: `session-${turnId}`, + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + modelFactory: () => model, + tools, + loadTurnRuntimeEvents: durable.loadTurnRuntimeEvents, + newId: idGenerator(), + now: monotonicClock(), + }); + return drainDurably(backend.send(durable.input()), durable); +} + +async function executionCountFor( + delivery: 'incremental' | 'atomic', + finishReason: FinishReason, +): Promise { + let executions = 0; + const model = twoStepModel(toolCallChunks(delivery, finishReason)); + await runModel(model, [ + writeTool(() => { + executions += 1; + }), + ]).catch(() => []); + return executions; +} + +describe('tool execution safety (real production path)', () => { + for (const delivery of ['incremental', 'atomic'] as const) { + for (const scenario of [ + { unified: 'stop', raw: 'stop', executions: 1 }, + { unified: 'tool-calls', raw: 'tool_calls', executions: 1 }, + { unified: 'length', raw: 'length', executions: 0 }, + { unified: 'content-filter', raw: 'content_filter', executions: 0 }, + { unified: 'error', raw: 'error', executions: 0 }, + { unified: 'other', raw: 'other', executions: 0 }, + ] as const) { + test(`${delivery} + ${scenario.unified}: executes ${scenario.executions} time(s)`, async () => { + const executions = await executionCountFor(delivery, { + unified: scenario.unified, + raw: scenario.raw, + }); + assert.equal(executions, scenario.executions); + }); + } + } + + // Finish-reason authority: chunkFinishReason (model-adapter.ts) already + // resolves "other" through the provider's own raw spelling for step + // settlement. This proves the guard's own terminal classification for an + // incrementally-streamed call now agrees with that resolution end to end, + // through the real ModelAdapter -> resolveToolCallSafety wiring, rather + // than only in the unit-level tracker tests. ("unknown" is exercised at + // the unit level only — it is Maka's own settlement-layer fallback, never + // a value a real raw SDK finish chunk's own `unified` field carries.) + test('incremental + finish reason unified "other" but raw "stop": executes once', async () => { + const executions = await executionCountFor('incremental', { unified: 'other', raw: 'stop' }); + assert.equal(executions, 1); + }); + + test('missing terminal event executes zero times', async () => { + let executions = 0; + const chunks = toolCallChunks('incremental', { unified: 'stop', raw: 'stop' }).slice(0, -1); + await runModel(twoStepModel(chunks), [ + writeTool(() => { + executions += 1; + }), + ]).catch(() => []); + assert.equal(executions, 0); + }); + + test('truncated raw arguments without tool-input-end execute zero times', async () => { + let executions = 0; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-1', delta: '{"path":"notes.md","content":"unf' }, + { + type: 'finish', + finishReason: { unified: 'length', raw: 'length' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + executions += 1; + }), + ]).catch(() => []); + assert.equal(executions, 0); + }); + + test('abort after a complete incremental call but before terminal finish executes zero times', async () => { + let executions = 0; + const durable = durableTurnHarness('turn-abort', 'write it'); + const chunksEnqueued = makeGate(); + const model = new MockLanguageModelV4({ + doStream: async (options) => { + const chunks = toolCallChunks('incremental', { unified: 'stop', raw: 'stop' }).slice(0, -1); + const stream = hangingProviderStream(chunks, options.abortSignal); + chunksEnqueued.release(); + return { stream }; + }, + }); + const backend = createTestAiSdkBackend({ + sessionId: 'session-abort', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + modelFactory: () => model, + tools: [ + writeTool(() => { + executions += 1; + }), + ], + loadTurnRuntimeEvents: durable.loadTurnRuntimeEvents, + newId: idGenerator(), + now: monotonicClock(), + }); + const drainPromise = drainDurably(backend.send(durable.input()), durable).catch(() => []); + await chunksEnqueued.promise; + await backend.stop('user_stop'); + await drainPromise; + assert.equal(executions, 0); + }); + + test('concurrent incremental requests sharing call_1 stay isolated', async () => { + let safeExecutions = 0; + let unsafeExecutions = 0; + const safe = runModel( + twoStepModel(toolCallChunks('incremental', { unified: 'stop', raw: 'stop' })), + [ + writeTool(() => { + safeExecutions += 1; + }), + ], + 'turn-safe', + ); + const unsafe = runModel( + twoStepModel(toolCallChunks('incremental', { unified: 'length', raw: 'length' })), + [ + writeTool(() => { + unsafeExecutions += 1; + }), + ], + 'turn-unsafe', + ).catch(() => []); + await unsafe; + await safe; + assert.equal(safeExecutions, 1); + assert.equal(unsafeExecutions, 0); + }); + + test('concurrent atomic requests sharing call_1 stay isolated', async () => { + let safeExecutions = 0; + let unsafeExecutions = 0; + const safe = runModel( + twoStepModel(toolCallChunks('atomic', { unified: 'stop', raw: 'stop' })), + [ + writeTool(() => { + safeExecutions += 1; + }), + ], + 'turn-safe-atomic', + ); + const unsafe = runModel( + twoStepModel(toolCallChunks('atomic', { unified: 'length', raw: 'length' })), + [ + writeTool(() => { + unsafeExecutions += 1; + }), + ], + 'turn-unsafe-atomic', + ).catch(() => []); + await unsafe; + await safe; + assert.equal(safeExecutions, 1); + assert.equal(unsafeExecutions, 0); + }); + + test('raw evidence for call_1 cannot authorize a resolved call_2', async () => { + let executions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { + rawId: 'call_1', + resolvedId: 'call_2', + }, + ), + ); + await runModel(model, [ + writeTool(() => { + executions += 1; + }), + ]).catch(() => []); + assert.equal(executions, 0); + }); + + // This is the installed Google adapter's real mixed-delivery shape: an + // argument-bearing tool call streams start/delta.../end/final-call, while a + // genuinely zero-argument sibling in the SAME physical request legitimately + // streams start/end/final-call with ZERO tool-input-delta chunks. Each + // call's eligibility must come only from its own lifecycle — the + // argument-bearing sibling having raw evidence must not poison the + // zero-argument sibling's own, independently-proved atomic completion. + // A. A genuinely zero-argument sibling -- the installed Google adapter's + // real wire shape, input "{}" -- executes from its own atomic proof, + // isolated from an argument-bearing sibling in the same physical request. + test('a legitimate zero-argument sibling executes from its own atomic proof, isolated from an argument-bearing sibling', async () => { + let writeExecutions = 0; + let pingExecutions = 0; + let writeReceived: unknown; + let pingReceived: unknown; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { + type: 'tool-input-delta', + id: 'call-incremental', + delta: JSON.stringify(writeInput), + }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + { type: 'tool-input-start', id: 'call-atomic', toolName: 'Ping' }, + { type: 'tool-input-end', id: 'call-atomic' }, + { type: 'tool-call', toolCallId: 'call-atomic', toolName: 'Ping', input: '{}' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool((input) => { + writeExecutions += 1; + writeReceived = input; + }), + pingTool((input) => { + pingExecutions += 1; + pingReceived = input; + }), + ]); + assert.equal(writeExecutions, 1); + assert.equal(pingExecutions, 1); + assert.deepEqual(writeReceived, writeInput); + assert.deepEqual(pingReceived, {}); + }); + + // B. THE CRITICAL REGRESSION CASE: a zero-delta call whose SDK-resolved + // final `tool-call` carries a NON-EMPTY projected input must still not + // execute with that value. Zero raw delta bytes for this id is itself the + // proof that the provider supplied no arguments (see ai-sdk-backend.ts); + // a divergent non-empty projection is never trusted, so this executes with + // the canonical empty object -- which, against Notify's *required* + // `message` field, fails schema validation and never reaches `impl`. + test('a zero-delta sibling with a non-empty final projected input still executes zero times', async () => { + let writeExecutions = 0; + let notifyExecutions = 0; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + // Zero tool-input-delta chunks for call-atomic, yet the SDK's own + // resolved tool-call somehow carries a non-empty payload -- a stale + // repair, a provider bug, or worse. It must never reach `impl`. + { type: 'tool-input-start', id: 'call-atomic', toolName: 'Notify' }, + { type: 'tool-input-end', id: 'call-atomic' }, + { + type: 'tool-call', + toolCallId: 'call-atomic', + toolName: 'Notify', + input: JSON.stringify({ message: 'done' }), + }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + notifyTool(() => { + notifyExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 1); + assert.equal(notifyExecutions, 0); + }); + + // C. Name substitution under a zero-delta, canonically-empty call is still + // rejected -- isolated from case B by using an empty projected input here, + // so this test fails only on identity, never on argument content. + test('a zero-delta sibling whose final name disagrees with its own observed start is still rejected', async () => { + let writeExecutions = 0; + let notifyExecutions = 0; + let shellExecutions = 0; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + // Observed as "Notify" with zero deltas, but the AI SDK resolves the + // same id under a DIFFERENT tool name. The guard's own observed name + // for this id must still win — exactly like the raw-byte proof case + // below — never the SDK's post-hoc projection under either name. + { type: 'tool-input-start', id: 'call-atomic', toolName: 'Notify' }, + { type: 'tool-input-end', id: 'call-atomic' }, + { type: 'tool-call', toolCallId: 'call-atomic', toolName: 'Shell', input: '{}' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + notifyTool(() => { + notifyExecutions += 1; + }), + shellTool(() => { + shellExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 1); + assert.equal(notifyExecutions, 0); + assert.equal(shellExecutions, 0); + }); + + // D. Id substitution under a zero-delta, canonically-empty call is still + // rejected: the started id never got its own resolved tool-call, and the + // resolved id never got its own start/end, so neither has an atomic proof + // -- exactly the raw-byte proof's "call_1 evidence cannot authorize a + // resolved call_2" shape, one level down. + test('a zero-delta start/end pair cannot authorize a resolved call under a different id', async () => { + let writeExecutions = 0; + let pingExecutions = 0; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + { type: 'tool-input-start', id: 'call-atomic-started', toolName: 'Ping' }, + { type: 'tool-input-end', id: 'call-atomic-started' }, + // Resolved under a different id than the one that streamed start/end. + { type: 'tool-call', toolCallId: 'call-atomic-resolved', toolName: 'Ping', input: '{}' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + pingTool(() => { + pingExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 1); + assert.equal(pingExecutions, 0); + }); + + // E. A zero-delta call missing its own tool-input-end (empty projected + // input, so only the incomplete lifecycle -- not argument content -- is + // under test here) is still rejected, even next to a proved sibling. + test('a zero-delta sibling missing its own tool-input-end is still rejected, even next to a proved sibling', async () => { + let writeExecutions = 0; + let pingExecutions = 0; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + // No tool-input-end for call-atomic at all: an incomplete lifecycle, + // not a proof of atomicity. + { type: 'tool-input-start', id: 'call-atomic', toolName: 'Ping' }, + { type: 'tool-call', toolCallId: 'call-atomic', toolName: 'Ping', input: '{}' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + pingTool(() => { + pingExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 1); + assert.equal(pingExecutions, 0); + }); + + // F. A complete, unpoisoned zero-delta lifecycle next to an unsafe + // terminal reason still executes zero times -- terminal safety gates the + // atomic proof exactly like it gates the raw-byte proof. + test('a zero-delta sibling is withheld when the physical request terminates unsafely', async () => { + let writeExecutions = 0; + let pingExecutions = 0; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + { type: 'tool-input-start', id: 'call-atomic', toolName: 'Ping' }, + { type: 'tool-input-end', id: 'call-atomic' }, + { type: 'tool-call', toolCallId: 'call-atomic', toolName: 'Ping', input: '{}' }, + { + type: 'finish', + finishReason: { unified: 'length', raw: 'length' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + pingTool(() => { + pingExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 0); + assert.equal(pingExecutions, 0); + }); + + // G. Composes with P1: a genuinely empty raw lifecycle, run through + // ToolRuntime's own schema, reaches `impl` with the schema's declared + // default filled in -- and the SDK's own (divergent, non-empty) projected + // input is never what gets used, proving the atomic path never falls back + // to `toolCall.input` even when a value happens to be available there. + // + // Mixed with a real incremental sibling on purpose: the per-call atomic + // proof's canonical-empty-value rule is scoped to `hadRawArgumentEvidence` + // being true (see ai-sdk-backend.ts) -- a genuinely whole-request-atomic + // delivery is a separate, unchanged policy that trusts `toolCall.input` + // directly, so this composition only exercises the new path with a + // sibling that streamed real delta bytes. + test('a zero-delta call composes with a schema default: the defaulted value reaches impl, never the divergent SDK projection', async () => { + let writeExecutions = 0; + let listTodosExecutions = 0; + let listTodosReceived: unknown; + const writeInput = { path: 'notes.md', content: 'hello' }; + const model = twoStepModel([ + { type: 'stream-start', warnings: [] }, + { type: 'tool-input-start', id: 'call-incremental', toolName: 'Write' }, + { type: 'tool-input-delta', id: 'call-incremental', delta: JSON.stringify(writeInput) }, + { type: 'tool-input-end', id: 'call-incremental' }, + { + type: 'tool-call', + toolCallId: 'call-incremental', + toolName: 'Write', + input: JSON.stringify(writeInput), + }, + { type: 'tool-input-start', id: 'call-atomic', toolName: 'ListTodos' }, + { type: 'tool-input-end', id: 'call-atomic' }, + { + type: 'tool-call', + toolCallId: 'call-atomic', + toolName: 'ListTodos', + // Divergent from the canonical empty proof on purpose: must never + // reach `impl`, whether or not it happens to satisfy the schema. + input: JSON.stringify({ limit: 999 }), + }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: ZERO_USAGE, + }, + ]); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + listTodosTool((input) => { + listTodosExecutions += 1; + listTodosReceived = input; + }), + ]); + assert.equal(writeExecutions, 1); + assert.equal(listTodosExecutions, 1); + assert.deepEqual(listTodosReceived, { limit: 10 }); + }); + + test('proved Write identity cannot be substituted with Shell under the same id', async () => { + let writeExecutions = 0; + let shellExecutions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { + rawToolName: 'Write', + resolvedToolName: 'Shell', + rawInput: { path: 'notes.md', content: 'hello' }, + projectedInput: { command: 'echo hi' }, + }, + ), + ); + await runModel(model, [ + writeTool(() => { + writeExecutions += 1; + }), + shellTool(() => { + shellExecutions += 1; + }), + ]).catch(() => []); + assert.equal(writeExecutions, 0); + assert.equal(shellExecutions, 0); + }); + + test('ToolRuntime receives the raw-proved object, never divergent projected input', async () => { + let receivedInput: unknown; + let executions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { + rawInput: { path: 'safe.md', content: 'hello' }, + projectedInput: { path: 'evil.md', content: 'hello' }, + }, + ), + ); + await runModel(model, [ + writeTool((input) => { + executions += 1; + receivedInput = input; + }), + ]); + assert.equal(executions, 1); + assert.deepEqual(receivedInput, { path: 'safe.md', content: 'hello' }); + }); + + // P1: the schema's own output (defaults filled in) must reach `impl`, and + // it must be derived from the guard's raw-proved value, never the SDK's + // divergent projected input — the two security properties compose. + test('a schema default is applied on top of the raw-proved value, never the divergent SDK projection', async () => { + let receivedInput: unknown; + let executions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { + rawInput: { path: 'safe.md' }, + projectedInput: { path: 'evil.md', content: 'untrusted' }, + }, + ), + ); + const tool: MakaTool = { + name: 'Write', + description: 'Write file contents', + parameters: z.object({ + path: z.string(), + content: z.string().default('placeholder'), + }), + impl: async (input) => { + executions += 1; + receivedInput = input; + return { ok: true }; + }, + }; + await runModel(model, [tool]); + assert.equal(executions, 1); + assert.deepEqual(receivedInput, { path: 'safe.md', content: 'placeholder' }); + }); + + // P1 CASE E: structurally valid JSON that the tool's own schema rejects + // (a required field is missing) must never reach `impl`, through the full + // production dispatch chain — not just at the ToolRuntime unit level. + test('raw-proved arguments that fail the declared schema execute zero times', async () => { + let executions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { rawInput: { path: 'notes.md' } }, + ), + ); + await runModel(model, [ + writeTool(() => { + executions += 1; + }), + ]).catch(() => []); + assert.equal(executions, 0); + }); + + test('matching id/name/value executes exactly once with the proved object', async () => { + let receivedInput: unknown; + let executions = 0; + const model = twoStepModel( + toolCallChunks( + 'incremental', + { unified: 'stop', raw: 'stop' }, + { + rawInput: { path: 'notes.md', content: 'hello' }, + }, + ), + ); + await runModel(model, [ + writeTool((input) => { + executions += 1; + receivedInput = input; + }), + ]); + assert.equal(executions, 1); + assert.deepEqual(receivedInput, { path: 'notes.md', content: 'hello' }); + }); +}); diff --git a/packages/runtime/src/__tests__/tool-args-violation.test.ts b/packages/runtime/src/__tests__/tool-args-violation.test.ts index d5b6942c0a..50763281aa 100644 --- a/packages/runtime/src/__tests__/tool-args-violation.test.ts +++ b/packages/runtime/src/__tests__/tool-args-violation.test.ts @@ -230,7 +230,7 @@ test('ToolRuntime enforces the declared schema before impl without a permissionA assert.match(said, /Read takes `file_path`, `offset`, `limit`\./); }); -test('ToolRuntime validates without rewriting arguments at permission and implementation boundaries', async () => { +test('ToolRuntime uses the schema-derived value (defaults filled in, transforms applied) at permission and implementation boundaries', async () => { const observed: unknown[] = []; const runtime = createTestToolRuntime({ sessionId: 'session-1', @@ -272,11 +272,56 @@ test('ToolRuntime validates without rewriting arguments at permission and implem }, }); - const expected = { file_path: ' /workspace/a.ts ' }; + // The raw call omitted `limit` and padded `file_path` with whitespace; the + // schema's `default()`/`transform()` output — not the literal input — is + // what reaches both the permission check and the implementation. + const expected = { file_path: '/workspace/a.ts', limit: 25 }; assert.deepEqual(observed, [expected, expected]); assert.deepEqual(result, expected); }); +test('ToolRuntime applies z.preprocess output the same way as transform/default output', async () => { + const observed: unknown[] = []; + const runtime = createTestToolRuntime({ + sessionId: 'session-1', + header: header(), + connection: connection(), + modelId: 'mock-model', + appendMessage: async () => {}, + newId: nextId(), + now: () => 1, + getPermissionPauseTarget: () => null, + }); + const parameters = z.object({ + limit: z.preprocess((value) => (value === undefined ? '10' : value), z.coerce.number()), + }); + const tool: MakaTool = { + name: 'List', + description: 'test', + parameters, + impl: async (args) => { + observed.push(structuredClone(args)); + return args; + }, + }; + + const { result } = await runtime.settleToolCall({ + tool, + turnId: 'turn-1', + toolCallId: 'tool-preprocess', + input: {}, + abortSignal: new AbortController().signal, + eventSink: { + push: () => {}, + pushAndWaitUntilConsumed: async () => {}, + }, + }); + + const expected = { limit: 10 }; + assert.deepEqual(observed, [expected]); + assert.deepEqual(result, expected); +}); + test('a sandbox denial names the tool that widens the boundary', async () => { const messages: StoredMessage[] = []; const events: SessionEvent[] = []; diff --git a/packages/runtime/src/__tests__/tool-call-execution-guard.test.ts b/packages/runtime/src/__tests__/tool-call-execution-guard.test.ts new file mode 100644 index 0000000000..2d70e2b2e0 --- /dev/null +++ b/packages/runtime/src/__tests__/tool-call-execution-guard.test.ts @@ -0,0 +1,460 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import { describe, test } from 'node:test'; + +import { + createToolCallSafetyTracker, + isSafeToolExecutionStepOutcome, + observeRawChunk, + resolveToolCallSafety, +} from '../tool-call-execution-guard.js'; +import type { ModelFailure, ModelStepOutcome } from '../model-protocol.js'; + +function completeToolCallParts(id: string, input: unknown, toolName = 'Write') { + return [ + { type: 'tool-input-start', id, toolName }, + { type: 'tool-input-delta', id, delta: JSON.stringify(input) }, + { type: 'tool-input-end', id }, + { type: 'tool-call', toolCallId: id, toolName, input: JSON.stringify(input) }, + ]; +} + +function pushCompleteCall( + tracker: ReturnType, + id = 'call-1', + input: unknown = { path: 'a.md' }, +): void { + for (const part of completeToolCallParts(id, input)) observeRawChunk(tracker, part); +} + +function hasProof( + tracker: ReturnType, + id = 'call-1', + meta?: { providerReason?: string }, +): boolean { + return resolveToolCallSafety(tracker, meta).proofs.has(id); +} + +describe('tool-call-execution-guard', () => { + test('stop positively proves the raw-stream name and value', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker, 'call-1', { path: 'a.md', content: 'hello' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.hadRawArgumentEvidence, true); + assert.deepEqual(safety.proofs.get('call-1'), { + name: 'Write', + value: { path: 'a.md', content: 'hello' }, + }); + }); + + test('tool-calls is also an explicitly safe terminal reason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'tool-calls' }); + assert.equal(hasProof(tracker), true); + }); + + for (const finishReason of ['length', 'content-filter', 'error', 'other', 'unknown']) { + test(`${finishReason} never authorizes a fully streamed call`, () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason }); + assert.equal(hasProof(tracker), false); + }); + } + + test('provider error part fails closed', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'error', error: new Error('upstream 500') }); + assert.equal(hasProof(tracker), false); + }); + + test('abort part fails closed', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'abort' }); + assert.equal(hasProof(tracker), false); + }); + + test('missing terminal event fails closed even when fallback metadata says stop', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + // Finish-reason authority: model-adapter.ts already resolves a stronger + // finish reason (chunkFinishReason) that falls back to the provider's own + // spelling when the SDK's unified reason is "other"/"unknown", and passes + // it as providerReason. Before this, the tracker's own local (unified-only) + // read of the same finish chunk could disagree with that stronger answer + // for the exact same physical request. + test('a real finish event with an ambiguous unified reason is rescued by a real providerReason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: { unified: 'other', raw: 'stop' } }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), true); + }); + + test('unknown is also rescuable when providerReason resolves to tool-calls', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { + type: 'finish', + finishReason: { unified: 'unknown', raw: 'tool-calls' }, + }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'tool-calls' }), true); + }); + + test('an ambiguous reason with no rescuing providerReason still fails closed', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: { unified: 'other', raw: 'other' } }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'other' }), false); + }); + + test('an abort terminal event is never rescued by providerReason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'abort' }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + test('a provider error terminal event is never rescued by providerReason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'error', error: new Error('upstream 500') }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + test('a directly observed length finish reason is never rescued by providerReason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'length' }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + test('a directly observed content-filter finish reason is never rescued by providerReason', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'content-filter' }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + test('a second finish-shaped event is treated as poisoning, not a softer verdict', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: { unified: 'other', raw: 'stop' } }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(hasProof(tracker, 'call-1', { providerReason: 'stop' }), false); + }); + + test('missing tool-input-end fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { + type: 'tool-input-delta', + id: 'call-1', + delta: '{"path":"a.md"}', + }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(hasProof(tracker), false); + }); + + test('malformed JSON fails closed despite start/end and a safe finish', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-delta', id: 'call-1', delta: '{"path":' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(hasProof(tracker), false); + }); + + test('delta before start is contradictory evidence and fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-delta', id: 'call-1', delta: '{"path":"a.md"}' }); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(hasProof(tracker), false); + }); + + test('duplicate start fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-delta', id: 'call-1', delta: '{"path":"a.md"}' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(hasProof(tracker), false); + }); + + test('tool evidence after a terminal event poisons the request', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + observeRawChunk(tracker, { type: 'tool-input-delta', id: 'call-1', delta: ' ' }); + assert.equal(hasProof(tracker), false); + }); + + test('pure atomic delivery has no proof and keeps request-level evidence false', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'Write', + input: JSON.stringify({ path: 'a.md' }), + }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.hadRawArgumentEvidence, false); + assert.equal(safety.proofs.has('call-1'), false); + }); + + test('a mixed request records raw evidence globally while leaving the atomic sibling unproved', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker, 'call-incremental', { path: 'a.md' }); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-atomic', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-atomic' }); + observeRawChunk(tracker, { + type: 'tool-call', + toolCallId: 'call-atomic', + toolName: 'Write', + input: JSON.stringify({ path: 'b.md' }), + }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.hadRawArgumentEvidence, true); + assert.equal(safety.proofs.has('call-incremental'), true); + assert.equal(safety.proofs.has('call-atomic'), false); + }); + + test('concurrent request trackers can reuse call_1 without cross-resolution', () => { + const safe = createToolCallSafetyTracker(); + const unsafe = createToolCallSafetyTracker(); + pushCompleteCall(safe, 'call_1', { path: 'safe.md' }); + pushCompleteCall(unsafe, 'call_1', { path: 'unsafe.md' }); + observeRawChunk(unsafe, { type: 'finish', finishReason: 'length' }); + observeRawChunk(safe, { type: 'finish', finishReason: 'stop' }); + + assert.equal(hasProof(safe, 'call_1'), true); + assert.equal(hasProof(unsafe, 'call_1'), false); + assert.deepEqual(resolveToolCallSafety(safe).proofs.get('call_1')?.value, { + path: 'safe.md', + }); + }); + + test('resolution is stable when read more than once', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + const first = resolveToolCallSafety(tracker); + const second = resolveToolCallSafety(tracker); + assert.equal(second, first); + }); +}); + +// The installed Google adapter can legitimately deliver a zero-argument tool +// call as start -> end -> final-call with ZERO tool-input-delta chunks. These +// cases prove that shape gets its own per-id positive proof, and that the +// proof requires the same completeness/identity discipline `proofs` does — +// never "no deltas" alone, and never contamination from a sibling's evidence. +describe('atomicProofs (per-call zero-argument evidence)', () => { + function pushAtomicCall( + tracker: ReturnType, + id = 'call-1', + toolName = 'ListTodos', + ): void { + observeRawChunk(tracker, { type: 'tool-input-start', id, toolName }); + observeRawChunk(tracker, { type: 'tool-input-end', id }); + } + + test('a genuinely atomic zero-argument call gets its own positive proof', () => { + const tracker = createToolCallSafetyTracker(); + pushAtomicCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.deepEqual(safety.atomicProofs.get('call-1'), { name: 'ListTodos' }); + assert.equal(safety.proofs.has('call-1'), false); + }); + + test('an argument-bearing sibling no longer poisons a legitimate zero-argument call', () => { + const tracker = createToolCallSafetyTracker(); + pushCompleteCall(tracker, 'call-args', { path: 'a.md' }); + pushAtomicCall(tracker, 'call-zero', 'ListTodos'); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + // Sibling state is exactly as before: request-global evidence is true, + // and the zero-arg call still has no entry in `proofs` (it never + // streamed argument bytes) — but it now has its own atomic proof, which + // is what a caller must consult before falling back to + // `hadRawArgumentEvidence`. + assert.equal(safety.hadRawArgumentEvidence, true); + assert.deepEqual(safety.proofs.get('call-args'), { + name: 'Write', + value: { path: 'a.md' }, + }); + assert.equal(safety.proofs.has('call-zero'), false); + assert.deepEqual(safety.atomicProofs.get('call-zero'), { name: 'ListTodos' }); + }); + + test('an id with any raw delta is never atomic-eligible, even if the bytes never proved', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'Write' }); + observeRawChunk(tracker, { type: 'tool-input-delta', id: 'call-1', delta: '{"path":' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.proofs.has('call-1'), false); // malformed JSON: no raw-byte proof either. + assert.equal(safety.atomicProofs.has('call-1'), false); + }); + + test('zero deltas but missing start fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(resolveToolCallSafety(tracker).atomicProofs.has('call-1'), false); + }); + + test('zero deltas but missing end fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'ListTodos' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(resolveToolCallSafety(tracker).atomicProofs.has('call-1'), false); + }); + + test('duplicate start with zero deltas fails closed', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'ListTodos' }); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-1', toolName: 'ListTodos' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-1' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + assert.equal(resolveToolCallSafety(tracker).atomicProofs.has('call-1'), false); + }); + + test('end observed for a different id than start leaves neither id atomic-proved', () => { + const tracker = createToolCallSafetyTracker(); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-a', toolName: 'ListTodos' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-b' }); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.atomicProofs.has('call-a'), false); // started, never ended. + assert.equal(safety.atomicProofs.has('call-b'), false); // ended without starting: invalid. + }); + + test('tool evidence after a terminal event poisons an otherwise-atomic call too', () => { + const tracker = createToolCallSafetyTracker(); + pushAtomicCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: 'stop' }); + observeRawChunk(tracker, { type: 'tool-input-start', id: 'call-2', toolName: 'ListTodos' }); + observeRawChunk(tracker, { type: 'tool-input-end', id: 'call-2' }); + + const safety = resolveToolCallSafety(tracker); + assert.equal(safety.atomicProofs.has('call-1'), false); + assert.equal(safety.atomicProofs.has('call-2'), false); + }); + + for (const finishReason of ['length', 'content-filter', 'error', 'other', 'unknown']) { + test(`${finishReason} never authorizes an atomic call either`, () => { + const tracker = createToolCallSafetyTracker(); + pushAtomicCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason }); + assert.equal(resolveToolCallSafety(tracker).atomicProofs.has('call-1'), false); + }); + } + + test('an ambiguous finish reason is rescued by providerReason for an atomic call too', () => { + const tracker = createToolCallSafetyTracker(); + pushAtomicCall(tracker); + observeRawChunk(tracker, { type: 'finish', finishReason: { unified: 'other', raw: 'stop' } }); + assert.deepEqual( + resolveToolCallSafety(tracker, { providerReason: 'stop' }).atomicProofs.get('call-1'), + { + name: 'ListTodos', + }, + ); + }); +}); + +describe('isSafeToolExecutionStepOutcome', () => { + const request = {}; + const failure: ModelFailure = { + type: 'model_failure', + kind: 'provider_unavailable', + message: 'boom', + retryable: false, + }; + + for (const finishReason of ['stop', 'tool-calls'] as const) { + test(`allows ${finishReason}`, () => { + const outcome: ModelStepOutcome = { + kind: 'completed', + finishReason, + request, + continuation: 'none', + }; + assert.equal(isSafeToolExecutionStepOutcome(outcome), true); + }); + } + + for (const finishReason of ['length', 'content-filter', 'other', 'unknown'] as const) { + test(`rejects completed/${finishReason}`, () => { + const outcome: ModelStepOutcome = { + kind: 'completed', + finishReason, + request, + continuation: 'none', + }; + assert.equal(isSafeToolExecutionStepOutcome(outcome), false); + }); + } + + for (const outcome of [ + { kind: 'truncated', failure, request, continuation: 'none' }, + { kind: 'terminal-failure', failure, request, continuation: 'none' }, + { + kind: 'retryable-failure', + failure: { ...failure, retryable: true }, + request, + continuation: 'none', + }, + { + kind: 'aborted', + failure: { ...failure, kind: 'abort' as const }, + request, + continuation: 'none', + }, + ] satisfies ModelStepOutcome[]) { + test(`rejects ${outcome.kind}`, () => { + assert.equal(isSafeToolExecutionStepOutcome(outcome), false); + }); + } +}); diff --git a/packages/runtime/src/ai-sdk-backend.ts b/packages/runtime/src/ai-sdk-backend.ts index 48bb09af68..7da77a7269 100644 --- a/packages/runtime/src/ai-sdk-backend.ts +++ b/packages/runtime/src/ai-sdk-backend.ts @@ -119,6 +119,7 @@ import type { NormalizedUsage, ModelFailureKind, ToolCallPart, + ToolCallExecutionSafety, ToolResultOutput, UserContent, } from './model-protocol.js'; @@ -276,6 +277,7 @@ import { type HistoryCompactCheckpoint, } from './history-compact-checkpoint.js'; import { resolveSelectedModelContextWindow } from './context-budget-policy.js'; +import { isSafeToolExecutionStepOutcome } from './tool-call-execution-guard.js'; export { DEFAULT_PERMISSION_TIMEOUT_MS, MAX_ACTIVE_CHILD_AGENT_RUNS_PER_TURN, @@ -2104,6 +2106,17 @@ export class AiSdkBackend implements AgentBackend { let overflowRetryUsed = false; let result: ModelStreamResult; let providerOutcome: ModelStepOutcome; + // Per-tool-call positive-completion proof for the physical provider + // request that produced `returnedToolCalls` below — see + // tool-call-execution-guard.ts. Reset every loop iteration alongside + // `providerOutcome`/`result`: it describes only the most recent + // request, matching `returnedToolCalls`'s own per-request lifetime + // (settled once per iteration, never carried across steps). + let toolCallSafety: ToolCallExecutionSafety = { + hadRawArgumentEvidence: false, + proofs: new Map(), + atomicProofs: new Map(), + }; let finishReason: ModelFinishReason = 'stop'; let terminalProviderError: unknown; agentLoop: for (;;) { @@ -2414,6 +2427,7 @@ export class AiSdkBackend implements AgentBackend { // must not be reported as the already-handled watchdog timeout. const settledWatchdogTimeout = consumeWatchdogTimeout(); providerOutcome = await result.outcome; + toolCallSafety = await result.toolCallSafety; const incompleteStreamTerminal = providerOutcome.kind === 'truncated'; const incompleteStreamHasNoObservableOutput = incompleteStreamTerminal && @@ -2597,7 +2611,120 @@ export class AiSdkBackend implements AgentBackend { `Provider-executed tool call "${toolCall.toolName}" is outside the main-agent tool loop`, ); } - const requestedTool = toolsByName.get(toolCall.toolName); + // A structurally complete tool call is not proof the raw + // stream that produced it was ever confirmed safe: + // `settleModelStepOutcome` treats `finishReason: "length"` as + // `{ kind: 'completed' }` (the same branch `"stop"`/ + // `"tool-calls"` take, for that function's own + // continuation/retry bookkeeping purpose), so without this + // gate a call cut off by a token limit could still reach + // ToolRuntime. See tool-call-execution-guard.ts for the full + // raw-evidence/proof/atomic-proof/whole-request-fallback + // contract this reads, including why a missing proof and a + // failed-verification proof are the same "no entry" case. + // + // A proof (raw-byte OR atomic) is only trusted when its own + // proved tool name agrees with this call's — case- + // insensitively, since `repairMakaToolCall` legitimately + // corrects a mis-cased name (streamed as "WRITE", dispatched + // as "Write") between what the guard observed at + // `tool-input-start` and the final resolved `tool-call`. A + // proved id whose name disagrees even case-insensitively is + // the identity-substitution this gate exists to catch (e.g. + // "Write" proved, "Shell" dispatched), not something to + // execute under either name. The one exception is + // `INVALID_TOOL_NAME`: repair routes an unrepairable call + // there deliberately, and that handler never does more than + // format an error — so it is exempt from the identity check + // but, per the value rule below, never eligible for the + // guard's proved value either, since that value describes the + // ORIGINAL (now-irrelevant) tool call, not the + // repair-synthesized `{tool, error}` payload + // `buildInvalidMakaTool` expects. + // + // An atomic proof is checked only when there is no raw-byte + // proof for this id AND some OTHER call in the same physical + // request did stream real delta bytes + // (`hadRawArgumentEvidence`) — i.e. only inside a genuinely + // mixed-delivery request, the installed Google adapter's own + // shape for one argument-bearing call plus one zero-argument + // sibling. It is a PER-CALL fact: it says nothing about any + // OTHER sibling in the same request, which is what lets the + // zero-argument sibling execute independently of whether the + // argument-bearing one streamed real bytes. + // + // A genuinely whole-request-atomic delivery (no call in the + // request streamed any delta bytes at all — some providers + // never emit granular per-call lifecycle chunks and instead + // hand off a complete, possibly non-empty call in one shot) + // is a separate, pre-existing policy this leaves untouched: + // `atomicProofs` is deliberately NOT consulted when + // `hadRawArgumentEvidence` is false, so that case keeps + // falling through to the unscoped whole-request fallback + // below exactly as it always has, trusting `toolCall.input` + // verbatim with no emptiness requirement. That policy is + // sound specifically because those calls have no + // real-delta-carrying sibling in the request to be confused + // with — the ambiguity this scoping resolves does not arise. + const proof = toolCallSafety.proofs.get(toolCall.toolCallId); + const atomicProof = toolCallSafety.hadRawArgumentEvidence + ? toolCallSafety.atomicProofs.get(toolCall.toolCallId) + : undefined; + const provedNameMatches = + proof !== undefined && + proof.name.toLowerCase() === toolCall.toolName.toLowerCase(); + const atomicNameMatches = + atomicProof !== undefined && + atomicProof.name.toLowerCase() === toolCall.toolName.toLowerCase(); + const confirmedSafe = + proof !== undefined + ? toolCall.toolName === INVALID_TOOL_NAME || provedNameMatches + : atomicProof !== undefined + ? toolCall.toolName === INVALID_TOOL_NAME || atomicNameMatches + : !toolCallSafety.hadRawArgumentEvidence && + isSafeToolExecutionStepOutcome(providerOutcome); + // The guard's own proved value — decoded from this call's + // raw bytes, never the SDK-projected `toolCall.input` a + // later repair/coercion could have substituted — is the + // sole payload authority whenever the proved identity + // genuinely matches the tool about to run. + // + // A mixed-delivery atomic proof (the branch above, only + // reached when `hadRawArgumentEvidence` is true) has no raw + // bytes to decode a value from, but it does not fall back to + // `toolCall.input` either: zero `tool-input-delta` chunks for + // this id, next to a sibling that DID stream real bytes, is + // itself the raw evidence, and what it proves is that the + // provider supplied no arguments at all for this one id + // (confirmed against the installed Google adapter's own + // source — its ONLY zero-delta path is `isNoArgsCompleteCall`, + // gated on the model producing no arguments; its complete-call + // path with real arguments always emits exactly one + // `tool-input-delta` carrying them, so a real zero-delta id + // and a real non-empty id are never the same wire shape from + // that adapter). The canonical empty object is therefore this + // proof's own value, in the same sense `JSON.parse(state.raw)` + // is the raw-byte proof's — never `toolCall.input`, which + // could disagree with that proof (a divergent SDK projection, + // a stale repair, or simply a bug) with nothing here able to + // tell. Schema defaults still apply from here exactly as they + // do for any other call: the guard proves the id's own + // arguments were empty, and `ToolRuntime`'s own schema parsing + // is what may add trusted local defaults/transforms on top of + // that proven value. + // + // The whole-request atomic fallback — reached both when a + // call has no proof of any kind AND `hadRawArgumentEvidence` + // is false (a genuinely all-atomic request, where a + // non-empty `toolCall.input` is an ordinary complete call, + // not a contradiction) — and the invalid-tool bypass above + // have no proved value either and fall back to + // `toolCall.input`, per that separate, pre-existing policy. + const provedValue = provedNameMatches ? proof.value : undefined; + const atomicValue = atomicNameMatches ? {} : undefined; + const requestedTool = confirmedSafe + ? toolsByName.get(toolCall.toolName) + : undefined; const tool = requestedTool ?? toolsByName.get(INVALID_TOOL_NAME); if (!tool) throw new Error('Runtime invalid-tool fallback is unavailable'); return await toolRuntime.settleToolCall({ @@ -2618,10 +2745,16 @@ export class AiSdkBackend implements AgentBackend { : {}), input: requestedTool !== undefined - ? toolCall.input + ? provedValue !== undefined + ? provedValue + : atomicValue !== undefined + ? atomicValue + : toolCall.input : { tool: toolCall.toolName, - error: 'returned tool is unavailable', + error: confirmedSafe + ? 'returned tool is unavailable' + : 'the stream that produced this call was not confirmed to complete safely', }, abortSignal: turnAbortController.signal, eventSink: queue, diff --git a/packages/runtime/src/model-adapter.ts b/packages/runtime/src/model-adapter.ts index a21f89f3e1..b09d6233e7 100644 --- a/packages/runtime/src/model-adapter.ts +++ b/packages/runtime/src/model-adapter.ts @@ -28,6 +28,11 @@ import { lookupModelMetadata } from '@maka/core/model-metadata'; import { generalizedErrorMessage } from '@maka/core/redaction'; import type { CacheMissInputSource } from '@maka/core/usage-stats/types'; import { rawFinishReasonString } from './model-protocol.js'; +import { + createToolCallSafetyTracker, + observeRawChunk, + resolveToolCallSafety, +} from './tool-call-execution-guard.js'; import type { ModelMessage, NormalizedUsage, @@ -41,6 +46,7 @@ import type { ModelRequestMetadata, ModelToolSet, ToolCallPart, + ToolCallExecutionSafety, } from './model-protocol.js'; export type { NormalizedUsage, @@ -331,6 +337,16 @@ export class ModelAdapter { const outcome = new Promise((resolve) => { settleOutcome = resolve; }); + let settleToolCallSafety!: (safety: ToolCallExecutionSafety) => void; + const toolCallSafety = new Promise((resolve) => { + settleToolCallSafety = resolve; + }); + // One tracker per physical provider request (this method's own + // "one physical provider request" contract — see ModelStreamResult) so + // two concurrent requests, even with colliding provider-issued + // toolCallIds, can never share or cross-resolve state. See + // tool-call-execution-guard.ts. + const toolCallGuard = createToolCallSafetyTracker(); const request = { messages: continuation.requestMessages }; const events: AsyncIterable = { async *[Symbol.asyncIterator]() { @@ -340,6 +356,11 @@ export class ModelAdapter { try { for await (const chunk of sdk.stream as AsyncIterable) { onStreamActivity(); + // Raw, unmodified chunk — before translateChunk's semantic + // narrowing — so the guard proves each tool call's argument + // completeness from its own raw byte stream, not from + // translateChunk's already-trusted `tool-call.input`. + observeRawChunk(toolCallGuard, chunk); for (const event of translateChunk(chunk, openAiChatReasoningTransportState)) { if (event.kind === 'error') failure = event.failure; if (event.kind === 'finish') sawFinish = true; @@ -390,12 +411,15 @@ export class ModelAdapter { } } } finally { + settleToolCallSafety( + resolveToolCallSafety(toolCallGuard, { providerReason: finishReason }), + ); settleOutcome(settled); } } }, }; - return { events, outcome }; + return { events, outcome, toolCallSafety }; } endContinuation(lane: string): void { diff --git a/packages/runtime/src/model-protocol.ts b/packages/runtime/src/model-protocol.ts index fe1a1db769..bc079b413d 100644 --- a/packages/runtime/src/model-protocol.ts +++ b/packages/runtime/src/model-protocol.ts @@ -427,8 +427,73 @@ export type ModelStepOutcome = continuation: 'none'; }; +/** + * One tool call's positive completion proof from + * `tool-call-execution-guard.ts`'s raw-byte verification, keyed by + * `toolCallId` in `ToolCallExecutionSafety.proofs`. `name`/`value` are the + * guard's own proof — the tool name and parsed argument value it verified + * from the raw stream — and are authoritative over whatever the AI SDK's + * post-hoc `tool-call` chunk claims for that same id. See + * `tool-call-execution-guard.ts` for the full contract, including why an + * absent entry (no proof) is the only negative representation: there is no + * separate rejected/retry state to consult. + */ +export interface ToolCallSafetyProof { + readonly name: string; + readonly value: unknown; +} + +/** + * One tool call's positive proof of an atomic, zero-delta-byte delivery: + * its own `tool-input-start`/`tool-input-end` lifecycle matched (same id, + * no contradictory evidence), it received zero raw `tool-input-delta` + * chunks, and the request terminated safely — see + * `tool-call-execution-guard.ts`. There is no `value` here (unlike + * `ToolCallSafetyProof`): this tracker never observes the resolved + * `tool-call` chunk, so it has no opinion on what the AI SDK's own + * projected input for this id contains — only that no argument bytes + * streamed for it. + * + * The name is an identity check only. `ai-sdk-backend.ts`'s dispatch is + * what turns this into an execution value: only when `hadRawArgumentEvidence` + * is true (a genuinely mixed-delivery request, e.g. the installed Google + * adapter's `isNoArgsCompleteCall` shape next to an argument-bearing + * sibling) does it consult this proof at all, and when it does, it executes + * the canonical empty object — never the SDK's projected `toolCall.input` — + * since the installed adapter's own source confirms a real zero-delta id and + * a real non-empty id are never the same wire shape, making that projection + * both untrustworthy and unnecessary here. A genuinely whole-request-atomic + * delivery (`hadRawArgumentEvidence` false) is a separate, pre-existing + * policy that still trusts `toolCall.input` verbatim and does not consult + * this map at all. + */ +export interface ToolCallAtomicProof { + readonly name: string; +} + +/** + * The full per-physical-request result of `tool-call-execution-guard.ts`'s + * raw-byte verification. See that file's header for the complete contract: + * `proofs` (non-empty raw-byte proofs, value included), `atomicProofs` + * (per-id zero-delta lifecycle proofs, value deliberately excluded — see + * `ToolCallAtomicProof`), and why a call with neither — falling back to + * `hadRawArgumentEvidence` being false for the whole request plus a safe + * step outcome — is the only remaining, narrowly-scoped fallback. + */ +export interface ToolCallExecutionSafety { + readonly hadRawArgumentEvidence: boolean; + readonly proofs: ReadonlyMap; + readonly atomicProofs: ReadonlyMap; +} + /** One physical provider request: live output plus one authoritative settlement. */ export interface ModelStreamResult { events: AsyncIterable; outcome: Promise; + /** + * Per-tool-call positive-completion proof for this same physical provider + * request — see `tool-call-execution-guard.ts` and `ToolCallExecutionSafety` + * above. + */ + toolCallSafety: Promise; } diff --git a/packages/runtime/src/tool-call-execution-guard.ts b/packages/runtime/src/tool-call-execution-guard.ts new file mode 100644 index 0000000000..8e9b50bbf5 --- /dev/null +++ b/packages/runtime/src/tool-call-execution-guard.ts @@ -0,0 +1,397 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * tool-call-execution-guard.ts — positive-completion proof for a provider + * step's tool calls before `ai-sdk-backend.ts` settles them through + * `ToolRuntime` (real side effects: filesystem writes, shell commands, + * apply_patch, SQL execution, dependency installs, and so on). + * + * Why this exists: `settleModelStepOutcome` (model-adapter.ts) classifies a + * step as `{ kind: 'completed' }` for `finishReason: "length"` — the same + * branch `"stop"` and `"tool-calls"` take. `ai-sdk-backend.ts`'s historical + * gate before settling `returnedToolCalls` only asked for that broader + * completed outcome, so a step that streamed a structurally complete tool + * call and was then cut off by a token limit while producing more content + * could still execute that call. The `length -> completed` classification is + * correct for that function's continuation/retry bookkeeping purpose and is + * deliberately not reused as an irreversible tool-execution gate here. + * + * `translateChunk` (model-adapter.ts) also does not establish execution + * authority for a final `'tool-call'` chunk's `input`: that value is the AI + * SDK's parsed/post-processed projection and may have passed through + * `repairToolCall` or another coercion path. Where the provider exposes raw + * `tool-input-delta` chunks, those bytes are stronger evidence of what the + * provider actually streamed. This tracker observes the raw SDK chunks + * verbatim before translation, keeps the argument bytes per physical request, + * requires matching start/end evidence plus a positively safe terminal event, + * and parses only the captured raw JSON. The resulting raw-stream tool name + * and parsed value are the execution authority for that call. + * + * Some provider-facing chunk sequences (the installed Google adapter's among + * them) may contain `tool-input-start` immediately followed by + * `tool-input-end` with zero `tool-input-delta` chunks, for a tool that + * genuinely takes no arguments — with the (empty) arguments present only in + * the trailing `tool-call` projection. That shape is legitimate on its own + * and, per call, indistinguishable from a truncated/mismatched stream only + * in the absence of its own start/end lifecycle — which this tracker does + * observe. The focused guard and production-path regression suites exercise + * it directly (`tool-call-execution-guard.test.ts`, + * `length-cutoff-tool-execution-repro.test.ts`, and `ai-sdk-backend.test.ts`), + * including the case where an argument-bearing sibling call shares the same + * physical request. + * + * `resolveToolCallSafety` represents only positive proof, in two disjoint + * maps: + * + * - `proofs` holds an entry for an id if and only if every condition below + * held — matching start, non-empty raw bytes, matching end, unpoisoned + * ordering, a raw-stream tool name, JSON that parses from exactly those + * bytes, and an execution-safe terminal classification. Its `value` is the + * guard's own parsed argument value. + * - `atomicProofs` holds an entry for an id whose OWN lifecycle positively + * proved an atomic, genuinely zero-argument delivery: matching start, + * matching end, unpoisoned ordering, a raw-stream tool name, zero raw + * `tool-input-delta` chunks for that id specifically, and an + * execution-safe terminal classification. It carries no `value` — there + * were no raw bytes to parse — only the name the tracker itself observed, + * for the caller's identity check. + * + * An id can appear in at most one of the two maps (an id that received any + * raw delta bytes is only ever eligible for `proofs`, never `atomicProofs`, + * even if those bytes failed to produce a proof). An id in neither map had + * no positive per-id evidence at all — a mismatched, out-of-order, or wholly + * unobserved lifecycle — and a caller must treat that the same as a call + * that never streamed anything: there is no separate rejected/retry state to + * distinguish "evidence existed but failed a condition" from "no evidence at + * all", because a caller can never act on that difference anyway. + * + * `atomicProofs` is deliberately silent on VALUE, and that is the caller's + * responsibility to get right, not this tracker's: this tracker only ever + * observes `tool-input-start`/`tool-input-delta`/`tool-input-end`, never the + * resolved `tool-call` chunk itself, so it has no way to know — and does not + * claim to know — what the SDK's own projected `toolCall.input` for that id + * turned out to be. A complete, unpoisoned, zero-delta id/name lifecycle + * proves the provider streamed no argument bytes; it does not, by itself, + * prove what the SDK-resolved input contains. The caller (see + * `ai-sdk-backend.ts`) closes that gap by additionally requiring + * `hadRawArgumentEvidence` to be true before it will consult `atomicProofs` + * at all, and by treating the canonical empty object — never + * `toolCall.input` — as the executed value once it does: see that file for + * why (the installed Google adapter's own source confirms a real zero-delta + * id and a real non-empty id are never the same wire shape, so the SDK's + * projection is not needed and is not trusted for this branch). + * + * A caller resolves a call with neither a `proofs` nor an (its own + * `hadRawArgumentEvidence`-gated) `atomicProofs` entry by falling back to + * "genuinely atomic; use the step-level fallback, trusting `toolCall.input` + * verbatim" ONLY when `hadRawArgumentEvidence` is false for the WHOLE + * physical request — i.e. this tracker observed no `tool-input-delta` bytes + * anywhere in the request, meaning either every call in it is legitimately + * atomic (a provider that hands off complete, possibly non-empty calls in + * one shot, with no incremental streaming at all) or the provider's protocol + * never emits granular per-call lifecycle chunks in the first place. That is + * a separate, pre-existing policy this tracker does not change. The moment + * any call anywhere in the request streamed real bytes, a call with no + * `proofs` entry AND no `atomicProofs` entry is indistinguishable from an id + * mismatch between this tracker's raw-chunk view and the SDK's resolved + * `tool-call`; that case must still fail closed. A call with a `proofs` + * entry of its own, or an `atomicProofs` entry the caller is willing to + * trust, never needs this fallback — which is exactly what fixed the case a + * purely request-global rule got wrong: an argument-bearing call and a + * legitimate zero-argument sibling in the same request, each proved from its + * own lifecycle, neither one's fate decided by the other's. + * + * A call with a present `proofs` entry gets the raw-stream name/value as + * execution authority. Neither case ever defers to the SDK's post-hoc + * projection for a DIFFERENT id, and neither lets one call's raw evidence + * stand in for another's. + * + * Concurrency note: nothing here is shared across requests or stored beyond + * one `ModelAdapter.startStream` result. `createToolCallSafetyTracker()` owns + * fresh maps and terminal state every time, so two concurrent provider + * requests that both use a provider-issued id like `"call_1"` can never + * observe or resolve each other's evidence. + * + * `isSafeToolExecutionStepOutcome` below is the fallback used for a call with + * no per-id evidence at all (see the whole-request fallback above). It is + * deliberately narrower than `settleModelStepOutcome`'s own + * `kind === 'completed'` — that classification also covers `"length"`, + * which is correct for that function's bookkeeping purpose but is not an + * execution-safe outcome. This helper exists only for the tool-execution + * gate; it does not change `settleModelStepOutcome` itself or anything else + * that reads `ModelStepOutcome`. + */ +import type { + ModelStepOutcome, + ToolCallAtomicProof, + ToolCallExecutionSafety, + ToolCallSafetyProof, +} from './model-protocol.js'; + +interface RawToolCallState { + name?: string; + raw: string; + started: boolean; + ended: boolean; + invalid: boolean; +} + +/** + * `pending`: no terminal stream event observed yet. + * `finish`: an actual `finish` chunk was observed; `reason` is this + * tracker's own local read of it (`.unified` only — see + * `normalizedFinishReason`), which may be ambiguous (`"other"`/`"unknown"`/ + * `undefined`) even when the provider's own spelling, available elsewhere, + * is not. See `isTerminalSafe`. + * `blocked`: an explicit `error`/`abort` part, or tool-call evidence + * arriving after any terminal event already happened (a poisoned request). + * Sticky: nothing can move a tracker out of `blocked`. + */ +type TerminalState = + | { readonly kind: 'pending' } + | { readonly kind: 'finish'; readonly reason: string | undefined } + | { readonly kind: 'blocked' }; + +export interface ToolCallSafetyTracker { + /** Per-id raw argument state for this physical provider request only. */ + readonly calls: Map; + /** ids that received at least one non-empty tool-input-delta chunk. */ + readonly idsWithRawDelta: Set; + terminal: TerminalState; + resolved?: ToolCallExecutionSafety; +} + +function toolPartId(part: { id?: unknown; toolCallId?: unknown }): string | undefined { + if (typeof part.id === 'string') return part.id; + if (typeof part.toolCallId === 'string') return part.toolCallId; + return undefined; +} + +function callState(tracker: ToolCallSafetyTracker, id: string): RawToolCallState { + let state = tracker.calls.get(id); + if (state === undefined) { + state = { raw: '', started: false, ended: false, invalid: false }; + tracker.calls.set(id, state); + } + return state; +} + +function normalizedFinishReason(value: unknown): string | undefined { + if (typeof value === 'string') return value; + if ( + value !== null && + typeof value === 'object' && + typeof (value as { unified?: unknown }).unified === 'string' + ) { + return (value as { unified: string }).unified; + } + return undefined; +} + +function blockTerminal(tracker: ToolCallSafetyTracker): void { + tracker.terminal = { kind: 'blocked' }; +} + +/** Starts tracking one physical provider request's raw stream. */ +export function createToolCallSafetyTracker(): ToolCallSafetyTracker { + return { calls: new Map(), idsWithRawDelta: new Set(), terminal: { kind: 'pending' } }; +} + +/** Feed one raw AI SDK stream chunk through, unchanged, as observed. */ +export function observeRawChunk(tracker: ToolCallSafetyTracker, chunk: unknown): void { + if (chunk === null || typeof chunk !== 'object') return; + const part = chunk as { + type?: unknown; + id?: unknown; + toolCallId?: unknown; + toolName?: unknown; + delta?: unknown; + finishReason?: unknown; + }; + + const isToolEvidence = + part.type === 'tool-input-start' || + part.type === 'tool-input-delta' || + part.type === 'tool-input-end' || + part.type === 'tool-error'; + if (tracker.terminal.kind !== 'pending' && isToolEvidence) blockTerminal(tracker); + + switch (part.type) { + case 'tool-input-start': { + const id = toolPartId(part); + if (id === undefined) return; + const state = callState(tracker, id); + if (state.started || state.ended || state.raw.length > 0) state.invalid = true; + state.started = true; + if (typeof part.toolName === 'string' && part.toolName.length > 0) { + if (state.name !== undefined && state.name !== part.toolName) state.invalid = true; + state.name = part.toolName; + } + return; + } + case 'tool-input-delta': { + const id = toolPartId(part); + if (id === undefined || typeof part.delta !== 'string' || part.delta.length === 0) return; + const state = callState(tracker, id); + if (!state.started || state.ended) state.invalid = true; + state.raw += part.delta; + tracker.idsWithRawDelta.add(id); + return; + } + case 'tool-input-end': { + const id = toolPartId(part); + if (id === undefined) return; + const state = callState(tracker, id); + if (!state.started || state.ended) state.invalid = true; + state.ended = true; + return; + } + case 'tool-error': { + const id = toolPartId(part); + if (id !== undefined) callState(tracker, id).invalid = true; + return; + } + case 'finish': { + // A second finish-shaped event is exactly as suspicious as tool + // evidence after a terminal event — treat it the same way (blocked), + // rather than letting a later finish silently replace an earlier one. + if (tracker.terminal.kind !== 'pending') { + blockTerminal(tracker); + return; + } + tracker.terminal = { kind: 'finish', reason: normalizedFinishReason(part.finishReason) }; + return; + } + case 'error': + case 'abort': + blockTerminal(tracker); + return; + default: + return; + } +} + +/** + * Whether this request's terminal state is execution-safe, given the + * stronger provider reason `resolveToolCallSafety`'s caller already resolved + * (see below). Mirrors `chunkFinishReason`'s (model-adapter.ts) own rule — + * fall back to the provider's own spelling only when the SDK's unified + * reason is ambiguous (`"other"`/`"unknown"`) — without importing it + * directly: `tool-call-execution-guard.ts` is imported by `model-adapter.ts`, + * so the reverse import would cycle. + * + * `providerReason` may only resolve an AMBIGUOUS classification belonging to + * a real terminal event this tracker itself witnessed. It can never promote + * `pending` (no terminal event observed at all — see the module doc comment + * for why `sdk.finishReason`'s own fallback means a caller can have a + * non-empty `providerReason` even then) or `blocked` (explicit error/abort, + * or tool evidence poisoning the request) to safe, and it can never override + * a terminal event this tracker directly classified as unsafe on its own + * (`length`, `content-filter`, `stop`-with-no-safe-match, etc.) — only an + * ambiguous one. + */ +function isTerminalSafe(terminal: TerminalState, providerReason: string | undefined): boolean { + if (terminal.kind !== 'finish') return false; + if (terminal.reason === 'stop' || terminal.reason === 'tool-calls') return true; + const ambiguous = + terminal.reason === undefined || terminal.reason === 'other' || terminal.reason === 'unknown'; + if (!ambiguous) return false; + return providerReason === 'stop' || providerReason === 'tool-calls'; +} + +/** + * Resolves every call this tracker has per-id evidence for into a positive + * proof (raw-byte or atomic) or no entry at all — see the module doc comment + * for the full `proofs`/`atomicProofs` contract. `meta.providerReason` is the + * same finish reason `ModelAdapter.startStream` resolves for step settlement + * (`chunkFinishReason`, model-adapter.ts) — passing it lets an ambiguous + * local classification agree with the stronger, already-computed answer + * instead of the two diverging for the same physical request. + */ +export function resolveToolCallSafety( + tracker: ToolCallSafetyTracker, + meta?: { providerReason?: string }, +): ToolCallExecutionSafety { + if (tracker.resolved !== undefined) return tracker.resolved; + + const terminalSafe = isTerminalSafe(tracker.terminal, meta?.providerReason); + const proofs = new Map(); + for (const id of tracker.idsWithRawDelta) { + const state = tracker.calls.get(id); + if ( + !terminalSafe || + state === undefined || + !state.started || + !state.ended || + state.invalid || + state.name === undefined + ) { + continue; + } + + try { + const value: unknown = JSON.parse(state.raw); + proofs.set(id, { name: state.name, value }); + } catch { + // No proof: the captured bytes did not parse as JSON. + } + } + + // A call with its own matching start + end, zero raw delta bytes, and no + // contradictory evidence has a complete, unambiguous lifecycle even though + // it never carries argument bytes to parse — this is the shape the + // installed Google adapter legitimately produces for a zero-argument tool + // call. `idsWithRawDelta` and this loop are disjoint by construction: an id + // that received any delta bytes is only ever eligible for `proofs` above, + // even when those bytes failed to produce one. + const atomicProofs = new Map(); + for (const [id, state] of tracker.calls) { + if (tracker.idsWithRawDelta.has(id)) continue; + if ( + !terminalSafe || + !state.started || + !state.ended || + state.invalid || + state.name === undefined + ) { + continue; + } + atomicProofs.set(id, { name: state.name }); + } + + tracker.resolved = { + hadRawArgumentEvidence: tracker.idsWithRawDelta.size > 0, + proofs, + atomicProofs, + }; + return tracker.resolved; +} + +/** + * Whether an all-atomic provider step's own termination is execution-safe. + * Only an unambiguous normal completion qualifies. `length` is intentionally + * excluded even though `settleModelStepOutcome` classifies it as completed. + */ +export function isSafeToolExecutionStepOutcome(outcome: ModelStepOutcome): boolean { + return ( + outcome.kind === 'completed' && + (outcome.finishReason === 'stop' || outcome.finishReason === 'tool-calls') + ); +} diff --git a/packages/runtime/src/tool-runtime.ts b/packages/runtime/src/tool-runtime.ts index f78e13240c..bd8b55a678 100644 --- a/packages/runtime/src/tool-runtime.ts +++ b/packages/runtime/src/tool-runtime.ts @@ -906,7 +906,13 @@ export class ToolRuntime { ? `Tool ${tool.name} is direct-only and cannot run inside exec.` : undefined; const admissionFailure = directOnlyFailure ?? this.admitToolForStep(tool, stepId); - const executionArgs = rawExecutionArgs; + // Reassigned below to the schema's own parsed/projected value (defaults + // filled in, transforms applied) once validation succeeds — everything + // downstream (permission check, persisted tool_start/tool_call args, + // loop-gate signature, and `tool.impl` itself) reads this one variable, + // so that reassignment is the single point that makes schema semantics + // reach execution instead of being discarded after a validate-only check. + let executionArgs = rawExecutionArgs; let permissionArgs = executionArgs; let permissionArgsError: unknown; if (directOnlyFailure === undefined) { @@ -919,7 +925,9 @@ export class ToolRuntime { !this.interactionRun() && (!this.input.createSandboxBoundaryRequest || !this.input.settleSandboxBoundaryRequest); if (!sandboxBoundaryUnavailable) { - await validateDeclaredToolArgs(tool.parameters, rawExecutionArgs); + executionArgs = snapshotToolArgs( + await validateDeclaredToolArgs(tool.parameters, rawExecutionArgs), + ); } permissionArgs = tool.permissionArgs ? snapshotToolArgs( @@ -2507,9 +2515,20 @@ export class ToolRuntime { } } -async function validateDeclaredToolArgs(parameters: unknown, args: unknown): Promise { +/** + * Validates `args` against the tool's declared schema and returns the + * schema's own output — defaults filled in, `transform`/`preprocess` + * applied — rather than the pre-validation input. The caller (`executeTool`) + * uses this returned value, not `args`, as the execution/persisted/ + * permission-check arguments from this point on: a schema that validates but + * discards its own parsed result silently drops `z.default()`/`.transform()` + * output before it ever reaches `tool.impl`. Schemas exposing none of the + * recognized validator interfaces (or no schema at all) return `args` + * unchanged, matching the pre-existing no-op behavior for those tools. + */ +async function validateDeclaredToolArgs(parameters: unknown, args: unknown): Promise { if (!parameters || (typeof parameters !== 'object' && typeof parameters !== 'function')) { - return; + return args; } const schema = parameters as { safeParseAsync?: ( @@ -2536,24 +2555,25 @@ async function validateDeclaredToolArgs(parameters: unknown, args: unknown): Pro if (typeof schema.safeParseAsync === 'function') { const parsed = await schema.safeParseAsync(args); - if (parsed.success) return; + if (parsed.success) return parsed.data; throw parsed.error; } if (typeof schema.safeParse === 'function') { const parsed = schema.safeParse(args); - if (parsed.success) return; + if (parsed.success) return parsed.data; throw parsed.error; } if (typeof schema.validate === 'function') { const parsed = await schema.validate(args); - if (parsed.success) return; + if (parsed.success) return parsed.value; throw parsed.error; } if (typeof schema['~standard']?.validate === 'function') { const parsed = await schema['~standard'].validate(args); - if ('value' in parsed) return; + if ('value' in parsed) return parsed.value; throw new Error('Tool arguments failed declared schema validation', { cause: parsed.issues }); } + return args; } function isInteractionControlError(error: unknown): boolean {