import * as aiModule from 'ai';
import type { JSONSchema7 } from 'json-schema';
import type { Mock, MockedFunction } from 'vitest';
import { z } from 'zod';

import { InMemoryFilesystem } from '../../__tests__/workspace/test-utils';
import { Agent } from '../../sdk/agent';
import { createCancellation } from '../../sdk/cancellation';
import { isLlmMessage } from '../../sdk/message';
import { Tool, Tool as ToolBuilder } from '../../sdk/tool';
import type { CheckpointStore, SerializableAgentState } from '../../types';
import { AgentEvent } from '../../types/runtime/event';
import type { AgentEventData } from '../../types/runtime/event';
import type { StreamChunk } from '../../types/sdk/agent';
import type { AgentDbMessage, ContentToolCall, Message } from '../../types/sdk/message';
import type { BuiltTool, InterruptibleToolContext, ToolContext } from '../../types/sdk/tool';
import type { BuiltTelemetry } from '../../types/telemetry';
import { Workspace, getToolResultRunDirectory } from '../../workspace';
import { AgentRuntime } from '../loop/agent-runtime';
import { InMemoryMemory } from '../memory/memory-store';
import { OBSERVATION_CONTINUATION_REMINDER } from '../model/message-list';
import { AgentEventBus } from '../state/event-bus';
import { StaleResumeError } from '../state/run-state';
import {
	DELEGATE_SUB_AGENT_TOOL_NAME,
	INLINE_DELEGATE_SUB_AGENT_TOOL_METADATA_KEY,
	createDelegateSubAgentTool,
	type DelegateSubAgentRunner,
} from '../tools/delegate-sub-agent-tool';
import { toAiSdkTools } from '../tools/tool-adapter';
import { MAX_MODEL_TOOL_RESULT_TOKENS } from '../tools/tool-result-guard';

// ---------------------------------------------------------------------------
// Module mocks
// ---------------------------------------------------------------------------

// Mock provider packages so createModel() doesn't fail when no API key is set
vi.mock('@ai-sdk/openai', () => ({
	createOpenAI: () =>
		Object.assign(() => ({ provider: 'openai', modelId: 'mock', specificationVersion: 'v3' }), {
			embeddingModel: () => ({ provider: 'openai', modelId: 'mock', specificationVersion: 'v2' }),
		}),
}));

vi.mock('@ai-sdk/anthropic', () => ({
	createAnthropic: () => () => ({
		provider: 'anthropic',
		modelId: 'mock',
		specificationVersion: 'v3',
	}),
}));

// eslint-disable-next-line @typescript-eslint/consistent-type-imports
type AiImport = typeof import('ai');

// Mock generateText and streamText from the 'ai' package
vi.mock('ai', async () => {
	const actual = await vi.importActual<AiImport>('ai');
	return {
		...actual,
		embed: vi.fn(),
		embedMany: vi.fn(),
		generateText: vi.fn(),
		streamText: vi.fn(),
		tool: vi.fn((config: unknown) => config),
		jsonSchema: vi.fn((schema: unknown) => ({ _type: 'jsonSchema', schema })),
		Output: {
			object: vi.fn(({ schema }: { schema: unknown }) => ({ _type: 'object', schema })),
		},
	};
});

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const { embed, embedMany, generateText, streamText, jsonSchema } = aiModule as unknown as {
	embed: Mock;
	embedMany: Mock;
	generateText: Mock;
	streamText: Mock;
	jsonSchema: Mock;
};

const makeCancellationSpy = () =>
	vi.fn<NonNullable<BuiltTool['onCancellation']>>().mockResolvedValue(undefined);

/** Minimal successful generateText response. */
function makeGenerateSuccess(text = 'OK') {
	return {
		finishReason: 'stop',
		usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
		response: {
			messages: [
				{
					role: 'assistant',
					content: [{ type: 'text', text }],
				},
			],
		},
		toolCalls: [],
	};
}

/** Generator that yields the given chunks then ends cleanly. */
function* makeChunkStream(
	chunks: Array<Record<string, unknown>>,
): Generator<Record<string, unknown>> {
	for (const c of chunks) {
		yield c;
	}
}

/**
 * Create a pre-rejected promise whose rejection is already handled, so it
 * never causes an unhandled-rejection crash even if nobody awaits it.
 */
// eslint-disable-next-line @typescript-eslint/promise-function-async
function silentReject<T = never>(error: Error): Promise<T> {
	const p = Promise.reject<T>(error);
	p.catch(() => {}); // suppress unhandled-rejection warning
	return p;
}

function makeErrorStream(error: Error) {
	return new ReadableStream({
		start(controller) {
			controller.error(error);
		},
	});
}

function makeExecutionCounter() {
	return {
		incrementMessageCount: vi.fn(),
		incrementToolCallCount: vi.fn(),
		incrementTokenCount: vi.fn(),
	};
}

/** Collect all chunks from a ReadableStream. */
async function collectChunks(stream: ReadableStream<unknown>): Promise<StreamChunk[]> {
	const chunks: StreamChunk[] = [];
	const reader = stream.getReader();
	while (true) {
		const { done, value } = await reader.read();
		if (done) break;
		chunks.push(value as StreamChunk);
	}
	return chunks;
}

/** Minimal successful streamText response. */
function makeStreamSuccess(text = 'Hello') {
	return {
		stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text }]),
		finishReason: Promise.resolve('stop'),
		usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
		response: Promise.resolve({
			messages: [{ role: 'assistant', content: [{ type: 'text', text }] }],
		}),
		toolCalls: Promise.resolve([]),
	};
}

/**
 * streamText response where the model invokes a provider-executed tool (e.g.
 * native web search): the SDK streams a `tool-call` and its terminal part
 * (`tool-result` on success, `tool-error` on failure) with `providerExecuted`,
 * then finishes with `stop` (the provider runs the tool server-side mid-step).
 */
function makeStreamWithProviderTool(opts: {
	toolCallId: string;
	toolName: string;
	input: unknown;
	output?: unknown;
	error?: unknown;
	text?: string;
}) {
	const terminal =
		opts.error !== undefined
			? {
					type: 'tool-error',
					toolCallId: opts.toolCallId,
					toolName: opts.toolName,
					input: opts.input,
					error: opts.error,
					providerExecuted: true,
				}
			: {
					type: 'tool-result',
					toolCallId: opts.toolCallId,
					toolName: opts.toolName,
					input: opts.input,
					output: opts.output,
					providerExecuted: true,
				};
	const text = opts.text ?? 'done';
	return {
		stream: makeChunkStream([
			{
				type: 'tool-call',
				toolCallId: opts.toolCallId,
				toolName: opts.toolName,
				input: opts.input,
				providerExecuted: true,
			},
			terminal,
			{ type: 'text-delta', id: 'text-1', text },
		]),
		finishReason: Promise.resolve('stop'),
		usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
		response: Promise.resolve({
			messages: [{ role: 'assistant', content: [{ type: 'text', text }] }],
		}),
		toolCalls: Promise.resolve([]),
	};
}

/** Build a default runtime wired to the shared eventBus for inspection. */
function createRuntime(eventBus?: AgentEventBus, model = 'openai/gpt-4o-mini') {
	const bus = eventBus ?? new AgentEventBus();
	const runtime = new AgentRuntime({
		name: 'test',
		model,
		instructions: 'You are a test assistant.',
		eventBus: bus,
	});
	return { runtime, bus };
}

const testSchema = z.object({ answer: z.string(), score: z.number() });

const testJsonSchema: JSONSchema7 = {
	type: 'object',
	properties: { answer: { type: 'string' }, score: { type: 'number' } },
	required: ['answer', 'score'],
};

/** Build a runtime configured with structuredOutput. */
function createStructuredRuntime(options?: { tools?: BuiltTool[]; eventBus?: AgentEventBus }) {
	const bus = options?.eventBus ?? new AgentEventBus();
	const runtime = new AgentRuntime({
		name: 'test-structured',
		model: 'openai/gpt-4o-mini',
		instructions: 'You are a test assistant.',
		structuredOutput: testSchema,
		eventBus: bus,
		...(options?.tools ? { tools: options.tools } : {}),
	});
	return { runtime, bus };
}

/** Build a runtime configured with a raw JSON Schema structuredOutput. */
function createJsonSchemaStructuredRuntime() {
	const bus = new AgentEventBus();
	const runtime = new AgentRuntime({
		name: 'test-structured-json',
		model: 'openai/gpt-4o-mini',
		instructions: 'You are a test assistant.',
		structuredOutput: testJsonSchema,
		eventBus: bus,
	});
	return { runtime, bus };
}

/** Minimal successful generateText response with structured output. */
function makeGenerateSuccessWithOutput(output: unknown, text = 'OK') {
	return { ...makeGenerateSuccess(text), output };
}

/** Minimal successful streamText response with structured output. */
function makeStreamSuccessWithOutput(output: unknown, text = 'Hello') {
	return { ...makeStreamSuccess(text), output: Promise.resolve(output) };
}

/** Mock generateText response that triggers a tool call. */
function makeGenerateWithToolCall(toolCallId: string, toolName: string, input: unknown) {
	return {
		finishReason: 'tool-calls',
		usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
		response: {
			messages: [
				{
					role: 'assistant',
					content: [{ type: 'tool-call', toolCallId, toolName, args: input, input }],
				},
			],
		},
		toolCalls: [{ toolCallId, toolName, input }],
	};
}

function getFirstReplayedToolCallInput(callIndex: number) {
	const callArgs = generateText.mock.calls[callIndex][0] as {
		messages: Array<{ role: string; content: unknown }>;
	};
	const assistantMessage = callArgs.messages.find((message) => message.role === 'assistant');
	if (!assistantMessage || !Array.isArray(assistantMessage.content)) return undefined;

	const toolCall = assistantMessage.content.find(
		(part): part is { type: 'tool-call'; input?: unknown } =>
			typeof part === 'object' && part !== null && Reflect.get(part, 'type') === 'tool-call',
	);
	return toolCall?.input;
}

/** Build an interruptible tool that suspends on first call and returns on resume. */
function makeInterruptibleTool(): BuiltTool {
	return {
		name: 'approve',
		description: 'Requires approval',
		inputSchema: z.object({ question: z.string() }),
		suspendSchema: z.object({ question: z.string() }),
		resumeSchema: z.object({ approved: z.boolean() }),
		handler: async (_input: unknown, ctx: unknown) => {
			const { suspend, resumeData } = ctx as InterruptibleToolContext;
			if (!resumeData) {
				return await suspend({ question: 'approve?' });
			}
			return { approved: true };
		},
	};
}

// ---------------------------------------------------------------------------
// execution counters
// ---------------------------------------------------------------------------

describe('AgentRuntime — execution counters', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	it('counts one fresh generate turn and token usage', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const counter = makeExecutionCounter();

		const { runtime } = createRuntime();
		await runtime.generate('hi', { executionCounter: counter });

		expect(counter.incrementMessageCount).toHaveBeenCalledTimes(1);
		expect(counter.incrementToolCallCount).not.toHaveBeenCalled();
		expect(counter.incrementTokenCount).toHaveBeenCalledWith(15);
	});

	it('counts one fresh stream turn and streamed token usage', async () => {
		streamText.mockReturnValue(makeStreamSuccess());
		const counter = makeExecutionCounter();

		const { runtime } = createRuntime();
		const result = await runtime.stream('hi', { executionCounter: counter });
		await collectChunks(result.stream);

		expect(counter.incrementMessageCount).toHaveBeenCalledTimes(1);
		expect(counter.incrementToolCallCount).not.toHaveBeenCalled();
		expect(counter.incrementTokenCount).toHaveBeenCalledWith(15);
	});

	it('forwards onStepStart and onStepEnd to generateText and streamText', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		streamText.mockReturnValue(makeStreamSuccess());
		const onStepStart = vi.fn();
		const onStepEnd = vi.fn();

		const { runtime } = createRuntime();
		await runtime.generate('hi', { onStepStart, onStepEnd });
		const streamResult = await runtime.stream('hi', { onStepStart, onStepEnd });
		await collectChunks(streamResult.stream);

		for (const call of generateText.mock.calls) {
			const args = call[0] as Record<string, unknown>;
			expect(args.onStepStart).toBe(onStepStart);
			expect(args.onStepEnd).toBe(onStepEnd);
		}
		for (const call of streamText.mock.calls) {
			const args = call[0] as Record<string, unknown>;
			expect(args.onStepStart).toBe(onStepStart);
			expect(args.onStepEnd).toBe(onStepEnd);
		}
	});

	it('allows system-role messages in generateText and streamText history', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		streamText.mockReturnValue(makeStreamSuccess());
		const input: Message[] = [
			{ role: 'system', content: [{ type: 'text', text: 'Historical instruction' }] },
			{ role: 'user', content: [{ type: 'text', text: 'Hello' }] },
		];

		const { runtime } = createRuntime();
		await runtime.generate(input);
		const streamResult = await runtime.stream(input);
		await collectChunks(streamResult.stream);

		expect(generateText.mock.calls[0]?.[0]).toEqual(
			expect.objectContaining({ allowSystemInMessages: true }),
		);
		expect(streamText.mock.calls[0]?.[0]).toEqual(
			expect.objectContaining({ allowSystemInMessages: true }),
		);
	});

	it('counts provider-executed tool calls when surfaced by the model', async () => {
		generateText
			.mockResolvedValueOnce({
				...makeGenerateWithToolCall('tc-provider', 'openai.web_search', { query: 'n8n' }),
				toolCalls: [
					{
						toolCallId: 'tc-provider',
						toolName: 'openai.web_search',
						input: { query: 'n8n' },
						providerExecuted: true,
					},
				],
			})
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));
		const counter = makeExecutionCounter();

		const { runtime } = createRuntime();
		await runtime.generate('search', { executionCounter: counter });

		expect(counter.incrementToolCallCount).toHaveBeenCalledTimes(1);
	});

	it('does not count a resumed suspended tool call as a new tool invocation', async () => {
		const suspendTool = makeInterruptibleTool();
		const counter = makeExecutionCounter();
		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'approve', args: { question: 'continue?' } },
			]),
		);

		const first = await runtime.generate('needs approval', { executionCounter: counter });
		const { runId, toolCallId } = first.pendingSuspend![0];

		generateText.mockResolvedValueOnce(makeGenerateSuccess('approved'));
		await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId, executionCounter: counter },
		);

		expect(counter.incrementMessageCount).toHaveBeenCalledTimes(1);
		expect(counter.incrementToolCallCount).toHaveBeenCalledTimes(1);
		expect(counter.incrementTokenCount).toHaveBeenCalledTimes(2);
	});

	it('exposes the checkpointed suspend payload to a resumed tool handler via ctx.suspendPayload', async () => {
		const suspendPayload = { question: 'approve?', marker: 'xyz' };
		let observedCtx: InterruptibleToolContext | undefined;
		const suspendTool: BuiltTool = {
			name: 'approve',
			description: 'Requires approval',
			inputSchema: z.object({ question: z.string() }),
			suspendSchema: z.object({ question: z.string(), marker: z.string() }),
			resumeSchema: z.object({ approved: z.boolean() }),
			handler: async (_input: unknown, ctx: unknown) => {
				const interruptibleCtx = ctx as InterruptibleToolContext;
				if (!interruptibleCtx.resumeData) {
					return await interruptibleCtx.suspend(suspendPayload);
				}
				observedCtx = interruptibleCtx;
				return { approved: true };
			},
		};

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'approve', args: { question: 'continue?' } },
			]),
		);

		const first = await runtime.generate('needs approval');
		const { runId, toolCallId } = first.pendingSuspend![0];

		generateText.mockResolvedValueOnce(makeGenerateSuccess('approved'));
		await runtime.resume('generate', { approved: true }, { runId, toolCallId });

		expect(observedCtx?.suspendPayload).toEqual(suspendPayload);
	});

	it('persists dynamic suspension state privately and validates resume before claiming', async () => {
		const continuation = { childRunId: 'child-run-1', taskPath: '/root/research_0' };
		const dynamicResumeSchema = z.object({ decision: z.literal('continue') });
		let observedCtx: InterruptibleToolContext | undefined;
		const suspendTool: BuiltTool = {
			name: 'dynamic_suspend',
			description: 'Suspends with invocation-specific state',
			inputSchema: z.object({}),
			suspendSchema: z.object({ prompt: z.string() }),
			resumeSchema: z.object({ fallback: z.string() }),
			handler: async (_input, ctx) => {
				const interruptibleCtx = ctx as InterruptibleToolContext;
				if (!interruptibleCtx.resumeData) {
					return await interruptibleCtx.suspend(
						{ prompt: 'Continue the child run?' },
						{ resumeSchema: dynamicResumeSchema, continuation },
					);
				}
				observedCtx = interruptibleCtx;
				return { resumed: true };
			},
		};
		const checkpointStore = makeClaimingCheckpointStore();
		const runtime = createRuntimeWithCheckpointStore([suspendTool], checkpointStore);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-dynamic', 'dynamic_suspend', {}),
		);

		const first = await runtime.generate('start');
		const publicSuspension = first.pendingSuspend?.[0];
		expect(publicSuspension?.resumeSchema).toEqual(
			expect.objectContaining({
				type: 'object',
				properties: { decision: expect.objectContaining({ const: 'continue' }) },
				required: ['decision'],
			}),
		);
		expect(publicSuspension).not.toHaveProperty('continuation');

		const checkpoint = await checkpointStore.load(first.runId);
		expect(checkpoint?.pendingToolCalls['tc-dynamic']).toEqual(
			expect.objectContaining({
				continuation,
				resumeSchema: publicSuspension?.resumeSchema,
			}),
		);

		await expect(
			runtime.resume(
				'generate',
				{ fallback: 'accepted only by the static schema' },
				{ runId: first.runId, toolCallId: 'tc-dynamic' },
			),
		).rejects.toThrow('Invalid resume payload');
		expect(checkpointStore.claimForResume).not.toHaveBeenCalled();

		generateText.mockResolvedValueOnce(makeGenerateSuccess('resumed'));
		await runtime.resume(
			'generate',
			{ decision: 'continue' },
			{ runId: first.runId, toolCallId: 'tc-dynamic' },
		);

		expect(observedCtx?.continuation).toEqual(continuation);
		expect(observedCtx?.resumeSchema).toEqual(publicSuspension?.resumeSchema);
	});

	it('keeps delegate_subagent output usage per tool call without adding it to generate result usage', async () => {
		const delegateTool: BuiltTool = {
			name: DELEGATE_SUB_AGENT_TOOL_NAME,
			description: 'Delegate work',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () =>
				await Promise.resolve({
					status: 'completed',
					answer: 'child answer',
					usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7, cost: 0.01 },
				}),
		};
		const counter = makeExecutionCounter();
		const { runtime } = createRuntimeWithTools([delegateTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-delegate',
						toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
						args: { value: 'research' },
					},
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const result = await runtime.generate('delegate', { executionCounter: counter });

		expect(result.usage).toMatchObject({
			promptTokens: 20,
			completionTokens: 10,
			totalTokens: 30,
		});
		expect(result.toolCalls?.[0]?.output).toEqual(
			expect.objectContaining({
				usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7, cost: 0.01 },
			}),
		);
		expect(counter.incrementTokenCount).toHaveBeenCalledTimes(2);
		expect(counter.incrementTokenCount).toHaveBeenNthCalledWith(1, 15);
		expect(counter.incrementTokenCount).toHaveBeenNthCalledWith(2, 15);
	});

	it('does not roll normal tool output usage-like fields into generate result usage', async () => {
		const normalTool: BuiltTool = {
			name: 'normal_tool',
			description: 'Normal tool',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () =>
				await Promise.resolve({
					usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7 },
				}),
		};
		const { runtime } = createRuntimeWithTools([normalTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-normal', toolName: 'normal_tool', args: { value: 'run' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const result = await runtime.generate('use normal tool');

		expect(result.usage).toMatchObject({
			promptTokens: 20,
			completionTokens: 10,
			totalTokens: 30,
		});
	});

	it('keeps delegate_subagent output usage per stream tool result without adding it to finish usage', async () => {
		const delegateTool: BuiltTool = {
			name: DELEGATE_SUB_AGENT_TOOL_NAME,
			description: 'Delegate work',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () =>
				await Promise.resolve({
					status: 'completed',
					answer: 'child answer',
					usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7, cost: 0.01 },
				}),
		};
		const { runtime } = createRuntimeWithTools([delegateTool], 1);

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-delegate',
									toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
									args: { value: 'research' },
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{
						toolCallId: 'tc-delegate',
						toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
						input: { value: 'research' },
					},
				]),
			})
			.mockReturnValueOnce(makeStreamSuccess('done'));

		const result = await runtime.stream('delegate');
		const chunks = await collectChunks(result.stream);
		const finishChunks = chunks.filter((chunk) => chunk.type === 'finish');
		const finish = finishChunks[finishChunks.length - 1] as
			| (StreamChunk & { type: 'finish' })
			| undefined;
		const toolResult = chunks.find(
			(chunk) => chunk.type === 'tool-result' && chunk.toolName === DELEGATE_SUB_AGENT_TOOL_NAME,
		) as (StreamChunk & { type: 'tool-result' }) | undefined;

		expect(finish?.usage).toMatchObject({
			promptTokens: 20,
			completionTokens: 10,
			totalTokens: 30,
		});
		expect(toolResult?.output).toEqual(
			expect.objectContaining({
				usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7, cost: 0.01 },
			}),
		);
	});
});

// ---------------------------------------------------------------------------
// empty stop turn retry
// ---------------------------------------------------------------------------

describe('AgentRuntime — empty stop turn retry', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	/** A `stop` turn with no output at all — the provider stall the loop retries. */
	function makeGenerateEmpty() {
		return {
			finishReason: 'stop',
			usage: { inputTokens: 10, outputTokens: 0, totalTokens: 10 },
			response: { messages: [] },
			toolCalls: [],
		};
	}

	it('retries an empty stop turn and keeps the retry output', async () => {
		generateText
			.mockResolvedValueOnce(makeGenerateEmpty())
			.mockResolvedValueOnce(makeGenerateSuccess('Recovered'));

		const { runtime } = createRuntime();
		const result = await runtime.generate('hi');

		expect(generateText).toHaveBeenCalledTimes(2);
		expect(result.finishReason).toBe('stop');
		const assistantText = result.messages
			.flatMap((m) => ('content' in m && Array.isArray(m.content) ? m.content : []))
			.find((c) => c.type === 'text');
		expect(assistantText).toMatchObject({ text: 'Recovered' });
	});

	it('bills usage from discarded empty attempts', async () => {
		generateText
			.mockResolvedValueOnce(makeGenerateEmpty())
			.mockResolvedValueOnce(makeGenerateSuccess('Recovered'));
		const counter = makeExecutionCounter();

		const { runtime } = createRuntime();
		await runtime.generate('hi', { executionCounter: counter });

		expect(counter.incrementTokenCount).toHaveBeenCalledWith(10);
		expect(counter.incrementTokenCount).toHaveBeenCalledWith(15);
	});

	it('accepts the empty turn once the retry budget is exhausted', async () => {
		generateText.mockResolvedValue(makeGenerateEmpty());

		const { runtime } = createRuntime();
		const result = await runtime.generate('hi');

		// 1 initial call + 2 retries
		expect(generateText).toHaveBeenCalledTimes(3);
		expect(result.finishReason).toBe('stop');
		expect(result.error).toBeUndefined();
	});

	it('does not retry a whitespace-free stop turn with text', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Done'));

		const { runtime } = createRuntime();
		await runtime.generate('hi');

		expect(generateText).toHaveBeenCalledTimes(1);
	});
});

// ---------------------------------------------------------------------------
// reasoning-only turn
// ---------------------------------------------------------------------------

describe('AgentRuntime — reasoning-only turn', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	/**
	 * The provider streamed thinking and then the stream ended without a
	 * terminal chunk, so the SDK synthesizes the finish from its defaults: a
	 * reasoning-only message, `other`, and no usage. Nothing the user can see
	 * and nothing the loop can act on.
	 */
	function makeStreamReasoningOnly(finishReason = 'other') {
		return {
			stream: makeChunkStream([{ type: 'reasoning-delta', id: 'r-1', text: 'thinking...' }]),
			finishReason: Promise.resolve(finishReason),
			usage: Promise.resolve({
				inputTokens: undefined,
				outputTokens: undefined,
				totalTokens: undefined,
			}),
			response: Promise.resolve({
				messages: [{ role: 'assistant', content: [{ type: 'reasoning', text: 'thinking...' }] }],
			}),
			toolCalls: Promise.resolve([]),
		};
	}

	it('retries a reasoning-only turn instead of ending the run', async () => {
		streamText
			.mockReturnValueOnce(makeStreamReasoningOnly())
			.mockReturnValueOnce(makeStreamSuccess('Recovered'));

		const { runtime } = createRuntime();
		const result = await runtime.stream('make my workflow smarter');
		const chunks = await collectChunks(result.stream);

		expect(streamText).toHaveBeenCalledTimes(2);
		const text = chunks
			.filter((chunk): chunk is StreamChunk & { type: 'text-delta' } => chunk.type === 'text-delta')
			.map((chunk) => chunk.delta)
			.join('');
		expect(text).toBe('Recovered');
	});

	it('reports an error rather than finishing silently when every attempt is reasoning-only', async () => {
		streamText.mockImplementation(() => makeStreamReasoningOnly());

		const { runtime } = createRuntime();
		const result = await runtime.stream('make my workflow smarter');
		const chunks = await collectChunks(result.stream);

		const finish = chunks.find(
			(chunk): chunk is StreamChunk & { type: 'finish' } => chunk.type === 'finish',
		);
		expect(finish?.finishReason).toBe('error');
		const error = chunks.find(
			(chunk): chunk is StreamChunk & { type: 'error' } => chunk.type === 'error',
		);
		expect(String((error?.error as Error).message)).toContain('no output');
	});

	it.each([
		{ finishReason: 'length', errorText: 'output token limit' },
		{ finishReason: 'stop', errorText: 'without returning an answer' },
	])(
		'fails a reasoning-only $finishReason turn without retrying or persisting it',
		async ({ finishReason, errorText }) => {
			streamText.mockReturnValue(makeStreamReasoningOnly(finishReason));
			const memory = new InMemoryMemory();
			const runtime = new AgentRuntime({
				name: 'test',
				model: 'anthropic/claude-opus-5',
				instructions: 'You are a test assistant.',
				memory,
			});

			const result = await runtime.stream('make my workflow smarter', {
				persistence: { threadId: 'thread-1', resourceId: 'user-1' },
			});
			const chunks = await collectChunks(result.stream);

			expect(streamText).toHaveBeenCalledTimes(1);
			expect(chunks).toEqual(
				expect.arrayContaining([
					expect.objectContaining({ type: 'error' }),
					expect.objectContaining({ type: 'finish', finishReason: 'error' }),
				]),
			);
			const error = chunks.find((chunk) => chunk.type === 'error');
			expect(String(error?.error)).toContain(errorText);
			const persisted = await memory.getMessages('thread-1', { resourceId: 'user-1' });
			expect(
				persisted.filter((message) => (message as { role?: string }).role === 'assistant'),
			).toHaveLength(0);
		},
	);
});

// ---------------------------------------------------------------------------
// generate() — graceful error contract
// ---------------------------------------------------------------------------

describe('AgentRuntime.generate() — graceful error contract', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('resolves (never rejects) when the LLM call throws', async () => {
		generateText.mockRejectedValue(new Error('API failure'));

		const { runtime } = createRuntime();
		await expect(runtime.generate('hello')).resolves.toBeDefined();
	});

	it('returns finishReason "error" when the LLM call throws', async () => {
		generateText.mockRejectedValue(new Error('API failure'));

		const { runtime } = createRuntime();
		const result = await runtime.generate('hello');

		expect(result.finishReason).toBe('error');
	});

	it('surfaces the original error in result.error when the LLM call throws', async () => {
		const cause = new Error('API failure');
		generateText.mockRejectedValue(cause);

		const { runtime } = createRuntime();
		const result = await runtime.generate('hello');

		expect(result.error).toBe(cause);
	});

	it('sets state to "failed" when the LLM call throws', async () => {
		generateText.mockRejectedValue(new Error('API failure'));

		const { runtime } = createRuntime();
		await runtime.generate('hello');

		expect(runtime.getState().status).toBe('failed');
	});

	it('emits AgentEvent.Error (not AgentEnd) when the LLM call throws', async () => {
		generateText.mockRejectedValue(new Error('API failure'));

		const { runtime, bus } = createRuntime();
		const errorEvents: unknown[] = [];
		const endEvents: unknown[] = [];
		bus.on(AgentEvent.Error, (d) => errorEvents.push(d));
		bus.on(AgentEvent.AgentEnd, (d) => endEvents.push(d));

		await runtime.generate('hello');

		expect(errorEvents.length).toBe(1);
		expect(endEvents.length).toBe(0);
	});

	it('does NOT emit AgentEvent.Error for abort — only sets cancelled state', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const { runtime, bus } = createRuntime();
		const errorEvents: unknown[] = [];
		bus.on(AgentEvent.Error, (d) => errorEvents.push(d));

		// Abort during AgentStart so the loop's first abort-check fires before generateText is called
		bus.on(AgentEvent.AgentStart, () => bus.abort());

		await runtime.generate('hello');

		expect(errorEvents.length).toBe(0);
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('returns finishReason "error" and sets cancelled status on abort', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const { runtime, bus } = createRuntime();
		// Abort during AgentStart so the loop's first abort-check fires before generateText is called
		bus.on(AgentEvent.AgentStart, () => bus.abort());

		const result = await runtime.generate('hello');

		expect(result.finishReason).toBe('error');
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('is reusable after an error — subsequent call with a good LLM response succeeds', async () => {
		generateText
			.mockRejectedValueOnce(new Error('transient error'))
			.mockResolvedValue(makeGenerateSuccess());

		const { runtime } = createRuntime();

		const first = await runtime.generate('hello');
		expect(first.finishReason).toBe('error');

		const second = await runtime.generate('hello');
		expect(second.finishReason).toBe('stop');
		expect(second.error).toBeUndefined();
	});

	it('returns an empty messages array (no partial data) when abort fires before the first LLM call', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('first turn response'));

		const { runtime, bus } = createRuntime();
		bus.on(AgentEvent.AgentStart, () => bus.abort());

		const result = await runtime.generate('hello');
		// Abort fired before any LLM response was accumulated
		expect(result.finishReason).toBe('error');
		expect(result.messages).toEqual([]);
	});

	it('succeeds normally when there is no error', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Great answer'));

		const { runtime } = createRuntime();
		const result = await runtime.generate('hello');

		expect(result.finishReason).toBe('stop');
		expect(result.error).toBeUndefined();
		expect(result.messages.length).toBeGreaterThan(0);
	});
});

// ---------------------------------------------------------------------------
// eager input persistence — input survives an uncompleted turn
// ---------------------------------------------------------------------------

describe('AgentRuntime — eager input persistence', () => {
	const PERSIST = { persistence: { threadId: 't1', resourceId: 'u1' } };

	function textOf(message: AgentDbMessage): string | undefined {
		const content = (message as { content?: unknown }).content;
		if (!Array.isArray(content)) return undefined;
		const block = content.find(
			(c): c is { type: 'text'; text: string } =>
				typeof c === 'object' && c !== null && Reflect.get(c, 'type') === 'text',
		);
		return block?.text;
	}

	function hasToolCall(message: AgentDbMessage, toolCallId: string): boolean {
		const content = (message as { content?: unknown }).content;
		if (!Array.isArray(content)) return false;
		return content.some(
			(c) =>
				typeof c === 'object' &&
				c !== null &&
				Reflect.get(c, 'type') === 'tool-call' &&
				Reflect.get(c, 'toolCallId') === toolCallId,
		);
	}

	function makeMemoryRuntime() {
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			memory,
		});
		return { runtime, bus, memory };
	}

	it('persists the user input when the turn aborts before completion', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const { runtime, bus, memory } = makeMemoryRuntime();
		// Abort during AgentStart: the input is added (and eagerly persisted) in
		// buildMessageList, but the loop's first abort-check throws before finishComplete.
		bus.on(AgentEvent.AgentStart, () => bus.abort());

		const result = await runtime.generate('remember this prompt', PERSIST);

		expect(result.finishReason).toBe('error');
		expect(runtime.getState().status).toBe('cancelled');
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		expect(persisted.map(textOf)).toContain('remember this prompt');
	});

	it('stores exactly one user-message row on normal completion (idempotent)', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('done'));
		const { runtime, memory } = makeMemoryRuntime();

		await runtime.generate('hello once', PERSIST);

		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const userRows = persisted.filter(
			(m) => (m as { role?: string }).role === 'user' && textOf(m) === 'hello once',
		);
		expect(userRows).toHaveLength(1);
	});

	it('does not duplicate the user row across a suspend -> resume cycle', async () => {
		const suspendTool = makeInterruptibleTool();
		const memory = new InMemoryMemory();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [suspendTool],
			checkpointStorage: 'memory',
			memory,
		});

		// Turn 1: the model calls the interruptible tool, so the run suspends on HITL
		// before reaching finishComplete.
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'approve', args: { question: 'continue?' } },
			]),
		);
		const first = await runtime.generate('please build it', PERSIST);
		const { runId, toolCallId } = first.pendingSuspend![0];

		// Eager persist wrote the input even though the turn has not completed.
		const afterSuspend = await memory.getMessages('t1', { resourceId: 'u1' });
		expect(afterSuspend.filter((m) => textOf(m) === 'please build it')).toHaveLength(1);

		// Resume to completion: the end-of-turn save writes the full turnDelta, which
		// includes the input again (same id) — it must upsert, not create a 2nd row.
		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		await runtime.resume('generate', { approved: true }, { runId, toolCallId });

		const afterResume = await memory.getMessages('t1', { resourceId: 'u1' });
		expect(afterResume.filter((m) => textOf(m) === 'please build it')).toHaveLength(1);
	});

	it("persists the turn's assistant output when the turn suspends on HITL", async () => {
		const suspendTool = makeInterruptibleTool();
		const memory = new InMemoryMemory();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [suspendTool],
			checkpointStorage: 'memory',
			memory,
		});

		// The model calls the interruptible tool, so the run suspends on HITL before
		// reaching finishComplete — the only path that previously saved the turn.
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'approve', args: { question: 'continue?' } },
			]),
		);

		await runtime.generate('please build it', PERSIST);

		// The assistant's tool-call output from the suspended turn is now in memory, so a
		// later cancel/abandon cannot drop it — the next turn sees the work was done.
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const assistantRows = persisted.filter(
			(m) => (m as { role?: string }).role === 'assistant' && hasToolCall(m, 'tc-1'),
		);
		expect(assistantRows).toHaveLength(1);
	});

	it('stores exactly one assistant row across a suspend -> resume cycle (idempotent)', async () => {
		const suspendTool = makeInterruptibleTool();
		const memory = new InMemoryMemory();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [suspendTool],
			checkpointStorage: 'memory',
			memory,
		});

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'approve', args: { question: 'continue?' } },
			]),
		);
		const first = await runtime.generate('please build it', PERSIST);
		const { runId, toolCallId } = first.pendingSuspend![0];

		// Saved on suspend (pending tool-call).
		const afterSuspend = await memory.getMessages('t1', { resourceId: 'u1' });
		expect(afterSuspend.filter((m) => hasToolCall(m, 'tc-1'))).toHaveLength(1);

		// Resume to completion: the end-of-turn save rewrites the same assistant row (now
		// with the resolved tool-call) under the same id — upsert, not a 2nd row.
		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		await runtime.resume('generate', { approved: true }, { runId, toolCallId });

		const afterResume = await memory.getMessages('t1', { resourceId: 'u1' });
		expect(afterResume.filter((m) => hasToolCall(m, 'tc-1'))).toHaveLength(1);
	});

	it("persists the turn's assistant output when an active turn is aborted (generate)", async () => {
		const lookupTool: BuiltTool = {
			name: 'lookup',
			description: 'looks something up',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve({ ok: true }),
		};
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [lookupTool],
			eventBus: bus,
			memory,
		});

		// The model issues a tool call, so the assistant message is added to the list
		// before the tool runs. Aborting once the tool starts mirrors a user hitting
		// "stop" mid-run — after work was done, before the turn could complete.
		generateText.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'lookup', {}));
		bus.on(AgentEvent.ToolExecutionStart, () => bus.abort());

		const result = await runtime.generate('please build it', PERSIST);

		expect(result.finishReason).toBe('error');
		expect(runtime.getState().status).toBe('cancelled');
		// The assistant work done before the stop is now in memory, so the next turn sees it.
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const assistantRows = persisted.filter(
			(m) => (m as { role?: string }).role === 'assistant' && hasToolCall(m, 'tc-1'),
		);
		expect(assistantRows).toHaveLength(1);
	});

	it("persists the turn's assistant output when an active turn is aborted (stream)", async () => {
		const lookupTool: BuiltTool = {
			name: 'lookup',
			description: 'looks something up',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve({ ok: true }),
		};
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [lookupTool],
			eventBus: bus,
			memory,
		});

		// Same shape as the generate case, on the streaming path (the one Instance AI uses):
		// the abort is handled by the StreamSession, which must persist the turn-so-far.
		streamText.mockReturnValueOnce({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'on it...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [{ type: 'tool-call', toolCallId: 'tc-1', toolName: 'lookup', args: {} }],
					},
				],
			}),
			toolCalls: Promise.resolve([{ toolCallId: 'tc-1', toolName: 'lookup', input: {} }]),
		});
		bus.on(AgentEvent.ToolExecutionStart, () => bus.abort());

		const { stream } = await runtime.stream('please build it', PERSIST);
		await collectChunks(stream);

		expect(runtime.getState().status).toBe('cancelled');
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const assistantRows = persisted.filter(
			(m) => (m as { role?: string }).role === 'assistant' && hasToolCall(m, 'tc-1'),
		);
		expect(assistantRows).toHaveLength(1);
	});

	it('persists the text already streamed when a run is stopped mid-response', async () => {
		const { runtime, memory } = makeMemoryRuntime();
		const controller = new AbortController();
		const abortError = Object.assign(new Error('This operation was aborted'), {
			name: 'AbortError',
		});

		// The model streams two text deltas, then the stop lands — before the turn
		// completes, so its `newMessages` are never built (finishReason/response reject).
		streamText.mockReturnValueOnce({
			stream: (async function* () {
				yield { type: 'text-delta', id: 't1', text: 'Here is my ' };
				yield { type: 'text-delta', id: 't1', text: 'partial answer' };
				await Promise.resolve();
				controller.abort();
			})(),
			finishReason: silentReject(abortError),
			usage: silentReject(abortError),
			response: silentReject(abortError),
			toolCalls: silentReject(abortError),
		});

		const { stream } = await runtime.stream('hello', {
			...PERSIST,
			abortSignal: controller.signal,
		});
		await collectChunks(stream);

		expect(runtime.getState().status).toBe('cancelled');
		// The text the user already saw is in memory, not lost with the aborted turn.
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const assistantTexts = persisted
			.filter((m) => (m as { role?: string }).role === 'assistant')
			.map((m) => textOf(m));
		expect(assistantTexts).toContain('Here is my partial answer');
	});

	it('persists a completed streamed turn when the stop lands before it is folded in', async () => {
		const { runtime, memory } = makeMemoryRuntime();
		const controller = new AbortController();

		// The model finishes streaming (finishReason/response resolve), but the stop
		// fires as the stream closes — before the loop folds the turn into the list. The
		// completed text must still be recovered (the clear is deferred until the fold).
		streamText.mockReturnValueOnce({
			stream: (async function* () {
				yield { type: 'text-delta', id: 't1', text: 'Complete answer' };
				await Promise.resolve();
				controller.abort();
			})(),
			finishReason: Promise.resolve('stop'),
			usage: Promise.resolve({ inputTokens: 5, outputTokens: 3, totalTokens: 8 }),
			response: Promise.resolve({
				messages: [{ role: 'assistant', content: [{ type: 'text', text: 'Complete answer' }] }],
			}),
			toolCalls: Promise.resolve([]),
		});

		const { stream } = await runtime.stream('hi', { ...PERSIST, abortSignal: controller.signal });
		await collectChunks(stream);

		expect(runtime.getState().status).toBe('cancelled');
		const persisted = await memory.getMessages('t1', { resourceId: 'u1' });
		const assistantTexts = persisted
			.filter((m) => (m as { role?: string }).role === 'assistant')
			.map((m) => textOf(m));
		expect(assistantTexts).toContain('Complete answer');
	});
});

// ---------------------------------------------------------------------------
// stream() — fallback error observability
// ---------------------------------------------------------------------------

describe('AgentRuntime.stream() — fallback error observability', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('emits AgentEvent.Error and sets failed state when post-loop persistence throws', async () => {
		const memory = {
			getMessages: vi.fn().mockResolvedValue([]),
			saveMessages: vi.fn().mockRejectedValue(new Error('memory write failed')),
		} as unknown as InMemoryMemory;
		const { runtime, bus } = createRuntime();
		const errorEvents: unknown[] = [];
		bus.on(AgentEvent.Error, (event) => errorEvents.push(event));

		(runtime as unknown as { config: { memory: unknown } }).config.memory = memory;
		streamText.mockReturnValue(makeStreamSuccess('done'));

		const result = await runtime.stream('hello', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		const chunks = await collectChunks(result.stream);

		expect(chunks).toContainEqual(
			expect.objectContaining({
				type: 'finish',
				finishReason: 'error',
			}),
		);
		// The post-loop persistence failure surfaces as a sourceless error event. (The
		// eager input persist also fails against this memory mock and emits its own
		// source-tagged event, which has dedicated coverage elsewhere.)
		expect(errorEvents.filter((e) => !(e as { source?: string }).source)).toHaveLength(1);
		expect(runtime.getState().status).toBe('failed');
	});
});

// ---------------------------------------------------------------------------
// stream() — usage billing on abort
// ---------------------------------------------------------------------------

describe('AgentRuntime.stream() — usage billing on abort', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('emits the accumulated token usage in the terminal finish chunk when aborted after a model turn', async () => {
		const { runtime, bus } = createRuntime();
		streamText.mockReturnValue(makeStreamSuccess('partial'));

		// Stop the run as soon as the turn begins; the model call still resolves
		// usage before the loop's abort-check fires, mirroring a user clicking
		// "stop" right after the model produced output.
		bus.on(AgentEvent.TurnStart, () => bus.abort());

		const { stream } = await runtime.stream('hello');
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		expect(finish?.usage).toMatchObject({
			promptTokens: 10,
			completionTokens: 5,
			totalTokens: 15,
		});
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('bills a discarded empty stop attempt when aborted during the retry', async () => {
		const { runtime } = createRuntime();
		const controller = new AbortController();
		const abortError = Object.assign(new Error('This operation was aborted'), {
			name: 'AbortError',
		});

		streamText
			// Empty stop turn — usage must be published before the retry starts.
			.mockReturnValueOnce({
				stream: makeChunkStream([]),
				finishReason: Promise.resolve('stop'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 0, totalTokens: 10 }),
				response: Promise.resolve({ messages: [] }),
				toolCalls: Promise.resolve([]),
			})
			// Retry starts a new raw-usage reader; abort mid-retry must still bill
			// the discarded empty attempt via the usage already reported to the sink.
			.mockReturnValueOnce({
				stream: (function* () {
					controller.abort();
					yield* makeChunkStream([]);
				})(),
				finishReason: silentReject(abortError),
				usage: silentReject(abortError),
				response: silentReject(abortError),
				toolCalls: silentReject(abortError),
			});

		const { stream } = await runtime.stream('hello', { abortSignal: controller.signal });
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		expect(streamText).toHaveBeenCalledTimes(2);
		expect(finish?.usage).toMatchObject({
			promptTokens: 10,
			completionTokens: 0,
			totalTokens: 10,
		});
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('bills a discarded empty stop attempt when the retry fails with a non-abort error', async () => {
		const { runtime } = createRuntime();
		const streamError = new Error('stream broke');

		streamText
			// Empty stop turn — its usage is reported to the sink before the retry.
			.mockReturnValueOnce({
				stream: makeChunkStream([]),
				finishReason: Promise.resolve('stop'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 0, totalTokens: 10 }),
				response: Promise.resolve({ messages: [] }),
				toolCalls: Promise.resolve([]),
			})
			// Retry fails outright; the error finish chunk must still carry the
			// discarded empty attempt's usage.
			.mockReturnValueOnce({
				stream: makeErrorStream(streamError),
				finishReason: silentReject(streamError),
				usage: silentReject(streamError),
				response: silentReject(streamError),
				toolCalls: silentReject(streamError),
			});

		const { stream } = await runtime.stream('hello');
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish'; finishReason: string })
			| undefined;

		expect(streamText).toHaveBeenCalledTimes(2);
		expect(finish?.finishReason).toBe('error');
		expect(finish?.usage).toMatchObject({
			promptTokens: 10,
			completionTokens: 0,
			totalTokens: 10,
		});
		expect(runtime.getState().status).toBe('failed');
	});

	it('recovers usage from the raw provider stream when aborted before the model finishes', async () => {
		// Anthropic model id: raw `message_start`/`message_delta` events are
		// Anthropic-shaped, so the run must resolve the Anthropic raw-usage reader.
		const { runtime } = createRuntime(undefined, 'anthropic/claude-sonnet-4-6');
		const controller = new AbortController();
		const abortError = Object.assign(new Error('This operation was aborted'), {
			name: 'AbortError',
		});

		// The SDK exposes no usage on abort (finish/usage promises reject), but the
		// provider's `message_start` raw event already carried the input/cache +
		// initial output tokens. The run must bill those.
		streamText.mockReturnValue({
			stream: (async function* () {
				yield {
					type: 'raw',
					rawValue: {
						type: 'message_start',
						message: {
							usage: {
								input_tokens: 5,
								cache_creation_input_tokens: 100,
								cache_read_input_tokens: 10,
								output_tokens: 7,
							},
						},
					},
				};
				await Promise.resolve();
				controller.abort();
			})(),
			finishReason: silentReject(abortError),
			usage: silentReject(abortError),
			response: silentReject(abortError),
			toolCalls: silentReject(abortError),
		});

		const { stream } = await runtime.stream('hello', {
			abortSignal: controller.signal,
			recoverUsageOnAbort: true,
		});
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		// input_tokens(5) + cache_creation(100) + cache_read(10) = 115 prompt tokens.
		expect(finish?.usage).toMatchObject({
			promptTokens: 115,
			completionTokens: 7,
			totalTokens: 122,
			inputTokenDetails: { noCache: 5, cacheRead: 10, cacheWrite: 100 },
		});
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('adds the in-flight turn usage to the completed turns when aborted mid second turn', async () => {
		const tool = makeMockTool('lookup', async () => await Promise.resolve({ ok: true }));
		// Anthropic model id so the in-flight turn's Anthropic raw event is read.
		const { runtime } = createRuntimeWithTools([tool], 1, undefined, 'anthropic/claude-sonnet-4-6');
		const controller = new AbortController();
		const abortError = Object.assign(new Error('This operation was aborted'), {
			name: 'AbortError',
		});

		streamText
			// Turn 1 completes with a tool call; its usage is folded into the total.
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{ type: 'tool-call', toolCallId: 'tc-1', toolName: 'lookup', args: { value: 'x' } },
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-1', toolName: 'lookup', input: { value: 'x' } },
				]),
			})
			// Turn 2 aborts mid-stream after the provider reported its raw usage; the
			// SDK promises reject so the only signal is the captured raw event.
			.mockReturnValueOnce({
				stream: (async function* () {
					yield {
						type: 'raw',
						rawValue: {
							type: 'message_start',
							message: { usage: { input_tokens: 20, output_tokens: 3 } },
						},
					};
					await Promise.resolve();
					controller.abort();
				})(),
				finishReason: silentReject(abortError),
				usage: silentReject(abortError),
				response: silentReject(abortError),
				toolCalls: silentReject(abortError),
			});

		const { stream } = await runtime.stream('hello', {
			abortSignal: controller.signal,
			recoverUsageOnAbort: true,
		});
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		// Turn 1 (10/5) + in-flight turn 2 (20/3). The `??` bug would bill turn 1 only.
		expect(finish?.usage).toMatchObject({
			promptTokens: 30,
			completionTokens: 8,
			totalTokens: 38,
			inputTokenDetails: { noCache: 20 },
		});
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('does not double-count a completed turn when aborted between turns', async () => {
		const tool = makeMockTool('lookup', async () => {
			// Abort after the first turn folded its usage but before a second turn runs.
			controller.abort();
			return await Promise.resolve({ ok: true });
		});
		// Anthropic model id so the reader is created and the clear-on-fold path is
		// genuinely exercised (otherwise no reader exists and the guard is vacuous).
		const { runtime } = createRuntimeWithTools([tool], 1, undefined, 'anthropic/claude-sonnet-4-6');
		const controller = new AbortController();

		// Single turn: emits a raw usage event (captured) AND resolves its final
		// usage. The raw capture must not be re-added on top of the folded total.
		streamText.mockReturnValueOnce({
			stream: makeChunkStream([
				{
					type: 'raw',
					rawValue: {
						type: 'message_start',
						message: { usage: { input_tokens: 20, output_tokens: 3 } },
					},
				},
			]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 20, outputTokens: 3, totalTokens: 23 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{ type: 'tool-call', toolCallId: 'tc-1', toolName: 'lookup', args: { value: 'x' } },
						],
					},
				],
			}),
			toolCalls: Promise.resolve([
				{ toolCallId: 'tc-1', toolName: 'lookup', input: { value: 'x' } },
			]),
		});

		const { stream } = await runtime.stream('hello', {
			abortSignal: controller.signal,
			recoverUsageOnAbort: true,
		});
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		// Exactly the folded turn-1 usage — not doubled by the stale raw capture.
		expect(finish?.usage).toMatchObject({ promptTokens: 20, completionTokens: 3, totalTokens: 23 });
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('does not recover in-flight usage on abort when recoverUsageOnAbort is off (default)', async () => {
		// Anthropic model + raw event present, but the option is omitted: the run
		// must NOT recover the in-flight turn's usage from raw chunks.
		const { runtime } = createRuntime(undefined, 'anthropic/claude-sonnet-4-6');
		const controller = new AbortController();
		const abortError = Object.assign(new Error('This operation was aborted'), {
			name: 'AbortError',
		});

		streamText.mockReturnValue({
			// Sync generator: `for await` consumes it fine and there's nothing to await.
			stream: (function* () {
				yield {
					type: 'raw',
					rawValue: {
						type: 'message_start',
						message: { usage: { input_tokens: 5, output_tokens: 7 } },
					},
				};
				controller.abort();
			})(),
			finishReason: silentReject(abortError),
			usage: silentReject(abortError),
			response: silentReject(abortError),
			toolCalls: silentReject(abortError),
		});

		// No recoverUsageOnAbort: the option stays off (the behavior under test).
		const { stream } = await runtime.stream('hello', { abortSignal: controller.signal });
		const chunks = await collectChunks(stream);

		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish' })
			| undefined;

		// No completed turns and no raw recovery → no usage billed.
		expect(finish?.usage).toBeUndefined();
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('requests raw chunks only when recoverUsageOnAbort is set', async () => {
		streamText.mockReturnValue(makeStreamSuccess('ok'));

		const off = createRuntime(undefined, 'anthropic/claude-sonnet-4-6');
		await collectChunks((await off.runtime.stream('hello')).stream);
		expect(streamText.mock.calls.at(-1)?.[0]).not.toHaveProperty('include.rawChunks');

		streamText.mockClear();
		streamText.mockReturnValue(makeStreamSuccess('ok'));

		const on = createRuntime(undefined, 'anthropic/claude-sonnet-4-6');
		await collectChunks((await on.runtime.stream('hello', { recoverUsageOnAbort: true })).stream);
		expect(streamText.mock.calls.at(-1)?.[0]).toHaveProperty('include.rawChunks', true);
	});
});

// ---------------------------------------------------------------------------
// stream() — graceful error contract
// ---------------------------------------------------------------------------

describe('AgentRuntime.stream() — graceful error contract', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('resolves (never rejects) when the LLM stream throws', async () => {
		const streamError = new Error('stream broke');
		streamText.mockReturnValue({
			stream: makeErrorStream(streamError),
			finishReason: silentReject(streamError),
			usage: Promise.resolve(undefined),
			response: silentReject(streamError),
			toolCalls: Promise.resolve([]),
		});

		const { runtime } = createRuntime();
		await expect(runtime.stream('hello')).resolves.toBeDefined();
	});

	it('yields an error chunk when the LLM stream throws', async () => {
		const streamError = new Error('stream broke');
		streamText.mockReturnValue({
			stream: makeErrorStream(streamError),
			finishReason: silentReject(streamError),
			usage: Promise.resolve(undefined),
			response: silentReject(streamError),
			toolCalls: Promise.resolve([]),
		});

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		const errorChunks = chunks.filter((c) => c.type === 'error');
		expect(errorChunks.length).toBeGreaterThan(0);
	});

	it('closes the stream cleanly (with a finish chunk) after an error', async () => {
		const streamError = new Error('stream broke');
		streamText.mockReturnValue({
			stream: makeErrorStream(streamError),
			finishReason: silentReject(streamError),
			usage: Promise.resolve(undefined),
			response: silentReject(streamError),
			toolCalls: Promise.resolve([]),
		});

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		const finishChunks = chunks.filter((c) => c.type === 'finish');
		expect(finishChunks.length).toBeGreaterThan(0);

		const finish = finishChunks[finishChunks.length - 1] as StreamChunk & {
			type: 'finish';
			finishReason: string;
		};
		expect(finish.finishReason).toBe('error');
	});

	it('sets state to "failed" after a stream error', async () => {
		const streamError = new Error('stream broke');
		streamText.mockReturnValue({
			stream: makeErrorStream(streamError),
			finishReason: silentReject(streamError),
			usage: Promise.resolve(undefined),
			response: silentReject(streamError),
			toolCalls: Promise.resolve([]),
		});

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		await collectChunks(readableStream);

		expect(runtime.getState().status).toBe('failed');
	});

	it('yields error chunk and finishes cleanly on abort', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const { runtime, bus } = createRuntime();
		bus.on(AgentEvent.TurnStart, () => bus.abort());

		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		const errorChunks = chunks.filter((c) => c.type === 'error');
		expect(errorChunks.length).toBeGreaterThan(0);

		expect(runtime.getState().status).toBe('cancelled');
	});

	it('stream is reusable after an error', async () => {
		const streamError = new Error('transient');
		streamText
			.mockReturnValueOnce({
				stream: makeErrorStream(streamError),
				finishReason: silentReject(streamError),
				usage: Promise.resolve(undefined),
				response: silentReject(streamError),
				toolCalls: Promise.resolve([]),
			})
			.mockReturnValue(makeStreamSuccess('OK'));

		const { runtime } = createRuntime();

		const { stream: firstStream } = await runtime.stream('hello');
		const firstChunks = await collectChunks(firstStream);
		expect(firstChunks.some((c) => c.type === 'error')).toBe(true);

		const { stream: secondStream } = await runtime.stream('hello');
		const secondChunks = await collectChunks(secondStream);
		expect(secondChunks.some((c) => c.type === 'finish')).toBe(true);
		const finish = secondChunks.find((c) => c.type === 'finish') as StreamChunk & {
			type: 'finish';
			finishReason: string;
		};
		expect(finish.finishReason).toBe('stop');
	});

	it('never throws when consuming the stream after an error', async () => {
		const streamError = new Error('stream broke');
		streamText.mockReturnValue({
			stream: makeErrorStream(streamError),
			finishReason: silentReject(streamError),
			usage: Promise.resolve(undefined),
			response: silentReject(streamError),
			toolCalls: Promise.resolve([]),
		});

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');

		await expect(collectChunks(readableStream)).resolves.toBeDefined();
	});
});

// ---------------------------------------------------------------------------
// resume() — graceful error contract
// ---------------------------------------------------------------------------

describe('AgentRuntime.resume() — graceful error contract', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('rejects with an error when the runId is not found', async () => {
		const { runtime } = createRuntime();
		const streamPromise = runtime.resume(
			'stream',
			{ approved: true },
			{ runId: 'nonexistent-run-id', toolCallId: 'tc-1' },
		);
		await expect(streamPromise).rejects.toMatchObject({
			name: 'StaleResumeError',
			message: 'No suspended run found for runId: nonexistent-run-id',
		});
	});
});

// ---------------------------------------------------------------------------
// State transitions
// ---------------------------------------------------------------------------

describe('AgentRuntime — state transitions on error', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('starts idle, then reflects running→failed after a generate error', async () => {
		const { runtime } = createRuntime();

		expect(runtime.getState().status).toBe('idle');

		generateText.mockRejectedValue(new Error('oops'));
		const runDone = runtime.generate('hi');

		await runDone;
		expect(runtime.getState().status).toBe('failed');
	});

	it('starts idle, then reflects running→cancelled on abort', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const { runtime, bus } = createRuntime();
		bus.on(AgentEvent.AgentStart, () => bus.abort());

		await runtime.generate('hi');
		expect(runtime.getState().status).toBe('cancelled');
	});

	it('transitions to success on a clean run', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const { runtime } = createRuntime();
		await runtime.generate('hi');

		expect(runtime.getState().status).toBe('success');
	});

	it("reflects running→success on stream()'s own result.getState(), not just runtime.getState()", async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const { runtime } = createRuntime();
		const { stream, getState } = await runtime.stream('hi');
		await collectChunks(stream);

		expect(getState().status).toBe('success');
	});
});

// ---------------------------------------------------------------------------
// Concurrent tool execution
// ---------------------------------------------------------------------------

/** Create a simple non-interruptible mock tool. */
function makeMockTool(name: string, handler: (input: unknown) => Promise<unknown>): BuiltTool {
	return {
		name,
		description: `Mock tool ${name}`,
		inputSchema: z.object({ value: z.string().optional() }),
		handler: async (input) => await handler(input),
	};
}

/** Create an interruptible mock tool that suspends. */
function makeSuspendingTool(
	name: string,
	handler: (input: unknown, ctx: InterruptibleToolContext) => Promise<unknown>,
	onCancellation?: BuiltTool['onCancellation'],
): BuiltTool {
	return {
		name,
		description: `Suspending tool ${name}`,
		inputSchema: z.object({ value: z.string().optional() }),
		suspendSchema: z.object({ reason: z.string() }),
		resumeSchema: z.object({ approved: z.boolean() }),
		handler: async (input, ctx) => await handler(input, ctx as InterruptibleToolContext),
		...(onCancellation ? { onCancellation } : {}),
	};
}

/** Build a runtime with tools and concurrency config. */
function createRuntimeWithTools(
	tools: BuiltTool[],
	concurrency: number,
	eventBus?: AgentEventBus,
	model = 'openai/gpt-4o-mini',
) {
	const bus = eventBus ?? new AgentEventBus();
	const runtime = new AgentRuntime({
		name: 'test',
		model,
		instructions: 'You are a test assistant.',
		tools,
		eventBus: bus,
		toolCallConcurrency: concurrency,
		checkpointStorage: 'memory',
	});
	return { runtime, bus };
}

type ClaimForResume = (key: string, state: SerializableAgentState) => Promise<boolean>;
type ClaimingCheckpointStore = CheckpointStore & {
	claimForResume: MockedFunction<ClaimForResume>;
};

function createRuntimeWithCheckpointStore(
	tools: BuiltTool[],
	checkpointStorage: CheckpointStore,
	eventBus?: AgentEventBus,
): AgentRuntime {
	return new AgentRuntime({
		name: 'test',
		model: 'openai/gpt-4o-mini',
		instructions: 'You are a test assistant.',
		tools,
		checkpointStorage,
		...(eventBus ? { eventBus } : {}),
	});
}

function makeClaimingCheckpointStore(): ClaimingCheckpointStore {
	const checkpoints = new Map<string, SerializableAgentState>();

	return {
		save: vi.fn(async (key: string, state: SerializableAgentState): Promise<void> => {
			await Promise.resolve();
			checkpoints.set(key, structuredClone(state));
		}),
		load: vi.fn(async (key: string): Promise<SerializableAgentState | undefined> => {
			await Promise.resolve();
			const state = checkpoints.get(key);
			return state ? structuredClone(state) : undefined;
		}),
		delete: vi.fn(async (key: string): Promise<void> => {
			await Promise.resolve();
			checkpoints.delete(key);
		}),
		claimForResume: vi.fn<ClaimForResume>(async (key, state) => {
			await Promise.resolve();
			const current = checkpoints.get(key);
			if (!current || JSON.stringify(current) !== JSON.stringify(state)) return false;
			checkpoints.set(key, { ...state, status: 'running' });
			return true;
		}),
	};
}

/** Create a generateText response that triggers tool calls, then a follow-up stop. */
function makeGenerateWithToolCalls(
	toolCalls: Array<{ toolCallId: string; toolName: string; args: Record<string, unknown> }>,
) {
	return {
		finishReason: 'tool-calls',
		usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
		response: {
			messages: [
				{
					role: 'assistant',
					content: toolCalls.map((tc) => ({
						type: 'tool-call',
						toolCallId: tc.toolCallId,
						toolName: tc.toolName,
						args: tc.args,
						input: tc.args,
					})),
				},
			],
		},
		toolCalls: toolCalls.map((tc) => ({
			toolCallId: tc.toolCallId,
			toolName: tc.toolName,
			input: tc.args,
		})),
	};
}

describe('AgentRuntime — deferred tool loading', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	it('searches and loads deferred tools into the next generate iteration', async () => {
		const coreTool = makeMockTool('core_tool', async () => await Promise.resolve({ ok: true }));
		const deferredTool = makeMockTool(
			'deferred_capability',
			async () => await Promise.resolve({ ok: true }),
		);
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [coreTool],
			deferredTools: [deferredTool],
		});
		const counter = makeExecutionCounter();

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-search',
						toolName: 'search_tools',
						args: { query: 'deferred capability' },
					},
				]),
			)
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-load',
						toolName: 'load_tool',
						args: { toolName: 'deferred_capability' },
					},
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('ready'));

		const result = await runtime.generate('need the deferred capability', {
			executionCounter: counter,
		});

		expect(generateText).toHaveBeenCalledTimes(3);
		expect(counter.incrementMessageCount).toHaveBeenCalledTimes(1);
		expect(counter.incrementToolCallCount).toHaveBeenCalledTimes(2);
		expect(counter.incrementTokenCount).toHaveBeenCalledTimes(3);

		const searchCall = result.toolCalls?.find((toolCall) => toolCall.tool === 'search_tools');
		expect(searchCall?.output).toEqual({
			results: [
				{
					name: 'deferred_capability',
					description: 'Mock tool deferred_capability',
					loaded: false,
				},
			],
		});

		const loadCall = result.toolCalls?.find((toolCall) => toolCall.tool === 'load_tool');
		expect(loadCall?.output).toEqual({
			status: 'loaded',
			toolName: 'deferred_capability',
			tool: {
				name: 'deferred_capability',
				description: 'Mock tool deferred_capability',
				loaded: true,
			},
			message: 'Tool "deferred_capability" is loaded and will be available on the next model turn.',
		});

		const generateTextCalls = generateText.mock.calls as Array<
			[{ tools: Record<string, unknown> }]
		>;
		const firstCall = generateTextCalls[0][0];
		const firstTools = Object.keys(firstCall.tools);
		expect(firstTools).toEqual(expect.arrayContaining(['core_tool', 'search_tools', 'load_tool']));
		expect(firstTools).not.toContain('deferred_capability');

		const secondTools = Object.keys(generateTextCalls[1][0].tools);
		expect(secondTools).toEqual(expect.arrayContaining(['core_tool', 'search_tools', 'load_tool']));
		expect(secondTools).not.toContain('deferred_capability');

		const thirdTools = Object.keys(generateTextCalls[2][0].tools);
		expect(thirdTools).toEqual(
			expect.arrayContaining(['core_tool', 'search_tools', 'load_tool', 'deferred_capability']),
		);
	});

	it('does not leak loaded deferred tools into the next generate run', async () => {
		const coreTool = makeMockTool('core_tool', async () => await Promise.resolve({ ok: true }));
		const deferredTool = makeMockTool(
			'deferred_capability',
			async () => await Promise.resolve({ ok: true }),
		);
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [coreTool],
			deferredTools: [deferredTool],
		});

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-load',
						toolName: 'load_tool',
						args: { toolName: 'deferred_capability' },
					},
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('first done'));

		await runtime.generate('load the deferred capability');

		generateText.mockClear();
		generateText.mockResolvedValueOnce(makeGenerateSuccess('second done'));

		await runtime.generate('start a fresh run');

		const generateTextCalls = generateText.mock.calls as Array<
			[{ tools: Record<string, unknown> }]
		>;
		const freshRunTools = Object.keys(generateTextCalls[0][0].tools);
		expect(freshRunTools).toEqual(
			expect.arrayContaining(['core_tool', 'search_tools', 'load_tool']),
		);
		expect(freshRunTools).not.toContain('deferred_capability');
	});

	it('resumes a suspended deferred tool after it has been loaded', async () => {
		const deferredTool = makeSuspendingTool('deferred_approval', async (input, ctx) => {
			if (!ctx.resumeData) {
				return await ctx.suspend({ reason: 'approve deferred action?' });
			}

			const resumeData = ctx.resumeData as { approved: boolean };
			return {
				approved: resumeData.approved,
				value: (input as { value?: string }).value,
			};
		});

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			deferredTools: [deferredTool],
			checkpointStorage: 'memory',
		});

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-load',
						toolName: 'load_tool',
						args: { toolName: 'deferred_approval' },
					},
				]),
			)
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-deferred',
						toolName: 'deferred_approval',
						args: { value: 'needs approval' },
					},
				]),
			);

		const firstResult = await runtime.generate('load and run the deferred approval tool');

		expect(firstResult.finishReason).toBe('tool-calls');
		expect(firstResult.pendingSuspend).toHaveLength(1);
		expect(firstResult.pendingSuspend![0].toolName).toBe('deferred_approval');

		generateText.mockResolvedValueOnce(makeGenerateSuccess('approved'));

		const resumeResult = await runtime.resume(
			'generate',
			{ approved: true },
			{
				runId: firstResult.pendingSuspend![0].runId,
				toolCallId: firstResult.pendingSuspend![0].toolCallId,
			},
		);

		expect(resumeResult.finishReason).toBe('stop');
		expect(resumeResult.toolCalls).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					tool: 'deferred_approval',
					output: { approved: true, value: 'needs approval' },
				}),
			]),
		);
	});
});

describe('AgentRuntime — concurrent tool execution', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	it('keeps abort signals scoped to overlapping generate runs', async () => {
		let firstModelSignal: AbortSignal | undefined;
		let secondModelSignal: AbortSignal | undefined;
		let firstToolSignal: AbortSignal | undefined;
		let resolveFirstModel!: (value: unknown) => void;
		const firstModel = new Promise((resolve) => {
			resolveFirstModel = resolve;
		});
		let firstModelCalled!: () => void;
		const waitForFirstModelCall = new Promise<void>((resolve) => {
			firstModelCalled = resolve;
		});

		const signalTool: BuiltTool = {
			name: 'signal_tool',
			description: 'Captures the run signal',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async (_input, ctx) => {
				firstToolSignal = (ctx as ToolContext).abortSignal;
				return await Promise.resolve({ done: true });
			},
		};
		const { runtime } = createRuntimeWithTools([signalTool], 1);

		generateText
			.mockImplementationOnce(async (options: { abortSignal?: AbortSignal }) => {
				firstModelSignal = options.abortSignal;
				firstModelCalled();
				return await firstModel;
			})
			.mockImplementationOnce(async (options: { abortSignal?: AbortSignal }) => {
				secondModelSignal = options.abortSignal;
				return await Promise.resolve(makeGenerateSuccess('second done'));
			})
			.mockResolvedValueOnce(makeGenerateSuccess('first done'));

		const firstRun = runtime.generate('first');
		await waitForFirstModelCall;

		const secondRun = runtime.generate('second');
		await secondRun;

		resolveFirstModel(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'signal_tool', args: { value: 'first' } },
			]),
		);
		await firstRun;

		expect(firstToolSignal).toBe(firstModelSignal);
		expect(firstToolSignal).not.toBe(secondModelSignal);
	});

	it('runs tools concurrently when concurrency > 1', async () => {
		let peakConcurrency = 0;
		let activeConcurrency = 0;

		const slowTool = makeMockTool('slow_tool', async () => {
			activeConcurrency++;
			peakConcurrency = Math.max(peakConcurrency, activeConcurrency);
			await new Promise((r) => setTimeout(r, 50));
			activeConcurrency--;
			return { done: true };
		});

		const { runtime } = createRuntimeWithTools([slowTool], Infinity);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'slow_tool', args: { value: 'a' } },
					{ toolCallId: 'tc-2', toolName: 'slow_tool', args: { value: 'b' } },
					{ toolCallId: 'tc-3', toolName: 'slow_tool', args: { value: 'c' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('run tools');

		expect(result.finishReason).toBe('stop');
		expect(peakConcurrency).toBe(3);
	});

	it('suspensions stop later batches and save unexecuted tools as pending', async () => {
		const executedTools: string[] = [];

		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			executedTools.push('suspend_tool');
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const normalTool = makeMockTool('normal_tool', async () => {
			executedTools.push('normal_tool');
			return await Promise.resolve({ done: true });
		});

		const { runtime } = createRuntimeWithTools([suspendTool, normalTool], 2);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'x' } },
				{ toolCallId: 'tc-2', toolName: 'normal_tool', args: { value: 'y' } },
				{ toolCallId: 'tc-3', toolName: 'normal_tool', args: { value: 'z' } },
			]),
		);

		const result = await runtime.generate('run tools');

		// Batch 1 (tc-1 + tc-2) runs, but batch 2 (tc-3) is skipped because batch 1 had a suspension
		expect(executedTools).toContain('suspend_tool');
		expect(executedTools.filter((t) => t === 'normal_tool').length).toBe(1);
		expect(result.finishReason).toBe('tool-calls');

		// tc-1 is suspended, tc-3 is saved as unexecuted pending
		expect(result.pendingSuspend).toBeDefined();
		expect(result.pendingSuspend!.length).toBe(1);
		expect(result.pendingSuspend![0].toolCallId).toBe('tc-1');

		// Verify tc-3 is in the persisted state as a pending tool call (without suspendPayload)
		const state = runtime.getState();
		expect(state.pendingToolCalls['tc-3']).toBeDefined();
		expect(state.pendingToolCalls['tc-3'].suspended).toBe(false);
	});

	it('validates resume data when targeting an unexecuted pending tool', async () => {
		const handler = vi.fn(async (_input, ctx: InterruptibleToolContext) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const suspendTool = makeSuspendingTool('suspend_tool', handler);
		const { runtime } = createRuntimeWithTools([suspendTool], 1);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		expect(runtime.getState().pendingToolCalls['tc-2']).toMatchObject({ suspended: false });

		await expect(
			runtime.resume('generate', { approved: 'yes' }, { runId: first.runId, toolCallId: 'tc-2' }),
		).rejects.toThrow('Invalid resume payload');
		expect(handler).toHaveBeenCalledOnce();
	});

	it('pendingSuspend is an array of all suspended tools with preserved payloads', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const result = await runtime.generate('run tools');

		expect(result.pendingSuspend).toBeDefined();
		expect(result.pendingSuspend!.length).toBe(2);

		const ids = result.pendingSuspend!.map((s) => s.toolCallId);
		expect(ids).toContain('tc-1');
		expect(ids).toContain('tc-2');

		for (const s of result.pendingSuspend!) {
			expect(s.suspendPayload).toEqual({ reason: 'needs approval' });
			expect(s.runId).toBeTruthy();
			expect(s.toolName).toBe('suspend_tool');
		}
	});

	it('resume returns leftover suspends after resolving one', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		expect(first.pendingSuspend!.length).toBe(2);

		const { runId } = first.pendingSuspend![0];

		// Resume tc-1 — concurrent resume doesn't call generateText when pending tools remain
		const resumed = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId: 'tc-1' },
		);

		// tc-2 should still be pending
		expect(resumed.pendingSuspend).toBeDefined();
		expect(resumed.pendingSuspend!.length).toBe(1);
		expect(resumed.pendingSuspend![0].toolCallId).toBe('tc-2');
		// runId must not change across suspend/resume cycles
		expect(resumed.pendingSuspend![0].runId).toBe(runId);
	});

	it('already-suspended tools are NOT re-executed on resume', async () => {
		let executionCount = 0;

		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			executionCount++;
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		expect(executionCount).toBe(2); // both executed initially

		const { runId } = first.pendingSuspend![0];

		// Resume tc-1 — only tc-1 should re-execute, not tc-2
		generateText.mockResolvedValueOnce(makeGenerateSuccess('After resume'));
		await runtime.resume('generate', { approved: true }, { runId, toolCallId: 'tc-1' });

		expect(executionCount).toBe(3); // only tc-1 re-executed (2 initial + 1 resume)
	});

	it('after all suspends resolved, LLM loop continues', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		const { runId } = first.pendingSuspend![0];

		// Resume tc-1 — tc-2 still pending, so no LLM call
		const second = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId: 'tc-1' },
		);
		expect(second.pendingSuspend!.length).toBe(1);

		// runId must stay the same across suspend/resume cycles
		const runId2 = second.pendingSuspend![0].runId;
		expect(runId2).toBe(runId);

		// Resume tc-2 — no more pending, LLM loop should continue
		generateText.mockResolvedValueOnce(makeGenerateSuccess('all done'));
		const third = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId: runId2, toolCallId: 'tc-2' },
		);

		expect(third.pendingSuspend).toBeUndefined();
		expect(third.finishReason).toBe('stop');
	});

	it('cancels a suspended tool before resume validation and adds the user message', async () => {
		const handler = vi.fn(async (_input, ctx: InterruptibleToolContext) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const suspendTool = makeSuspendingTool('suspend_tool', handler);
		const receivedMessages: unknown[] = [];

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		const first = await runtime.generate('run tools');
		const { runId, toolCallId } = first.pendingSuspend![0];
		generateText.mockImplementationOnce(async ({ messages }: { messages: unknown[] }) => {
			receivedMessages.push(...messages);
			return await Promise.resolve(makeGenerateSuccess('Cancelled'));
		});

		const result = await runtime.resume('generate', createCancellation('Do not run this tool'), {
			runId,
			toolCallId,
		});

		expect(result.finishReason).toBe('stop');
		expect(handler).toHaveBeenCalledTimes(1);
		expect(result.toolCalls).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					tool: 'suspend_tool',
					output: '[Tool call cancelled. User said: "Do not run this tool"]',
					canceled: true,
				}),
			]),
		);
		expect(receivedMessages).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					role: 'user',
					content: expect.arrayContaining([
						expect.objectContaining({ type: 'text', text: 'Do not run this tool' }),
					]),
				}),
			]),
		);
	});

	it('streams cancellation as a normal tool result on resume', async () => {
		const handler = vi.fn(async (_input, ctx: InterruptibleToolContext) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const suspendTool = makeSuspendingTool('suspend_tool', handler);

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		const first = await runtime.generate('run tools');
		const { runId, toolCallId } = first.pendingSuspend![0];
		streamText.mockReturnValueOnce(makeStreamSuccess('Cancelled'));

		const resumed = await runtime.resume('stream', createCancellation('Stop this action'), {
			runId,
			toolCallId,
		});
		const chunks = await collectChunks(resumed.stream as ReadableStream<unknown>);

		expect(handler).toHaveBeenCalledTimes(1);
		expect(chunks).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					type: 'tool-result',
					toolCallId,
					toolName: 'suspend_tool',
					output: '[Tool call cancelled. User said: "Stop this action"]',
					canceled: true,
				}),
			]),
		);
	});

	it('streams skipped sibling tool results when cancelling one of multiple suspensions', async () => {
		const handler = vi.fn(async (input, ctx: InterruptibleToolContext) => {
			if (ctx.resumeData) return { approved: true };
			const value = (input as { value: string }).value;
			return await ctx.suspend(
				{ reason: 'needs approval' },
				{ continuation: { cleanupKey: `cleanup-${value}` } },
			);
		});
		const onCancellation = makeCancellationSpy();
		const suspendTool = makeSuspendingTool('suspend_tool', handler, onCancellation);

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		const { runId } = first.pendingSuspend![0];
		streamText.mockReturnValueOnce(makeStreamSuccess('Cancelled'));

		const resumed = await runtime.resume('stream', createCancellation('Stop this action'), {
			runId,
			toolCallId: 'tc-1',
		});
		const chunks = await collectChunks(resumed.stream as ReadableStream<unknown>);

		expect(handler).toHaveBeenCalledTimes(2);
		expect(onCancellation).toHaveBeenCalledTimes(2);
		expect(onCancellation.mock.calls.map(([toolInput]) => toolInput)).toEqual([
			{ value: 'a' },
			{ value: 'b' },
		]);
		expect(onCancellation.mock.calls.map(([, context]) => context.continuation)).toEqual([
			{ cleanupKey: 'cleanup-a' },
			{ cleanupKey: 'cleanup-b' },
		]);
		expect(chunks).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					type: 'tool-result',
					toolCallId: 'tc-1',
					toolName: 'suspend_tool',
					output: '[Tool call cancelled. User said: "Stop this action"]',
					canceled: true,
				}),
				expect.objectContaining({
					type: 'tool-result',
					toolCallId: 'tc-2',
					toolName: 'suspend_tool',
					output: '[Skipped: a sibling tool call was cancelled]',
					canceled: true,
				}),
			]),
		);
	});

	it('does not run suspension cleanup for an unexecuted sibling', async () => {
		const handler = vi.fn(async (_input, ctx: InterruptibleToolContext) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const onCancellation = makeCancellationSpy();
		const suspendTool = makeSuspendingTool('suspend_tool', handler, onCancellation);

		const { runtime } = createRuntimeWithTools([suspendTool], 1);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
			]),
		);

		const first = await runtime.generate('run tools');
		const { runId } = first.pendingSuspend![0];
		streamText.mockReturnValueOnce(makeStreamSuccess('Cancelled'));

		const resumed = await runtime.resume('stream', createCancellation('Stop this action'), {
			runId,
			toolCallId: 'tc-1',
		});
		await collectChunks(resumed.stream as ReadableStream<unknown>);

		expect(handler).toHaveBeenCalledOnce();
		expect(onCancellation).toHaveBeenCalledOnce();
		expect(onCancellation).toHaveBeenCalledWith(
			{ value: 'a' },
			expect.objectContaining({ toolCallId: 'tc-1' }),
		);
	});

	it('re-suspends when cancellation cleanup fails', async () => {
		const handler = vi.fn(async (_input, ctx: InterruptibleToolContext) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const cleanupError = new Error('checkpoint delete failed');
		const suspendTool = {
			...makeSuspendingTool('suspend_tool', handler),
			onCancellation: vi
				.fn<NonNullable<BuiltTool['onCancellation']>>()
				.mockRejectedValue(cleanupError),
		};

		const { runtime } = createRuntimeWithTools([suspendTool], 1);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		const first = await runtime.generate('run tool');
		const resumed = await runtime.resume('generate', createCancellation('Stop this action'), {
			runId: first.runId,
			toolCallId: 'tc-1',
		});

		expect(resumed.pendingSuspend).toEqual([
			expect.objectContaining({
				runId: first.runId,
				toolCallId: 'tc-1',
				suspendPayload: { reason: 'needs approval' },
			}),
		]);
		expect(resumed.getState().status).toBe('suspended');
	});

	it('bounded concurrency (2) batches respects the limit', async () => {
		const batchSizes: number[] = [];
		let activeConcurrency = 0;

		const slowTool = makeMockTool('slow_tool', async () => {
			activeConcurrency++;
			batchSizes.push(activeConcurrency);
			await new Promise((r) => setTimeout(r, 30));
			activeConcurrency--;
			return { done: true };
		});

		const { runtime } = createRuntimeWithTools([slowTool], 2);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'slow_tool', args: { value: 'a' } },
					{ toolCallId: 'tc-2', toolName: 'slow_tool', args: { value: 'b' } },
					{ toolCallId: 'tc-3', toolName: 'slow_tool', args: { value: 'c' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('run tools');

		expect(result.finishReason).toBe('stop');
		// Peak concurrency within any batch should never exceed 2
		expect(Math.max(...batchSizes)).toBeLessThanOrEqual(2);
	});

	it('runs consecutive delegate_subagent calls in parallel up to maxChildren when toolCallConcurrency is 1', async () => {
		let activeDelegations = 0;
		let peakDelegations = 0;

		const runSubAgent: DelegateSubAgentRunner = async (request) => {
			activeDelegations++;
			peakDelegations = Math.max(peakDelegations, activeDelegations);
			await new Promise((resolve) => setTimeout(resolve, 30));
			activeDelegations--;
			return {
				status: 'completed',
				taskPath: request.taskPath,
				answer: `done:${request.taskName}`,
			};
		};

		const delegateTool = createDelegateSubAgentTool({
			policy: { maxChildren: 5 },
			runSubAgent,
		});
		const { runtime } = createRuntimeWithTools([delegateTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls(
					Array.from({ length: 5 }, (_, index) => ({
						toolCallId: `tc-${index + 1}`,
						toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
						args: {
							subAgentId: 'inline',
							taskName: `ping_${index + 1}`,
							goal: 'Reply with ping.',
						},
					})),
				),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('spawn 5 sub agents');

		expect(result.finishReason).toBe('stop');
		expect(peakDelegations).toBe(5);
	});

	it('batches delegate_subagent parallelism by maxChildren when more calls are issued than the parallel limit', async () => {
		let activeDelegations = 0;
		let peakDelegations = 0;

		const runSubAgent: DelegateSubAgentRunner = async (request) => {
			activeDelegations++;
			peakDelegations = Math.max(peakDelegations, activeDelegations);
			await new Promise((resolve) => setTimeout(resolve, 30));
			activeDelegations--;
			return {
				status: 'completed',
				taskPath: request.taskPath,
				answer: `done:${request.taskName}`,
			};
		};

		const delegateTool = createDelegateSubAgentTool({
			policy: { maxChildren: 2 },
			runSubAgent,
		});
		const { runtime } = createRuntimeWithTools([delegateTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls(
					Array.from({ length: 4 }, (_, index) => ({
						toolCallId: `tc-${index + 1}`,
						toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
						args: {
							subAgentId: 'inline',
							taskName: `ping_${index + 1}`,
							goal: 'Reply with ping.',
						},
					})),
				),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('spawn 4 sub agents');

		expect(result.finishReason).toBe('stop');
		expect(peakDelegations).toBe(2);
		expect(result.toolCalls).toHaveLength(4);
		for (const entry of result.toolCalls ?? []) {
			expect(JSON.stringify(entry.output)).toContain('"status":"completed"');
		}
	});

	it('batches a renamed delegate tool by maxChildren via metadata, not by tool name', async () => {
		let activeDelegations = 0;
		let peakDelegations = 0;

		const runSubAgent: DelegateSubAgentRunner = async (request) => {
			activeDelegations++;
			peakDelegations = Math.max(peakDelegations, activeDelegations);
			await new Promise((resolve) => setTimeout(resolve, 30));
			activeDelegations--;
			return {
				status: 'completed',
				taskPath: request.taskPath,
				answer: `done:${request.taskName}`,
			};
		};

		const delegateTool = createDelegateSubAgentTool({
			name: 'agent',
			policy: { maxChildren: 2 },
			runSubAgent,
		});
		const { runtime } = createRuntimeWithTools([delegateTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls(
					Array.from({ length: 4 }, (_, index) => ({
						toolCallId: `tc-${index + 1}`,
						toolName: 'agent',
						args: {
							subAgentId: 'inline',
							taskName: `ping_${index + 1}`,
							goal: 'Reply with ping.',
						},
					})),
				),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('spawn 4 sub agents');

		expect(result.finishReason).toBe('stop');
		expect(peakDelegations).toBe(2);
	});

	it('fails fast when delegate metadata has an invalid maxChildren batch size', async () => {
		const malformedDelegateTool: BuiltTool = {
			name: DELEGATE_SUB_AGENT_TOOL_NAME,
			description: 'Malformed delegate tool',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve({ done: true }),
			metadata: {
				[INLINE_DELEGATE_SUB_AGENT_TOOL_METADATA_KEY]: {
					policy: { maxChildren: 0 },
				},
			},
		};
		const { runtime } = createRuntimeWithTools([malformedDelegateTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{
					toolCallId: 'tc-1',
					toolName: DELEGATE_SUB_AGENT_TOOL_NAME,
					args: {
						subAgentId: 'inline',
						taskName: 'ping',
						goal: 'Reply with ping.',
					},
				},
			]),
		);

		const result = await runtime.generate('spawn sub agent');

		expect(result.finishReason).toBe('error');
		expect(result.error).toBeInstanceOf(Error);
		if (result.error instanceof Error) {
			expect(result.error.message).toBe('Invalid tool-call batch size for delegate_subagent: 0');
		}
	});

	it('keeps non-delegate tools sequential when toolCallConcurrency is 1', async () => {
		let active = 0;
		let peak = 0;

		const slowTool = makeMockTool('slow_tool', async () => {
			active++;
			peak = Math.max(peak, active);
			await new Promise((resolve) => setTimeout(resolve, 30));
			active--;
			return { done: true };
		});

		const { runtime } = createRuntimeWithTools([slowTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'slow_tool', args: { value: 'a' } },
					{ toolCallId: 'tc-2', toolName: 'slow_tool', args: { value: 'b' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('run tools');

		expect(result.finishReason).toBe('stop');
		expect(peak).toBe(1);
	});

	it('tool error becomes LLM-visible message — loop continues and slow tool completes', async () => {
		const completedTools: string[] = [];

		const errorTool = makeMockTool('error_tool', async () => {
			return await Promise.reject(new Error('tool failed'));
		});

		const slowTool = makeMockTool('slow_tool', async () => {
			await new Promise((r) => setTimeout(r, 50));
			completedTools.push('slow_tool');
			return { done: true };
		});

		const { runtime } = createRuntimeWithTools([errorTool, slowTool], Infinity);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'error_tool', args: {} },
					{ toolCallId: 'tc-2', toolName: 'slow_tool', args: {} },
				]),
			)
			// The LLM is called again after seeing the error tool result
			.mockResolvedValueOnce(makeGenerateSuccess('Handled the error'));

		const result = await runtime.generate('run tools');

		// Loop continues after tool error — LLM sees the error and responds
		expect(result.finishReason).toBe('stop');
		// allSettled waits for all in-flight tools before returning
		expect(completedTools).toContain('slow_tool');
	});

	it('sequential mode (concurrency=1) wraps single suspension in an array', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		const result = await runtime.generate('run tools');

		expect(result.pendingSuspend).toBeDefined();
		expect(Array.isArray(result.pendingSuspend)).toBe(true);
		expect(result.pendingSuspend!.length).toBe(1);
		expect(result.pendingSuspend![0].toolCallId).toBe('tc-1');
	});

	it('tool error produces an error tool-result in the message list and loop continues', async () => {
		type ToolOutputContent = {
			type: string;
			output?: { type: string; value?: unknown };
		};
		type ToolMessage = { role: string; content: ToolOutputContent[] };
		const receivedMessages: unknown[] = [];

		const errorTool = makeMockTool('error_tool', async () => {
			return await Promise.reject(new Error('intentional failure'));
		});

		const { runtime } = createRuntimeWithTools([errorTool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'error_tool', args: { value: 'x' } },
				]),
			)
			.mockImplementationOnce(async ({ messages }: { messages: unknown[] }) => {
				receivedMessages.push(...messages);
				return await Promise.resolve(makeGenerateSuccess('I saw the error'));
			});

		const result = await runtime.generate('run tools');

		expect(result.finishReason).toBe('stop');
		// LLM was called a second time — it saw the error tool result and continued
		expect(generateText).toHaveBeenCalledTimes(2);
		// The second LLM call received a tool message whose output carries the error description.
		const toolMsg = receivedMessages.find(
			(m): m is ToolMessage =>
				typeof m === 'object' && m !== null && (m as ToolMessage).role === 'tool',
		);
		expect(toolMsg).toBeDefined();
		const hasErrorOutput = toolMsg!.content.some(
			(c) => c.output?.type === 'error-text' || c.output?.type === 'error-json',
		);
		expect(hasErrorOutput).toBe(true);
	});

	it('mixed success+error batch — successful results are preserved', async () => {
		const successTool = makeMockTool('success_tool', async () => {
			return await Promise.resolve({ value: 42 });
		});
		const errorTool = makeMockTool('error_tool', async () => {
			return await Promise.reject(new Error('tool failed'));
		});

		const { runtime } = createRuntimeWithTools([successTool, errorTool], Infinity);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'success_tool', args: {} },
					{ toolCallId: 'tc-2', toolName: 'error_tool', args: {} },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const result = await runtime.generate('run tools');

		expect(result.finishReason).toBe('stop');
		// Successful tool call is recorded in toolCalls summary
		expect(result.toolCalls).toBeDefined();
		expect(result.toolCalls!.some((tc) => tc.tool === 'success_tool')).toBe(true);
	});

	it('toolName is stored in pendingToolCalls and accessible without findToolName', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		await runtime.generate('run tools');

		const state = runtime.getState();
		expect(state.pendingToolCalls['tc-1']).toBeDefined();
		expect(state.pendingToolCalls['tc-1'].toolName).toBe('suspend_tool');
	});

	it('unexecuted pending tools also carry toolName', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});
		const normalTool = makeMockTool('normal_tool', async () => {
			return await Promise.resolve({ done: true });
		});

		// concurrency=1 so tc-1 runs first (suspends), tc-2 is left unexecuted
		const { runtime } = createRuntimeWithTools([suspendTool, normalTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'x' } },
				{ toolCallId: 'tc-2', toolName: 'normal_tool', args: { value: 'y' } },
			]),
		);

		await runtime.generate('run tools');

		const state = runtime.getState();
		expect(state.pendingToolCalls['tc-2']).toBeDefined();
		expect(state.pendingToolCalls['tc-2'].toolName).toBe('normal_tool');
		expect(state.pendingToolCalls['tc-2'].suspended).toBe(false);
	});

	it('tool-call-suspended chunks expose a dynamic resume schema without private continuation', async () => {
		const dynamicResumeSchema = z.object({ decision: z.enum(['continue', 'stop']) });
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend(
				{ reason: 'needs approval' },
				{
					resumeSchema: dynamicResumeSchema,
					continuation: { childRunId: 'private-child-run' },
				},
			);
		});

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		streamText.mockReturnValue({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{
								type: 'tool-call',
								toolCallId: 'tc-1',
								toolName: 'suspend_tool',
								args: { value: 'a' },
							},
						],
					},
				],
			}),
			toolCalls: Promise.resolve([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', input: { value: 'a' } },
			]),
		});

		const { stream: readableStream } = await runtime.stream('run tools');
		const chunks = await collectChunks(readableStream);

		const suspendedChunks = chunks.filter((c) => c.type === 'tool-call-suspended') as Array<
			StreamChunk & { type: 'tool-call-suspended'; resumeSchema?: unknown }
		>;
		expect(suspendedChunks.length).toBe(1);
		// resumeSchema should be a JSON Schema object (from zod-to-json-schema conversion)
		expect(suspendedChunks[0].resumeSchema).toEqual(
			expect.objectContaining({
				type: 'object',
				properties: {
					decision: expect.objectContaining({ enum: ['continue', 'stop'] }),
				},
				required: ['decision'],
			}),
		);
		expect(suspendedChunks[0]).not.toHaveProperty('continuation');
	});

	it('pendingSuspend entries include resumeSchema when tool defines one', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
			]),
		);

		const result = await runtime.generate('run tools');

		expect(result.pendingSuspend).toBeDefined();
		expect(result.pendingSuspend![0].resumeSchema).toBeDefined();
		expect(typeof result.pendingSuspend![0].resumeSchema).toBe('object');
	});

	it('stream mode emits multiple tool-call-suspended chunks for concurrent suspensions', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], Infinity);

		streamText.mockReturnValue({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{
								type: 'tool-call',
								toolCallId: 'tc-1',
								toolName: 'suspend_tool',
								args: { value: 'a' },
							},
							{
								type: 'tool-call',
								toolCallId: 'tc-2',
								toolName: 'suspend_tool',
								args: { value: 'b' },
							},
						],
					},
				],
			}),
			toolCalls: Promise.resolve([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', input: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', input: { value: 'b' } },
			]),
		});

		const { stream: readableStream } = await runtime.stream('run tools');
		const chunks = await collectChunks(readableStream);

		const suspendedChunks = chunks.filter((c) => c.type === 'tool-call-suspended');
		expect(suspendedChunks.length).toBe(2);

		const finishChunks = chunks.filter((c) => c.type === 'finish');
		expect(finishChunks.length).toBe(1);
		expect((finishChunks[0] as StreamChunk & { type: 'finish' }).finishReason).toBe('tool-calls');
	});

	it('stamps the consumed usage and model on the terminal finish chunk when a stream suspends', async () => {
		const suspendTool = makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			return await ctx.suspend({ reason: 'needs approval' });
		});

		const { runtime } = createRuntimeWithTools([suspendTool], 1);

		streamText.mockReturnValue({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{
								type: 'tool-call',
								toolCallId: 'tc-1',
								toolName: 'suspend_tool',
								args: { value: 'a' },
							},
						],
					},
				],
			}),
			toolCalls: Promise.resolve([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', input: { value: 'a' } },
			]),
		});

		const { stream: readableStream } = await runtime.stream('run tools');
		const chunks = await collectChunks(readableStream);

		const finishChunks = chunks.filter((c) => c.type === 'finish') as Array<
			StreamChunk & { type: 'finish' }
		>;
		expect(finishChunks.length).toBe(1);
		// The tokens consumed to reach the suspension must ride out on the finish
		// chunk so the host can bill them, mirroring the completion path.
		expect(finishChunks[0].usage).toMatchObject({
			promptTokens: 10,
			completionTokens: 5,
			totalTokens: 15,
		});
		expect(finishChunks[0].model).toBeTruthy();
	});

	it('bridges subagent lifecycle events from tool context into stream chunks', async () => {
		const lifecycleTool: BuiltTool = {
			name: 'delegate_subagent',
			description: 'Delegate work',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async (_input, ctx) => {
				const toolCtx = ctx as ToolContext;
				const base = {
					taskName: 'Research API',
					taskPath: '/root/research_api',
					...(toolCtx.runId !== undefined ? { parentRunId: toolCtx.runId } : {}),
					...(toolCtx.toolCallId !== undefined ? { parentToolCallId: toolCtx.toolCallId } : {}),
				};
				toolCtx.emitEvent?.({
					type: AgentEvent.SubAgentStarted,
					...base,
					startedAt: 100,
				});
				toolCtx.emitEvent?.({
					type: AgentEvent.SubAgentCompleted,
					...base,
					status: 'completed',
					startedAt: 100,
					finishedAt: 200,
					durationMs: 100,
					runId: 'child-run-1',
					finishReason: 'stop',
				});
				return await Promise.resolve({ ok: true });
			},
		};
		const { runtime } = createRuntimeWithTools([lifecycleTool], 1);

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'delegate_subagent',
									args: { value: 'a' },
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-1', toolName: 'delegate_subagent', input: { value: 'a' } },
				]),
			})
			.mockReturnValueOnce(makeStreamSuccess('done'));

		const { stream: readableStream } = await runtime.stream('run tools');
		const chunks = await collectChunks(readableStream);

		expect(chunks).toEqual(
			expect.arrayContaining([
				expect.objectContaining({ type: 'subagent-started', taskPath: '/root/research_api' }),
				expect.objectContaining({
					type: 'subagent-completed',
					status: 'completed',
					runId: 'child-run-1',
				}),
			]),
		);
	});
});

// Structured output — generate()
// ---------------------------------------------------------------------------

describe('AgentRuntime.generate() — structured output', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('returns structuredOutput when schema is configured', async () => {
		const expected = { answer: 'Paris', score: 0.95 };
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput(expected));

		const { runtime } = createStructuredRuntime();
		const result = await runtime.generate('What is the capital of France?');

		expect(result.structuredOutput).toEqual(expected);
		expect(result.finishReason).toBe('stop');
	});

	it('does not include structuredOutput when no schema is configured', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Just text'));

		const { runtime } = createRuntime();
		const result = await runtime.generate('hello');

		expect(result.structuredOutput).toBeUndefined();
	});

	it('passes the output spec to generateText when schema is configured', async () => {
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput({ answer: 'test', score: 1 }));

		const { runtime } = createStructuredRuntime();
		await runtime.generate('hello');

		const callArgs = (generateText.mock.calls[0] as [unknown, unknown])[0] as Record<
			string,
			unknown
		>;
		expect(callArgs.output).toBeDefined();
		expect(callArgs.output).toEqual(
			expect.objectContaining({ _type: 'object', schema: testSchema }),
		);
	});

	it('does not pass the output spec to generateText when no schema is configured', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const { runtime } = createRuntime();
		await runtime.generate('hello');

		const callArgs = (generateText.mock.calls[0] as [unknown, unknown])[0] as Record<
			string,
			unknown
		>;
		expect(callArgs.output).toBeUndefined();
	});

	it('returns structuredOutput when a raw JSON Schema is configured', async () => {
		const expected = { answer: 'Paris', score: 0.95 };
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput(expected));

		const { runtime } = createJsonSchemaStructuredRuntime();
		const result = await runtime.generate('What is the capital of France?');

		expect(result.structuredOutput).toEqual(expected);
		expect(result.finishReason).toBe('stop');
	});

	it('wraps a provider-strict JSON Schema output spec via jsonSchema() before passing to generateText', async () => {
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput({ answer: 'x', score: 1 }));

		const { runtime } = createJsonSchemaStructuredRuntime();
		await runtime.generate('hello');

		// Raw JSON Schema is made provider-strict (additionalProperties: false) and
		// wrapped with the AI SDK's jsonSchema() helper, rather than passed through
		// directly (which is what Zod schemas do).
		const strictTestJsonSchema = { ...testJsonSchema, additionalProperties: false };
		expect(jsonSchema).toHaveBeenCalledWith(strictTestJsonSchema);

		const callArgs = (generateText.mock.calls[0] as [unknown, unknown])[0] as Record<
			string,
			unknown
		>;
		expect(callArgs.output).toEqual(
			expect.objectContaining({
				_type: 'object',
				schema: { _type: 'jsonSchema', schema: strictTestJsonSchema },
			}),
		);
	});

	it('disables strict JSON Schema validation for OpenAI/Groq when a raw JSON Schema is configured', async () => {
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput({ answer: 'x', score: 1 }));

		const { runtime } = createJsonSchemaStructuredRuntime();
		await runtime.generate('hello');

		const callArgs = (generateText.mock.calls[0] as [unknown, unknown])[0] as Record<
			string,
			unknown
		>;
		expect(callArgs.providerOptions).toEqual(
			expect.objectContaining({
				openai: { strictJsonSchema: false },
				groq: { strictJsonSchema: false },
			}),
		);
	});

	it('keeps strict JSON Schema validation (no relaxation) for a Zod schema', async () => {
		generateText.mockResolvedValue(makeGenerateSuccessWithOutput({ answer: 'x', score: 1 }));

		const { runtime } = createStructuredRuntime();
		await runtime.generate('hello');

		const callArgs = (generateText.mock.calls[0] as [unknown, unknown])[0] as Record<
			string,
			unknown
		>;
		// No thinking/provider options configured and Zod output → no strict override.
		expect(callArgs.providerOptions).toBeUndefined();
	});

	it('preserves structuredOutput through tool-call iterations', async () => {
		const tool: BuiltTool = {
			name: 'add',
			description: 'Add numbers',
			inputSchema: z.object({ a: z.number(), b: z.number() }),
			handler: async (input: unknown) => {
				const { a, b } = input as { a: number; b: number };
				return await Promise.resolve({ sum: a + b });
			},
		};

		const expected = { answer: '5', score: 1 };
		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'add', { a: 2, b: 3 }))
			.mockResolvedValueOnce(makeGenerateSuccessWithOutput(expected, 'The sum is 5'));

		const { runtime } = createStructuredRuntime({ tools: [tool] });
		const result = await runtime.generate('What is 2 + 3?');

		expect(result.structuredOutput).toEqual(expected);
		expect(result.finishReason).toBe('stop');
	});
});

// ---------------------------------------------------------------------------
// Provider-executed tool timing — stream()
// ---------------------------------------------------------------------------

describe('AgentRuntime.stream() — provider-executed tool timing', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('emits tool-execution-start/end for a provider-executed tool result', async () => {
		streamText.mockReturnValue(
			makeStreamWithProviderTool({
				toolCallId: 'tc-ws',
				toolName: 'web_search',
				input: { query: 'n8n' },
				output: [{ url: 'https://n8n.io' }],
			}),
		);
		const { runtime } = createRuntime();

		const { stream } = await runtime.stream('search please');
		const chunks = await collectChunks(stream);

		const start = chunks.find(
			(c): c is Extract<StreamChunk, { type: 'tool-execution-start' }> =>
				c.type === 'tool-execution-start' && c.toolCallId === 'tc-ws',
		);
		const end = chunks.find(
			(c): c is Extract<StreamChunk, { type: 'tool-execution-end' }> =>
				c.type === 'tool-execution-end' && c.toolCallId === 'tc-ws',
		);

		expect(start).toBeDefined();
		expect(start?.toolName).toBe('web_search');
		expect(typeof start?.startTime).toBe('number');

		expect(end).toBeDefined();
		expect(end?.isError).toBe(false);
		expect(typeof end?.endTime).toBe('number');
	});

	it('emits tool-execution-end with isError on a provider-executed tool error', async () => {
		streamText.mockReturnValue(
			makeStreamWithProviderTool({
				toolCallId: 'tc-ws-err',
				toolName: 'web_search',
				input: { query: 'n8n' },
				error: new Error('search failed'),
			}),
		);
		const { runtime } = createRuntime();

		const { stream } = await runtime.stream('search please');
		const chunks = await collectChunks(stream);

		const end = chunks.find(
			(c): c is Extract<StreamChunk, { type: 'tool-execution-end' }> =>
				c.type === 'tool-execution-end' && c.toolCallId === 'tc-ws-err',
		);

		expect(end).toBeDefined();
		expect(end?.isError).toBe(true);
		expect(typeof end?.endTime).toBe('number');

		const toolResult = chunks.find(
			(c): c is Extract<StreamChunk, { type: 'tool-result' }> =>
				c.type === 'tool-result' && c.toolCallId === 'tc-ws-err',
		);
		expect(toolResult).toBeDefined();
		expect(toolResult?.isError).toBe(true);
		expect(toolResult?.output).toEqual(new Error('search failed'));
	});
});

// ---------------------------------------------------------------------------
// Structured output — stream()
// ---------------------------------------------------------------------------

describe('AgentRuntime.stream() — structured output', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('includes structuredOutput in the finish chunk when schema is configured', async () => {
		const expected = { answer: 'Paris', score: 0.95 };
		streamText.mockReturnValue(makeStreamSuccessWithOutput(expected));

		const { runtime } = createStructuredRuntime();
		const { stream } = await runtime.stream('What is the capital of France?');
		const chunks = await collectChunks(stream);

		const finish = chunks.find((c) => c.type === 'finish') as StreamChunk & {
			type: 'finish';
		};
		expect(finish).toBeDefined();
		expect(finish.structuredOutput).toEqual(expected);
	});

	it('does not include structuredOutput in the finish chunk when no schema is configured', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const { runtime } = createRuntime();
		const { stream } = await runtime.stream('hello');
		const chunks = await collectChunks(stream);

		const finish = chunks.find((c) => c.type === 'finish') as StreamChunk & {
			type: 'finish';
		};
		expect(finish).toBeDefined();
		expect(finish.structuredOutput).toBeUndefined();
	});

	it('passes the output spec to streamText when schema is configured', async () => {
		streamText.mockReturnValue(makeStreamSuccessWithOutput({ answer: 'test', score: 1 }));

		const { runtime } = createStructuredRuntime();
		const { stream } = await runtime.stream('hello');
		await collectChunks(stream);

		const callArgs = (streamText.mock.calls[0] as [unknown, unknown])[0] as Record<string, unknown>;
		expect(callArgs.output).toBeDefined();
		expect(callArgs.output).toEqual(
			expect.objectContaining({ _type: 'object', schema: testSchema }),
		);
	});

	it('does not pass the output spec to streamText when no schema is configured', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const { runtime } = createRuntime();
		const { stream } = await runtime.stream('hello');
		await collectChunks(stream);
		const callArgs = (streamText.mock.calls[0] as [unknown, unknown])[0] as Record<string, unknown>;
		expect(callArgs.output).toBeUndefined();
	});

	it('preserves structuredOutput through tool-call iterations', async () => {
		const tool: BuiltTool = {
			name: 'add',
			description: 'Add numbers',
			inputSchema: z.object({ a: z.number(), b: z.number() }),
			handler: async (input: unknown) => {
				const { a, b } = input as { a: number; b: number };
				return await Promise.resolve({ sum: a + b });
			},
		};

		const expected = { answer: '5', score: 1 };
		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([
					{ type: 'tool-call', toolCallId: 'tc-1', toolName: 'add', args: { a: 2, b: 3 } },
				]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'add',
									args: { a: 2, b: 3 },
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-1', toolName: 'add', input: { a: 2, b: 3 } },
				]),
			})
			.mockReturnValueOnce(makeStreamSuccessWithOutput(expected, 'The sum is 5'));

		const { runtime } = createStructuredRuntime({ tools: [tool] });
		const { stream } = await runtime.stream('What is 2 + 3?');
		const chunks = await collectChunks(stream);

		const finish = chunks.find((c) => c.type === 'finish') as StreamChunk & {
			type: 'finish';
		};
		expect(finish).toBeDefined();
		expect(finish.structuredOutput).toEqual(expected);
	});
});

// ---------------------------------------------------------------------------
// Structured output — resume()
// ---------------------------------------------------------------------------

describe('AgentRuntime.resume() — structured output', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('returns structuredOutput after resume in generate mode', async () => {
		const interruptibleTool = makeInterruptibleTool();
		const { runtime } = createStructuredRuntime({ tools: [interruptibleTool] });

		// First call: LLM triggers the interruptible tool
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'approve', { question: 'proceed?' }),
		);

		const firstResult = await runtime.generate('Should I proceed?');
		expect(firstResult.finishReason).toBe('tool-calls');
		expect(firstResult.pendingSuspend).toBeDefined();

		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		// Second call (after resume): LLM returns final answer with structured output
		const expected = { answer: 'Approved', score: 1 };
		generateText.mockResolvedValueOnce(makeGenerateSuccessWithOutput(expected, 'Approved'));

		const resumeResult = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId },
		);

		expect(resumeResult.structuredOutput).toEqual(expected);
		expect(resumeResult.finishReason).toBe('stop');
	});

	it('returns structuredOutput after resume in stream mode', async () => {
		const interruptibleTool = makeInterruptibleTool();
		const { runtime } = createStructuredRuntime({ tools: [interruptibleTool] });

		// First call: LLM triggers the interruptible tool (via generate to get the runId)
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'approve', { question: 'proceed?' }),
		);

		const firstResult = await runtime.generate('Should I proceed?');
		expect(firstResult.pendingSuspend).toBeDefined();

		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		// Second call (after resume): LLM streams final answer with structured output
		const expected = { answer: 'Approved', score: 1 };
		streamText.mockReturnValueOnce(makeStreamSuccessWithOutput(expected, 'Approved'));

		const resumed = await runtime.resume('stream', { approved: true }, { runId, toolCallId });

		const chunks = await collectChunks(resumed.stream as ReadableStream<unknown>);
		const finish = chunks.find((c) => c.type === 'finish') as StreamChunk & {
			type: 'finish';
		};
		expect(finish).toBeDefined();
		expect(finish.structuredOutput).toEqual(expected);
	});
});

// ---------------------------------------------------------------------------
// Eager input streaming
// ---------------------------------------------------------------------------

/* eslint-disable @typescript-eslint/require-await */
describe('providerOptions — tool adapter', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('forwards providerOptions to the AI SDK tool when set', () => {
		const ai = aiModule as unknown as { tool: Mock };
		const adapter = { toAiSdkTools };

		const builtTool: BuiltTool = {
			name: 'set_code',
			description: 'Set code in editor',
			inputSchema: z.object({ code: z.string() }),
			providerOptions: { anthropic: { eagerInputStreaming: true } },
			handler: async () => ({ ok: true }),
		};

		const result = adapter.toAiSdkTools([builtTool]);

		expect(ai.tool).toHaveBeenCalledWith(
			expect.objectContaining({
				providerOptions: {
					anthropic: { eagerInputStreaming: true },
				},
			}),
		);
		expect(result).toHaveProperty('set_code');
	});

	it('forwards arbitrary provider options (not just Anthropic)', () => {
		const ai = aiModule as unknown as { tool: Mock };
		const adapter = { toAiSdkTools };

		const builtTool: BuiltTool = {
			name: 'draw',
			description: 'Draw an image',
			inputSchema: z.object({ prompt: z.string() }),
			providerOptions: { openai: { strict: true } },
			handler: async () => ({ ok: true }),
		};

		adapter.toAiSdkTools([builtTool]);

		expect(ai.tool).toHaveBeenCalledWith(
			expect.objectContaining({
				providerOptions: { openai: { strict: true }, anthropic: { eagerInputStreaming: false } },
			}),
		);
	});

	it('defaults providerOptions to the Anthropic quirk default when the tool sets none', () => {
		const ai = aiModule as unknown as { tool: Mock };
		const adapter = { toAiSdkTools };

		const builtTool: BuiltTool = {
			name: 'search',
			description: 'Search',
			inputSchema: z.object({ query: z.string() }),
			handler: async () => ({ results: [] }),
		};

		adapter.toAiSdkTools([builtTool]);

		expect(ai.tool).toHaveBeenCalledWith(
			expect.objectContaining({
				providerOptions: { anthropic: { eagerInputStreaming: false } },
			}),
		);
	});
});

describe('Tool builder — providerOptions', () => {
	it('sets providerOptions via .providerOptions()', () => {
		const built = new Tool('set_code')
			.description('Set code')
			.input(z.object({ code: z.string() }))
			.providerOptions({ anthropic: { eagerInputStreaming: true } })
			.handler(async () => ({ ok: true }))
			.build();

		expect(built.providerOptions).toEqual({ anthropic: { eagerInputStreaming: true } });
	});

	it('merges multiple .providerOptions() calls', () => {
		const built = new Tool('multi')
			.description('Multi-provider tool')
			.input(z.object({ text: z.string() }))
			.providerOptions({ anthropic: { eagerInputStreaming: true } })
			.providerOptions({ openai: { strict: true } })
			.handler(async () => ({ ok: true }))
			.build();

		expect(built.providerOptions).toEqual({
			anthropic: { eagerInputStreaming: true },
			openai: { strict: true },
		});
	});

	it('defaults to undefined when .providerOptions() is not called', () => {
		const built = new Tool('search')
			.description('Search')
			.input(z.object({ query: z.string() }))
			.handler(async () => ({ results: [] }))
			.build();

		expect(built.providerOptions).toBeUndefined();
	});
});
/* eslint-enable @typescript-eslint/require-await */

// ---------------------------------------------------------------------------
// Runtime validation — input schema
// ---------------------------------------------------------------------------

describe('AgentRuntime — runtime input schema validation', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('surfaces a ZodError as a tool error outcome when LLM provides invalid input', async () => {
		// Tool expects { id: z.string() } but LLM will provide { id: 123 } (wrong type)
		const strictTool: BuiltTool = {
			name: 'strict',
			description: 'strict',
			inputSchema: z.object({ id: z.string() }),
			handler: async () => {
				return await Promise.resolve({ ok: true });
			},
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'strict', { id: 123 }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtimeWithTool = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [strictTool],
		});

		const result = await runtimeWithTool.generate('go');
		// The ZodError is caught by Promise.allSettled; the loop continues and
		// the LLM responds with 'done' on the next turn.
		expect(result.finishReason).toBe('stop');

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		expect(assistantMsg).toBeDefined();
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain('Expected string, received number');
	});
});

// ---------------------------------------------------------------------------
// Runtime validation — JSON Schema input validation
// ---------------------------------------------------------------------------

describe('AgentRuntime — runtime JSON Schema input validation', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('passes valid input through without error', async () => {
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				properties: { id: { type: 'string' } },
				required: ['id'],
			},
			handler: async () => await Promise.resolve({ ok: true }),
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', { id: 'abc' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');

		// No error — the tool ran successfully
		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		expect(assistantMsg).toBeDefined();
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('resolved');
	});

	it('parses stringified JSON object tool input before validating and replaying', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				properties: { id: { type: 'string' } },
				required: ['id'],
			},
			handler: handlerFn,
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', '{"id":"abc"}'))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');
		expect(handlerFn).toHaveBeenCalledWith(
			expect.objectContaining({ id: 'abc' }),
			expect.anything(),
		);
		expect(getFirstReplayedToolCallInput(1)).toEqual({ id: 'abc' });
	});

	it('surfaces malformed string tool input as a retryable tool error', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			handler: handlerFn,
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', '{bad json'))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');
		expect(handlerFn).not.toHaveBeenCalled();
		expect(getFirstReplayedToolCallInput(1)).toEqual({});

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain(
			'Tool input must be a valid JSON object string.',
		);
	});

	it('surfaces stringified non-object tool input as a retryable tool error', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			handler: handlerFn,
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', '["not an object"]'))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');
		expect(handlerFn).not.toHaveBeenCalled();
		expect(getFirstReplayedToolCallInput(1)).toEqual({});

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain(
			'Tool input must be a JSON object, got array.',
		);
	});

	it('surfaces a validation error as a tool error outcome when LLM provides the wrong type', async () => {
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				properties: { id: { type: 'string' } },
				required: ['id'],
			},
			handler: async () => await Promise.resolve({ ok: true }),
		};

		generateText
			// LLM sends { id: 123 } — a number where a string is required
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', { id: 123 }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		expect(assistantMsg).toBeDefined();
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain('Invalid tool input');
	});

	it('surfaces a validation error when a required property is missing', async () => {
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				properties: {
					name: { type: 'string' },
					age: { type: 'integer' },
				},
				required: ['name', 'age'],
			},
			handler: async () => await Promise.resolve({ ok: true }),
		};

		generateText
			// LLM omits the required 'age' field
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', { name: 'Alice' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');
		expect(result.finishReason).toBe('stop');

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain('Invalid tool input');
	});

	it('does not invoke the handler when JSON Schema validation fails', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				properties: { count: { type: 'integer', minimum: 1 } },
				required: ['count'],
			},
			handler: handlerFn,
		};

		generateText
			// LLM sends count: 0, which fails the minimum: 1 constraint
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', { count: 0 }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		await runtime.generate('go');
		expect(handlerFn).not.toHaveBeenCalled();
	});

	it('rejects unknown keys against a closed JSON Schema', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: {
				type: 'object',
				additionalProperties: false,
				properties: { query: { type: 'string' } },
			},
			handler: handlerFn,
		};

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCall('tc-1', 'json_tool', { query: 'cats', extra: true }),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		await runtime.generate('go');
		expect(handlerFn).not.toHaveBeenCalled();
	});

	it('names the tool schema as the defect when it cannot be compiled', async () => {
		const tool: BuiltTool = {
			name: 'json_tool',
			description: 'json tool',
			inputSchema: { type: 'objct' } as unknown as JSONSchema7,
			handler: async () => await Promise.resolve({ ok: true }),
		};

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'json_tool', { query: 'cats' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state === 'rejected' && call.error).toContain(
			'input schema that could not be compiled',
		);
	});
});

// ---------------------------------------------------------------------------
// Tool builder — JSON Schema input integration
//
// Mirrors the resolveNodeTool() code path in node-tool-factory.ts where the
// input schema is a raw JSON Schema object (converted from Zod by ToolFromNode).
// ---------------------------------------------------------------------------

describe('AgentRuntime — Tool builder with JSON Schema input', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('passes valid input to the handler when built via Tool builder', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ found: true });

		const tool = new Tool('lookup')
			.description('Look up a record by id')
			.input({
				type: 'object',
				properties: { id: { type: 'string' } },
				required: ['id'],
			})
			.handler(handlerFn)
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'lookup', { id: 'abc-123' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');

		expect(result.finishReason).toBe('stop');
		expect(handlerFn).toHaveBeenCalledWith(
			expect.objectContaining({ id: 'abc-123' }),
			expect.anything(),
		);

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('resolved');
	});

	it('produces a tool error when the LLM sends input that fails JSON Schema validation', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ found: true });

		const tool = new Tool('lookup')
			.description('Look up a record by id')
			.input({
				type: 'object',
				properties: {
					id: { type: 'string' },
					count: { type: 'integer', minimum: 1 },
				},
				required: ['id', 'count'],
			})
			.handler(handlerFn)
			.build();

		generateText
			// LLM sends count: 0 (violates minimum: 1) and id as a number (wrong type)
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'lookup', { id: 42, count: 0 }))
			.mockResolvedValueOnce(makeGenerateSuccess('corrected'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');

		expect(result.finishReason).toBe('stop');
		// Handler must not be called — validation should block execution
		expect(handlerFn).not.toHaveBeenCalled();

		const assistantMsg = result.messages.find(
			(m) =>
				isLlmMessage(m) && m.role === 'assistant' && m.content.some((c) => c.type === 'tool-call'),
		) as Message;
		const call = assistantMsg.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(call.state).toBe('rejected');
		expect(call.state === 'rejected' && call.error).toContain('Invalid tool input');
	});

	it('validates enum and pattern constraints defined in JSON Schema', async () => {
		const handlerFn = vi.fn().mockResolvedValue({ ok: true });

		const tool = new Tool('set_status')
			.description('Set the status of a record')
			.input({
				type: 'object',
				properties: {
					status: { type: 'string', enum: ['active', 'inactive', 'pending'] },
				},
				required: ['status'],
			})
			.handler(handlerFn)
			.build();

		// First call: invalid enum value
		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'set_status', { status: 'deleted' }))
			// Second call: valid enum value after self-correction
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-2', 'set_status', { status: 'inactive' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
		});

		const result = await runtime.generate('go');

		expect(result.finishReason).toBe('stop');
		// Handler called exactly once — only for the valid input
		expect(handlerFn).toHaveBeenCalledTimes(1);
		expect(handlerFn).toHaveBeenCalledWith(
			expect.objectContaining({ status: 'inactive' }),
			expect.anything(),
		);
	});
});

// ---------------------------------------------------------------------------
// Runtime validation — resume data schema
// ---------------------------------------------------------------------------

describe('AgentRuntime — runtime resume data schema validation', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('surfaces a persisted-schema error when consumer provides invalid resume data', async () => {
		const tool = makeInterruptibleTool(); // has resumeSchema: z.object({ approved: z.boolean() })

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'approve', { question: 'ok?' }),
		);

		const runtimeWithTool = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
			checkpointStorage: 'memory',
		});

		const firstResult = await runtimeWithTool.generate('go');
		expect(firstResult.pendingSuspend).toBeDefined();
		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		// Resume with invalid data — missing required 'approved' field
		const resumeResultPromise = runtimeWithTool.resume(
			'generate',
			{ notApproved: 'wrong' },
			{ runId, toolCallId },
		);
		await expect(resumeResultPromise).rejects.toThrow('Invalid resume payload');
	});

	it('strips keys the persisted schema does not declare instead of rejecting', async () => {
		let observedResumeData: unknown;
		const tool: BuiltTool = {
			name: 'approve',
			description: 'Requires approval',
			inputSchema: z.object({ question: z.string() }),
			suspendSchema: z.object({ question: z.string() }),
			resumeSchema: z.object({ approved: z.boolean() }),
			handler: async (_input: unknown, ctx: unknown) => {
				const { suspend, resumeData } = ctx as InterruptibleToolContext;
				if (!resumeData) return await suspend({ question: 'approve?' });
				observedResumeData = resumeData;
				return { approved: true };
			},
		};

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'approve', { question: 'ok?' }),
		);

		const runtimeWithTool = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [tool],
			checkpointStorage: 'memory',
		});

		const first = await runtimeWithTool.generate('go');
		const { runId, toolCallId } = first.pendingSuspend![0];

		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		const payload = { approved: true, userInput: 'looks good' };
		await runtimeWithTool.resume('generate', payload, { runId, toolCallId });

		expect(observedResumeData).toEqual({ approved: true });
		expect(payload).toEqual({ approved: true, userInput: 'looks good' });
	});
});

// ---------------------------------------------------------------------------
// AgentRuntime — tool approval (HITL wrapper)
// ---------------------------------------------------------------------------

describe('AgentRuntime — tool approval (HITL wrapper)', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('suspends when a tool has .requireApproval() set', async () => {
		const builtApprovalTool = new ToolBuilder('delete')
			.description('Delete a record')
			.input(z.object({ id: z.string() }))
			.requireApproval()
			.handler(async ({ id }: { id: string }) => {
				return await Promise.resolve({ deleted: id });
			})
			.build();
		const approvalTool = {
			...builtApprovalTool,
			metadata: { displayName: 'Delete record' },
		};

		generateText.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'delete', { id: 'rec-1' }));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [approvalTool],
			checkpointStorage: 'memory',
		});

		const result = await runtime.generate('Delete record rec-1');
		expect(result.finishReason).toBe('tool-calls');
		expect(result.pendingSuspend).toHaveLength(1);
		expect(result.pendingSuspend![0].toolName).toBe('delete');
		expect(result.pendingSuspend![0].suspendPayload).toMatchObject({
			type: 'approval',
			toolName: 'delete',
			displayName: 'Delete record',
			args: { id: 'rec-1' },
		});
	});

	it('executes the original handler after approval', async () => {
		const approvalTool = new ToolBuilder('delete')
			.description('Delete a record')
			.input(z.object({ id: z.string() }))
			.requireApproval()
			.handler(async ({ id }: { id: string }) => {
				return await Promise.resolve({ deleted: id });
			})
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'delete', { id: 'rec-1' }))
			.mockResolvedValueOnce(makeGenerateSuccess('Done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [approvalTool],
			checkpointStorage: 'memory',
		});

		const firstResult = await runtime.generate('Delete record rec-1');
		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		const resumeResult = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId },
		);
		expect(resumeResult.finishReason).toBe('stop');
		expect(resumeResult.toolCalls).toEqual(
			expect.arrayContaining([
				expect.objectContaining({ tool: 'delete', output: { deleted: 'rec-1' } }),
			]),
		);
	});

	it('returns declined message to LLM when approval is rejected', async () => {
		const approvalTool = new ToolBuilder('delete')
			.description('Delete a record')
			.input(z.object({ id: z.string() }))
			.requireApproval()
			.handler(async ({ id }: { id: string }) => {
				return await Promise.resolve({ deleted: id });
			})
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'delete', { id: 'rec-1' }))
			.mockResolvedValueOnce(makeGenerateSuccess('Declined'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [approvalTool],
			checkpointStorage: 'memory',
		});

		const firstResult = await runtime.generate('Delete record rec-1');
		expect(firstResult.finishReason).toBe('tool-calls');
		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		const resumeResult = await runtime.resume(
			'generate',
			{ approved: false },
			{ runId, toolCallId },
		);
		expect(resumeResult.finishReason).toBe('stop');
	});

	it('does not start tool execution before approval is granted', async () => {
		const handler = vi.fn(async ({ id }: { id: string }) => {
			return await Promise.resolve({ deleted: id });
		});
		const approvalTool = new ToolBuilder('delete')
			.description('Delete a record')
			.input(z.object({ id: z.string() }))
			.requireApproval()
			.handler(handler)
			.build();

		streamText.mockReturnValue({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'checking...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{
								type: 'tool-call',
								toolCallId: 'tc-1',
								toolName: 'delete',
								args: { id: 'rec-1' },
							},
						],
					},
				],
			}),
			toolCalls: Promise.resolve([
				{ toolCallId: 'tc-1', toolName: 'delete', input: { id: 'rec-1' } },
			]),
		});

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [approvalTool],
			checkpointStorage: 'memory',
		});

		const { stream: readableStream } = await runtime.stream('Delete record rec-1');
		const chunks = await collectChunks(readableStream);

		expect(handler).not.toHaveBeenCalled();
		expect(chunks.map((c) => c.type)).toContain('tool-call-suspended');
		expect(chunks.map((c) => c.type)).not.toContain('tool-execution-start');
	});

	it('starts tool execution when conditional approval allows the call immediately', async () => {
		const handler = vi.fn(async ({ id }: { id: string }) => {
			return await Promise.resolve({ found: id });
		});
		const conditionalApprovalTool = new ToolBuilder('lookup')
			.description('Look up a record')
			.input(
				z.object({
					id: z.string(),
					password: z.string(),
					nested: z.object({ apiKey: z.string() }),
				}),
			)
			.needsApprovalFn(async ({ id }: { id: string }) => {
				return await Promise.resolve(id === 'secret');
			})
			.handler(handler)
			.build();
		const startEvents: Array<AgentEventData & { type: AgentEvent.ToolExecutionStart }> = [];
		const eventBus = new AgentEventBus();
		eventBus.on(AgentEvent.ToolExecutionStart, (event) => {
			startEvents.push(event as AgentEventData & { type: AgentEvent.ToolExecutionStart });
		});
		const toolInput = {
			id: 'public',
			password: 'plain-secret-password',
			nested: { apiKey: 'secret-api-key' },
		};

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'checking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'lookup',
									args: toolInput,
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([{ toolCallId: 'tc-1', toolName: 'lookup', input: toolInput }]),
			})
			.mockReturnValueOnce(makeStreamSuccess('Done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			tools: [conditionalApprovalTool],
			eventBus,
			checkpointStorage: 'memory',
		});

		const { stream: readableStream } = await runtime.stream('Look up public');
		const chunks = await collectChunks(readableStream);

		expect(handler).toHaveBeenCalledWith(
			toolInput,
			expect.objectContaining({ toolCallId: 'tc-1' }),
		);
		expect(startEvents[0]?.args).toEqual({
			id: 'public',
			password: 'plain-secret-password',
			nested: { apiKey: 'secret-api-key' },
		});
		expect(chunks.map((c) => c.type)).toContain('tool-execution-start');
		expect(chunks.map((c) => c.type)).toContain('tool-execution-end');
		expect(chunks.map((c) => c.type)).not.toContain('tool-call-suspended');
	});
});

// ---------------------------------------------------------------------------
// External abort signal
// ---------------------------------------------------------------------------

describe('external abort signal', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('should cancel run when external signal fires', async () => {
		const external = new AbortController();

		generateText.mockImplementation(
			async (opts: { abortSignal?: AbortSignal }) =>
				await new Promise((_, reject) => {
					const signal = opts.abortSignal;
					if (signal?.aborted) {
						reject(new Error('aborted'));
						return;
					}
					signal?.addEventListener('abort', () => {
						reject(new Error('aborted'));
					});
				}),
		);

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
		});

		const resultPromise = runtime.generate('hello', {
			persistence: { resourceId: 'user1', threadId: 'thread1' },
			abortSignal: external.signal,
		});

		external.abort();
		const result = await resultPromise;
		expect(result.finishReason).toBe('error');
	});

	it('marks tool ctx abortSignal as aborted when a generate run is aborted', async () => {
		const external = new AbortController();
		let toolSignal: AbortSignal | undefined;
		let toolStarted!: () => void;
		const waitForToolStart = new Promise<void>((resolve) => {
			toolStarted = resolve;
		});
		let toolObservedAbort!: () => void;
		const waitForToolAbort = new Promise<void>((resolve) => {
			toolObservedAbort = resolve;
		});
		const abortAwareTool: BuiltTool = {
			name: 'wait_for_abort',
			description: 'Waits until the run aborts',
			inputSchema: z.object({}),
			handler: async (_input, ctx) => {
				toolSignal = (ctx as ToolContext).abortSignal;
				toolStarted();
				await new Promise<void>((resolve) => {
					toolSignal?.addEventListener(
						'abort',
						() => {
							toolObservedAbort();
							resolve();
						},
						{ once: true },
					);
				});
				return { aborted: toolSignal?.aborted };
			},
		};

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'wait_for_abort', args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const { runtime } = createRuntimeWithTools([abortAwareTool], 1);
		const resultPromise = runtime.generate('start tool', { abortSignal: external.signal });

		await waitForToolStart;
		expect(toolSignal?.aborted).toBe(false);

		external.abort();
		await waitForToolAbort;
		expect(toolSignal?.aborted).toBe(true);

		await resultPromise;
	});

	it('marks tool ctx abortSignal as aborted when a stream run is aborted', async () => {
		const external = new AbortController();
		let toolSignal: AbortSignal | undefined;
		let toolStarted!: () => void;
		const waitForToolStart = new Promise<void>((resolve) => {
			toolStarted = resolve;
		});
		let toolObservedAbort!: () => void;
		const waitForToolAbort = new Promise<void>((resolve) => {
			toolObservedAbort = resolve;
		});
		const abortAwareTool: BuiltTool = {
			name: 'wait_for_abort',
			description: 'Waits until the run aborts',
			inputSchema: z.object({}),
			handler: async (_input, ctx) => {
				toolSignal = (ctx as ToolContext).abortSignal;
				toolStarted();
				await new Promise<void>((resolve) => {
					toolSignal?.addEventListener(
						'abort',
						() => {
							toolObservedAbort();
							resolve();
						},
						{ once: true },
					);
				});
				return { aborted: toolSignal?.aborted };
			},
		};

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'checking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'wait_for_abort',
									args: {},
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([{ toolCallId: 'tc-1', toolName: 'wait_for_abort', input: {} }]),
			})
			.mockReturnValueOnce(makeStreamSuccess('done'));

		const { runtime } = createRuntimeWithTools([abortAwareTool], 1);
		const { stream } = await runtime.stream('start tool', { abortSignal: external.signal });
		const collectPromise = collectChunks(stream);

		await waitForToolStart;
		expect(toolSignal?.aborted).toBe(false);

		external.abort();
		await waitForToolAbort;
		expect(toolSignal?.aborted).toBe(true);

		await collectPromise;
	});

	it('removes external abort listener after a stream completes', async () => {
		const external = new AbortController();
		const removeListener = vi.spyOn(external.signal, 'removeEventListener');
		streamText.mockReturnValue(makeStreamSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
		});

		const result = await runtime.stream('hello', {
			persistence: { resourceId: 'user1', threadId: 'thread1' },
			abortSignal: external.signal,
		});
		await collectChunks(result.stream);
		await new Promise((resolve) => setTimeout(resolve, 0));

		expect(removeListener).toHaveBeenCalledTimes(1);
		external.abort();
		expect(runtime.getState().status).toBe('success');
	});
});

// ---------------------------------------------------------------------------
// Abort during a tool batch (generate/stream parity)
// ---------------------------------------------------------------------------

describe('AgentRuntime — abort during a tool batch', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	// Two sequential (concurrency 1) tool calls where the FIRST aborts the run.
	// Batch 0 settles, then the executor hits the abort at the next batch
	// boundary and throws. This is the path that used to make the stream report
	// `failed` + an Error event; it must now resolve to `cancelled` with no Error
	// event — matching how generate has always handled it.
	function makeAbortingToolPair(bus: AgentEventBus) {
		const secondHandler = vi.fn(async () => {
			await Promise.resolve();
			return 'second';
		});
		const tools: BuiltTool[] = [
			{
				name: 'abort_now',
				description: 'Aborts the run while the batch is executing',
				inputSchema: z.object({}),
				handler: async () => {
					await Promise.resolve();
					bus.abort();
					return 'ok';
				},
			},
			{
				name: 'second_tool',
				description: 'Should never run once an earlier tool in the batch aborts',
				inputSchema: z.object({}),
				handler: secondHandler,
			},
		];
		return { tools, secondHandler };
	}

	const TOOL_CALLS = [
		{ toolCallId: 'tc-1', toolName: 'abort_now', args: {} },
		{ toolCallId: 'tc-2', toolName: 'second_tool', args: {} },
	];

	function makeStreamWithToolCalls(
		toolCalls: Array<{ toolCallId: string; toolName: string; args: Record<string, unknown> }>,
	) {
		return {
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'calling tools...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: toolCalls.map((tc) => ({
							type: 'tool-call',
							toolCallId: tc.toolCallId,
							toolName: tc.toolName,
							args: tc.args,
						})),
					},
				],
			}),
			toolCalls: Promise.resolve(
				toolCalls.map((tc) => ({
					toolCallId: tc.toolCallId,
					toolName: tc.toolName,
					input: tc.args,
				})),
			),
		};
	}

	it('stream: resolves to "cancelled" (not "failed") and emits no Error event', async () => {
		const bus = new AgentEventBus();
		const { tools, secondHandler } = makeAbortingToolPair(bus);
		const { runtime } = createRuntimeWithTools(tools, 1, bus);

		streamText.mockReturnValueOnce(makeStreamWithToolCalls(TOOL_CALLS));

		const errorEvents: AgentEventData[] = [];
		bus.on(AgentEvent.Error, (d) => errorEvents.push(d));

		const { stream } = await runtime.stream('go');
		await collectChunks(stream);

		expect(runtime.getState().status).toBe('cancelled');
		// A user abort is not an error: no AgentEvent.Error must be emitted.
		expect(errorEvents).toHaveLength(0);
		// Processing stops at the first suspension/abort, so the later tool in the
		// batch never executes.
		expect(secondHandler).not.toHaveBeenCalled();
	});

	it('stream: still closes cleanly with an error chunk then a finish chunk', async () => {
		const bus = new AgentEventBus();
		const { tools } = makeAbortingToolPair(bus);
		const { runtime } = createRuntimeWithTools(tools, 1, bus);

		streamText.mockReturnValueOnce(makeStreamWithToolCalls(TOOL_CALLS));

		const { stream } = await runtime.stream('go');
		const chunks = await collectChunks(stream);

		expect(chunks.some((c) => c.type === 'error')).toBe(true);

		const finishChunks = chunks.filter((c) => c.type === 'finish');
		expect(finishChunks.length).toBeGreaterThan(0);
		const finish = finishChunks[finishChunks.length - 1] as StreamChunk & {
			type: 'finish';
			finishReason: string;
		};
		expect(finish.finishReason).toBe('error');
	});

	it('generate: resolves to "cancelled" with no Error event (parity with stream)', async () => {
		const bus = new AgentEventBus();
		const { tools, secondHandler } = makeAbortingToolPair(bus);
		const { runtime } = createRuntimeWithTools(tools, 1, bus);

		generateText.mockResolvedValueOnce(makeGenerateWithToolCalls(TOOL_CALLS));

		const errorEvents: AgentEventData[] = [];
		bus.on(AgentEvent.Error, (d) => errorEvents.push(d));

		const result = await runtime.generate('go');

		expect(result.finishReason).toBe('error');
		expect(runtime.getState().status).toBe('cancelled');
		expect(errorEvents).toHaveLength(0);
		expect(secondHandler).not.toHaveBeenCalled();
	});

	it('cleans up a child that suspends after the parent abort race settles', async () => {
		const bus = new AgentEventBus();
		let releaseHandler!: () => void;
		let markHandlerStarted!: () => void;
		const handlerGate = new Promise<void>((resolve) => (releaseHandler = resolve));
		const handlerStarted = new Promise<void>((resolve) => (markHandlerStarted = resolve));
		const onCancellation = makeCancellationSpy();
		const tool = makeSuspendingTool(
			'suspend_tool',
			async (_input, ctx) => {
				markHandlerStarted();
				await handlerGate;
				return await ctx.suspend(
					{ reason: 'needs approval' },
					{ continuation: { childRunId: 'late-child-run' } },
				);
			},
			onCancellation,
		);
		const { runtime } = createRuntimeWithTools([tool], 1, bus);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'suspend_tool', { value: 'a' }),
		);

		const resultPromise = runtime.generate('go');
		await handlerStarted;
		bus.abort();
		await expect(resultPromise).resolves.toMatchObject({ finishReason: 'error' });
		releaseHandler();

		await vi.waitFor(() =>
			expect(onCancellation).toHaveBeenCalledWith(
				{ value: 'a' },
				expect.objectContaining({
					continuation: { childRunId: 'late-child-run' },
					suspendPayload: { reason: 'needs approval' },
				}),
			),
		);
	});

	it('cleans up a repeated suspension when abort lands during schema validation', async () => {
		const bus = new AgentEventBus();
		let releaseParse!: () => void;
		let markParseStarted!: () => void;
		const parseGate = new Promise<void>((resolve) => (releaseParse = resolve));
		const parseStarted = new Promise<void>((resolve) => (markParseStarted = resolve));
		const onCancellation = makeCancellationSpy();
		const tool: BuiltTool = {
			name: 'suspend_tool',
			description: 'Suspends twice',
			inputSchema: z.object({ value: z.string().optional() }),
			suspendSchema: z.object({ stage: z.string() }).superRefine(async ({ stage }) => {
				if (stage !== 'second') return;
				markParseStarted();
				await parseGate;
			}),
			resumeSchema: z.object({ approved: z.boolean() }),
			handler: async (_input, rawCtx) => {
				const ctx = rawCtx as InterruptibleToolContext;
				const stage = ctx.resumeData ? 'second' : 'first';
				return await ctx.suspend({ stage }, { continuation: { childRunId: stage } });
			},
			onCancellation,
		};
		const { runtime } = createRuntimeWithTools([tool], 1, bus);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'suspend_tool', { value: 'a' }),
		);

		const first = await runtime.generate('go');
		const resume = runtime.resume('generate', { approved: true }, first.pendingSuspend![0]);
		await parseStarted;
		bus.abort();
		releaseParse();
		const result = await resume;

		expect(result.pendingSuspend).toBeUndefined();
		expect(runtime.getState().status).toBe('cancelled');
		expect(onCancellation).toHaveBeenCalledWith(
			{ value: 'a' },
			expect.objectContaining({
				continuation: { childRunId: 'second' },
				suspendPayload: { stage: 'second' },
			}),
		);
	});

	it('cleans up a resumed child suspension when the parent aborts before re-suspending', async () => {
		const bus = new AgentEventBus();
		const continuation = { childRunId: 'child-run-on-resume' };
		const resumeSchema = z.object({ approved: z.boolean() });
		const onCancellation = makeCancellationSpy();
		const tool = makeSuspendingTool(
			'suspend_tool',
			async (_input, ctx) => {
				if (!ctx.resumeData) {
					return await ctx.suspend({ reason: 'needs approval' }, { continuation, resumeSchema });
				}
				bus.abort();
				return await new Promise<never>(() => undefined);
			},
			onCancellation,
		);
		const { runtime } = createRuntimeWithTools([tool], 1, bus);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'suspend_tool', { value: 'a' }),
		);

		const first = await runtime.generate('go');
		const { runId, toolCallId } = first.pendingSuspend![0];
		const result = await runtime.resume('generate', { approved: true }, { runId, toolCallId });

		expect(result.finishReason).toBe('error');
		expect(runtime.getState().status).toBe('cancelled');
		expect(onCancellation).toHaveBeenCalledWith(
			{ value: 'a' },
			expect.objectContaining({
				cancellation: { message: 'Run aborted' },
				continuation,
				toolCallId: 'tc-1',
				suspendPayload: { reason: 'needs approval' },
				resumeSchema: expect.objectContaining({
					type: 'object',
					properties: { approved: expect.objectContaining({ type: 'boolean' }) },
				}),
			}),
		);
	});

	it('cleans up suspended siblings when an unexecuted tool aborts during resume', async () => {
		const bus = new AgentEventBus();
		const onCancellation = makeCancellationSpy();
		const suspendTool = makeSuspendingTool(
			'suspend_tool',
			async (rawInput, ctx) => {
				if (ctx.resumeData) return { approved: true };
				const { value } = rawInput as { value: string };
				return await ctx.suspend(
					{ reason: 'needs approval' },
					{ continuation: { childRunId: `child-${value}` } },
				);
			},
			onCancellation,
		);
		const abortTool = makeMockTool('abort_tool', async () => {
			bus.abort();
			return await new Promise<never>(() => undefined);
		});
		const { runtime } = createRuntimeWithTools([suspendTool, abortTool], 2, bus);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: { value: 'a' } },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: { value: 'b' } },
				{ toolCallId: 'tc-3', toolName: 'abort_tool', args: {} },
			]),
		);

		const first = await runtime.generate('go');
		const result = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId: first.runId, toolCallId: 'tc-1' },
		);

		expect(result.pendingSuspend).toBeUndefined();
		expect(runtime.getState().status).toBe('cancelled');
		expect(onCancellation).toHaveBeenCalledOnce();
		expect(onCancellation).toHaveBeenCalledWith(
			{ value: 'b' },
			expect.objectContaining({
				continuation: { childRunId: 'child-b' },
				toolCallId: 'tc-2',
			}),
		);
	});

	it('rolls back a suspension when abort lands during checkpoint persistence', async () => {
		const bus = new AgentEventBus();
		let releaseSave!: () => void;
		let markSaveStarted!: () => void;
		const saveGate = new Promise<void>((resolve) => (releaseSave = resolve));
		const saveStarted = new Promise<void>((resolve) => (markSaveStarted = resolve));
		const checkpointStore: CheckpointStore = {
			save: vi.fn(async () => {
				markSaveStarted();
				await saveGate;
			}),
			load: vi.fn().mockResolvedValue(undefined),
			delete: vi.fn().mockResolvedValue(undefined),
		};
		const onCancellation = makeCancellationSpy();
		const tool = makeSuspendingTool(
			'suspend_tool',
			async (_input, ctx) =>
				await ctx.suspend(
					{ reason: 'needs approval' },
					{ continuation: { childRunId: 'child-during-save' } },
				),
			onCancellation,
		);
		const runtime = createRuntimeWithCheckpointStore([tool], checkpointStore, bus);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'suspend_tool', { value: 'a' }),
		);

		const resultPromise = runtime.generate('go');
		await saveStarted;
		bus.abort();
		releaseSave();
		const result = await resultPromise;

		expect(result.pendingSuspend).toBeUndefined();
		expect(runtime.getState().status).toBe('cancelled');
		expect(onCancellation).toHaveBeenCalledOnce();
		expect(checkpointStore.delete).toHaveBeenCalledWith(result.runId);
	});

	it('cleans up child continuations when checkpoint persistence fails', async () => {
		const persistenceError = new Error('checkpoint unavailable');
		const checkpointStore: CheckpointStore = {
			save: vi.fn(async () => await Promise.reject(persistenceError)),
			load: vi.fn().mockResolvedValue(undefined),
			delete: vi.fn().mockResolvedValue(undefined),
		};
		const onCancellation = makeCancellationSpy();
		const tool = makeSuspendingTool(
			'suspend_tool',
			async (_input, ctx) =>
				await ctx.suspend(
					{ reason: 'needs approval' },
					{ continuation: { childRunId: 'child-on-save-error' } },
				),
			onCancellation,
		);
		const runtime = createRuntimeWithCheckpointStore([tool], checkpointStore);
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-1', 'suspend_tool', { value: 'a' }),
		);

		const result = await runtime.generate('go');

		expect(result.error).toBe(persistenceError);
		expect(onCancellation).toHaveBeenCalledOnce();
		expect(checkpointStore.delete).toHaveBeenCalledWith(result.runId);
	});
});

// ---------------------------------------------------------------------------
// Resume checkpoint lifecycle (cleanup owned by the sinks)
// ---------------------------------------------------------------------------

describe('AgentRuntime.resume() — checkpoint lifecycle', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	function makeApprovalTool() {
		return makeSuspendingTool('suspend_tool', async (_input, ctx) => {
			if (ctx.resumeData) return { approved: true };
			return await ctx.suspend({ reason: 'needs approval' });
		});
	}

	it('deletes the checkpoint after a resumed generate run completes', async () => {
		const { runtime } = createRuntimeWithTools([makeApprovalTool()], 1);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];

		// Resume to completion (no re-suspend) — the sink should tear down the checkpoint.
		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		const resumed = await runtime.resume('generate', { approved: true }, { runId, toolCallId });
		expect(resumed.pendingSuspend).toBeUndefined();
		expect(resumed.finishReason).toBe('stop');

		// Checkpoint is gone: resuming the same runId again must fail.
		await expect(
			runtime.resume('generate', { approved: true }, { runId, toolCallId }),
		).rejects.toThrow(`No suspended run found for runId: ${runId}`);
	});

	it('keeps the checkpoint when a resumed generate run re-suspends', async () => {
		const { runtime } = createRuntimeWithTools([makeApprovalTool()], Infinity);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([
				{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} },
				{ toolCallId: 'tc-2', toolName: 'suspend_tool', args: {} },
			]),
		);
		const first = await runtime.generate('run tools');
		const { runId } = first.pendingSuspend![0];

		// Resume only tc-1 — tc-2 stays pending, so the run re-suspends and the
		// checkpoint must survive (finishSuspended must NOT clean up).
		const second = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId: 'tc-1' },
		);
		expect(second.pendingSuspend!.length).toBe(1);

		// Still resumable: tc-2 completes the run without throwing.
		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		const third = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId: 'tc-2' },
		);
		expect(third.pendingSuspend).toBeUndefined();
		expect(third.finishReason).toBe('stop');
	});

	it('deletes the checkpoint after a resumed stream run completes', async () => {
		const { runtime } = createRuntimeWithTools([makeApprovalTool()], 1);

		// Suspend via a generate run (so we can capture runId synchronously), then
		// resume as a stream and drain it.
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];

		streamText.mockReturnValueOnce(makeStreamSuccess('done'));
		const resumed = await runtime.resume('stream', { approved: true }, { runId, toolCallId });
		await collectChunks(resumed.stream);

		// Checkpoint is gone once the streamed run completed.
		await expect(
			runtime.resume('stream', { approved: true }, { runId, toolCallId }),
		).rejects.toThrow(`No suspended run found for runId: ${runId}`);
	});

	it('claims the checkpoint after resume validation passes', async () => {
		const checkpointStore = makeClaimingCheckpointStore();
		const runtime = createRuntimeWithCheckpointStore([makeApprovalTool()], checkpointStore);
		const onResumeClaimed = vi.fn();

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];

		generateText.mockResolvedValueOnce(makeGenerateSuccess('done'));
		const resumed = await runtime.resume(
			'generate',
			{ approved: true },
			{ runId, toolCallId, onResumeClaimed },
		);

		expect(resumed.finishReason).toBe('stop');
		expect(checkpointStore.claimForResume).toHaveBeenCalledTimes(1);
		expect(checkpointStore.claimForResume).toHaveBeenCalledWith(
			runId,
			expect.objectContaining({ status: 'suspended' }),
		);
		expect(onResumeClaimed).toHaveBeenCalledTimes(1);
	});

	it('does not emit a runtime error when another resume wins the checkpoint claim', async () => {
		const checkpointStore = makeClaimingCheckpointStore();
		const eventBus = new AgentEventBus();
		const onResumeClaimed = vi.fn();
		const errorEvents: AgentEventData[] = [];
		eventBus.on(AgentEvent.Error, (event) => errorEvents.push(event));
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			tools: [makeApprovalTool()],
			checkpointStorage: checkpointStore,
			eventBus,
		});

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];
		checkpointStore.claimForResume.mockResolvedValueOnce(false);

		await expect(
			runtime.resume('stream', { approved: true }, { runId, toolCallId, onResumeClaimed }),
		).rejects.toBeInstanceOf(StaleResumeError);
		expect(errorEvents).toEqual([]);
		expect(streamText).not.toHaveBeenCalled();
		expect(onResumeClaimed).not.toHaveBeenCalled();
	});

	it('does not invoke the claim hook when checkpoint claiming fails', async () => {
		const checkpointStore = makeClaimingCheckpointStore();
		const runtime = createRuntimeWithCheckpointStore([makeApprovalTool()], checkpointStore);
		const onResumeClaimed = vi.fn();

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];
		checkpointStore.claimForResume.mockRejectedValueOnce(new Error('checkpoint unavailable'));

		const resumed = await runtime.resume(
			'stream',
			{ approved: true },
			{ runId, toolCallId, onResumeClaimed },
		);
		await collectChunks(resumed.stream);

		expect(onResumeClaimed).not.toHaveBeenCalled();
		expect(streamText).not.toHaveBeenCalled();
	});

	it('does not claim the checkpoint when the resume payload is invalid', async () => {
		const checkpointStore = makeClaimingCheckpointStore();
		const runtime = createRuntimeWithCheckpointStore([makeApprovalTool()], checkpointStore);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId, toolCallId } = first.pendingSuspend![0];

		await expect(
			runtime.resume('generate', { approved: 'yes' }, { runId, toolCallId }),
		).rejects.toThrow('Invalid resume payload');

		expect(checkpointStore.claimForResume).not.toHaveBeenCalled();
	});

	it('does not claim the checkpoint when the tool call id is invalid', async () => {
		const checkpointStore = makeClaimingCheckpointStore();
		const runtime = createRuntimeWithCheckpointStore([makeApprovalTool()], checkpointStore);

		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
		);
		const first = await runtime.generate('run tool');
		const { runId } = first.pendingSuspend![0];

		await expect(
			runtime.resume('generate', { approved: true }, { runId, toolCallId: 'tc-missing' }),
		).rejects.toThrow('No tool call found for toolCallId: tc-missing');

		expect(checkpointStore.claimForResume).not.toHaveBeenCalled();
	});
});

// ---------------------------------------------------------------------------
// Generic reasoning
// ---------------------------------------------------------------------------

describe('reasoning', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('passes reasoning as a top-level generateText option without provider-specific settings', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			reasoning: 'high',
		});

		await runtime.generate('hello', {
			persistence: { resourceId: 'user1', threadId: 'thread1' },
			providerOptions: {
				anthropic: { cacheControl: { type: 'ephemeral' } },
			},
		});

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.reasoning).toBe('high');
		expect((callArgs.providerOptions as Record<string, Record<string, unknown>>).anthropic).toEqual(
			{ cacheControl: { type: 'ephemeral' } },
		);
	});

	it('passes reasoning as a top-level streamText option', async () => {
		streamText.mockReturnValue(makeStreamSuccess());
		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-5.1',
			instructions: 'You are a test assistant.',
			reasoning: 'low',
		});

		const result = await runtime.stream('hello');
		await collectChunks(result.stream);

		const callArgs = streamText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.reasoning).toBe('low');
		expect(callArgs.providerOptions).toBeUndefined();
	});
});

describe('provider-specific thinking', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('maps a fixed Anthropic thinking budget into providerOptions', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			thinking: { budgetTokens: 4096 },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.reasoning).toBeUndefined();
		expect(callArgs.providerOptions).toEqual({
			anthropic: { thinking: { type: 'enabled', budgetTokens: 4096 } },
		});
	});
});

// ---------------------------------------------------------------------------
// Instruction providerOptions
// ---------------------------------------------------------------------------

describe('tool systemInstruction merging', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	function getSystemMessageText(): string {
		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const systemMsg = callArgs.instructions;
		if (Array.isArray(systemMsg)) {
			return systemMsg.map((entry) => String((entry as { content: string }).content)).join('');
		}
		const system = systemMsg as { role: string; content: string };
		expect(system.role).toBe('system');
		return String(system.content);
	}

	it("wraps a tool's systemInstruction in a built_in_rules block above user instructions", async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const toolWithDirective: BuiltTool = {
			name: 'show_card',
			description: 'show a card',
			systemInstruction: 'Prefer this tool over plain text when posting images.',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () => await Promise.resolve('ok'),
		};

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a helpful assistant.',
			tools: [toolWithDirective],
		});

		await runtime.generate('hello');

		const text = getSystemMessageText();
		expect(text).toContain('<built_in_rules>');
		expect(text).toContain('- Prefer this tool over plain text when posting images.');
		expect(text).toContain('</built_in_rules>');
		expect(text).toContain('You are a helpful assistant.');
		expect(text.indexOf('<built_in_rules>')).toBeLessThan(text.indexOf('You are a helpful'));
	});

	it('joins multiple tools systemInstructions into a single block', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const toolA: BuiltTool = {
			name: 'a',
			description: 'a',
			systemInstruction: 'Rule A.',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve('ok'),
		};
		const toolB: BuiltTool = {
			name: 'b',
			description: 'b',
			systemInstruction: 'Rule B.',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve('ok'),
		};
		const toolC: BuiltTool = {
			name: 'c',
			description: 'c',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve('ok'),
		};

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'base',
			tools: [toolA, toolB, toolC],
		});

		await runtime.generate('hello');

		const text = getSystemMessageText();
		const block = text.match(/<built_in_rules>([\s\S]*?)<\/built_in_rules>/);
		expect(block).not.toBeNull();
		expect(block![1]).toContain('- Rule A.');
		expect(block![1]).toContain('- Rule B.');
		expect(block![1]).not.toContain('Rule C');
	});

	it('does not add a built_in_rules block when no tool sets a systemInstruction', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const plainTool: BuiltTool = {
			name: 'plain',
			description: 'plain',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve('ok'),
		};

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a helpful assistant.',
			tools: [plainTool],
		});

		await runtime.generate('hello');

		const text = getSystemMessageText();
		expect(text).not.toContain('<built_in_rules>');
		expect(text).toContain('You are a helpful assistant.');
	});

	it("moves a newly loaded deferred tool's systemInstruction into the uncached system message without changing the cached one", async () => {
		const deferredTool: BuiltTool = {
			name: 'deferred_capability',
			description: 'Deferred capability',
			systemInstruction: 'Always confirm before using deferred_capability.',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () => await Promise.resolve({ ok: true }),
		};

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			deferredTools: [deferredTool],
		});

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{
						toolCallId: 'tc-load',
						toolName: 'load_tool',
						args: { toolName: 'deferred_capability' },
					},
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		await runtime.generate('load the deferred capability');

		const calls = generateText.mock.calls as Array<
			[{ instructions: Array<{ content: string }> | { content: string } }]
		>;
		const beforeLoadSystem = calls[0][0].instructions;
		const afterLoadSystem = calls[1][0].instructions;

		const beforeFirst = Array.isArray(beforeLoadSystem)
			? beforeLoadSystem[0].content
			: beforeLoadSystem.content;
		const afterFirst = Array.isArray(afterLoadSystem)
			? afterLoadSystem[0].content
			: afterLoadSystem.content;
		const afterSecond = Array.isArray(afterLoadSystem) ? afterLoadSystem[1]?.content : undefined;

		// Before loading, the tool's instruction appears nowhere.
		expect(beforeFirst).not.toContain('Always confirm before using deferred_capability.');

		// The cached (first) system message must stay byte-identical after the load.
		expect(afterFirst).toBe(beforeFirst);

		// After loading, the instruction reaches the model via the uncached second message.
		expect(afterSecond).toContain('Always confirm before using deferred_capability.');
		expect(afterFirst).not.toContain('Always confirm before using deferred_capability.');
	});
});

describe('instruction providerOptions', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('should attach providerOptions to system message', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			instructionProviderOptions: {
				anthropic: { cacheControl: { type: 'ephemeral' } },
			},
		});

		await runtime.generate('hello', {
			persistence: { resourceId: 'user1', threadId: 'thread1' },
		});

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const systemMsg = callArgs.instructions as Record<string, unknown>;
		expect(systemMsg.role).toBe('system');
		expect(systemMsg.providerOptions).toEqual({
			anthropic: { cacheControl: { type: 'ephemeral' } },
		});
	});
});

// ---------------------------------------------------------------------------
// promptCaching()
// ---------------------------------------------------------------------------

describe('promptCaching', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('applies default 1h Anthropic cache breakpoints to the system message and the last conversation message', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const systemMsg = callArgs.instructions as Record<string, unknown>;
		expect(systemMsg.providerOptions).toEqual({
			anthropic: { cacheControl: { type: 'ephemeral', ttl: '1h' } },
		});
		const messages = callArgs.messages as Array<Record<string, unknown>>;
		expect(messages[messages.length - 1].providerOptions).toEqual({
			anthropic: { cacheControl: { type: 'ephemeral', ttl: '1h' } },
		});
	});

	it('defaults OpenAI to 24h retention with an auto-generated promptCacheKey', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-5.1',
			instructions: 'You are a test assistant.',
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const openaiOpts = (callArgs.providerOptions as Record<string, Record<string, unknown>>).openai;
		expect(openaiOpts.promptCacheRetention).toBe('24h');
		expect(typeof openaiOpts.promptCacheKey).toBe('string');
	});

	it('keeps generic reasoning separate from prompt caching and caller providerOptions', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-5.1',
			instructions: 'You are a test assistant.',
			reasoning: 'high',
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello', {
			providerOptions: { openai: { strictJsonSchema: false } },
		});

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const openaiOpts = (callArgs.providerOptions as Record<string, Record<string, unknown>>).openai;
		expect(callArgs.reasoning).toBe('high');
		expect(openaiOpts).not.toHaveProperty('reasoningEffort');
		expect(openaiOpts).not.toHaveProperty('reasoningSummary');
		expect(openaiOpts.promptCacheRetention).toBe('24h');
		expect(typeof openaiOpts.promptCacheKey).toBe('string');
		expect(openaiOpts.strictJsonSchema).toBe(false);
	});

	it('adds a cache breakpoint to the last tool definition for a fully-static Anthropic tool set', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const toolA = makeMockTool('tool_a', async () => await Promise.resolve({ ok: true }));
		const toolB = makeMockTool('tool_b', async () => await Promise.resolve({ ok: true }));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			tools: [toolA, toolB],
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const tools = callArgs.tools as Record<string, { providerOptions?: unknown }>;
		// Every tool also carries the Anthropic `eagerInputStreaming: false` quirk default.
		expect(tools.tool_a.providerOptions).toEqual({
			anthropic: { eagerInputStreaming: false },
		});
		expect(tools.tool_b.providerOptions).toEqual({
			anthropic: { eagerInputStreaming: false, cacheControl: { type: 'ephemeral', ttl: '1h' } },
		});
	});

	it('does not add a tool cache breakpoint when deferred tools are configured (Anthropic)', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const coreTool = makeMockTool('core_tool', async () => await Promise.resolve({ ok: true }));
		const deferredTool = makeMockTool(
			'deferred_capability',
			async () => await Promise.resolve({ ok: true }),
		);

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			tools: [coreTool],
			deferredTools: [deferredTool],
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const tools = callArgs.tools as Record<string, { providerOptions?: unknown }>;
		// No cache breakpoint is added, but every tool still carries the
		// Anthropic `eagerInputStreaming: false` quirk default.
		for (const tool of Object.values(tools)) {
			expect(tool.providerOptions).toEqual({ anthropic: { eagerInputStreaming: false } });
		}
	});

	it('adds a tool cache breakpoint on recall_memory for an episodic Anthropic agent (no deferred tools)', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const memory = new InMemoryMemory();
		const fakeEmbedder = { specificationVersion: 'v2' } as never;

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'anthropic/claude-sonnet-4-5',
			instructions: 'You are a test assistant.',
			memory,
			episodicMemory: { embedder: fakeEmbedder },
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});

		// recall_memory is static within a run, so it is eligible to anchor the
		// tool breakpoint (it is the last tool in getCurrentTools).
		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const tools = callArgs.tools as Record<string, { providerOptions?: unknown }>;
		expect(tools).toHaveProperty('recall_memory');
		expect(tools.recall_memory.providerOptions).toEqual({
			anthropic: { eagerInputStreaming: false, cacheControl: { type: 'ephemeral', ttl: '1h' } },
		});
	});

	it('does not add message or tool cache breakpoints for OpenAI', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const toolA = makeMockTool('tool_a', async () => await Promise.resolve({ ok: true }));

		const runtime = new AgentRuntime({
			name: 'test',
			model: 'openai/gpt-5.1',
			instructions: 'You are a test assistant.',
			tools: [toolA],
			promptCaching: { enabled: true },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const messages = callArgs.messages as Array<Record<string, unknown>>;
		expect(messages[messages.length - 1].providerOptions).toBeUndefined();
		const tools = callArgs.tools as Record<string, { providerOptions?: unknown }>;
		// The Anthropic `eagerInputStreaming: false` quirk default is applied to every
		// tool regardless of model; non-Anthropic providers simply ignore the namespace.
		expect(tools.tool_a.providerOptions).toEqual({ anthropic: { eagerInputStreaming: false } });
	});
});

// ---------------------------------------------------------------------------
// Observation log background jobs
// ---------------------------------------------------------------------------

describe('AgentRuntime — observation log jobs', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('schedules observation after a persisted stream turn', async () => {
		streamText.mockReturnValue(makeStreamSuccess('Remembered response'));
		const memory = new InMemoryMemory();
		await memory.saveThread({ id: 'thread-1', resourceId: 'resource-1' });

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 10_000,
				observationLogTailLimit: 20,
				observe: async () =>
					await Promise.resolve('* CRITICAL (14:30) User needs observation memory wired.'),
				reflect: async () => await Promise.resolve('{"drop":[],"merge":[]}'),
			},
		});

		const result = await runtime.stream('Please remember this.', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await collectChunks(result.stream);
		await runtime.dispose();

		const observations = await memory.getActiveObservationLog({
			observationScopeId: 'thread-1',
		});
		expect(observations).toMatchObject([
			{
				marker: 'critical',
				text: 'User needs observation memory wired.',
				parentId: null,
			},
		]);
		const cursor = await memory.getCursor('thread-1');
		expect(typeof cursor?.lastObservedMessageId).toBe('string');
	});

	it('schedules observation after a persisted generate turn', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Remembered response'));
		const memory = new InMemoryMemory();
		await memory.saveThread({ id: 'thread-1', resourceId: 'resource-1' });

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 10_000,
				observationLogTailLimit: 20,
				observe: async () =>
					await Promise.resolve('* IMPORTANT (14:30) User asked for generate observation memory.'),
				reflect: async () => await Promise.resolve('{"drop":[],"merge":[]}'),
			},
		});

		await runtime.generate('Please remember this.', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		const observations = await memory.getActiveObservationLog({
			observationScopeId: 'thread-1',
		});
		expect(observations).toMatchObject([
			{
				marker: 'important',
				text: 'User asked for generate observation memory.',
			},
		]);
	});

	it('indexes episodic memory after observation jobs complete', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Remembered response'));
		embed.mockResolvedValue({ embedding: [1, 0], usage: { tokens: 1 } });
		embedMany.mockResolvedValue({ embeddings: [[1, 0]], usage: { tokens: 1 } });
		const memory = new InMemoryMemory();
		const fakeEmbedder = { specificationVersion: 'v2' } as never;
		const observationLockSpy = vi.spyOn(memory, 'acquireObservationLogTaskLock');
		const episodicLockSpy = vi.spyOn(memory.episodic.taskLock!, 'acquire');

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				observationLogTailLimit: 20,
				observe: async () =>
					await Promise.resolve('* CRITICAL (14:30) User chose Postgres for memory storage.'),
			},
			episodicMemory: {
				embedder: fakeEmbedder,
				extract: async ({ observations }) =>
					await Promise.resolve({
						entries: [
							{
								content: 'User chose Postgres for memory storage.',
								sources: [
									{
										observationId: observations[0].id,
										evidence: 'User chose Postgres',
									},
								],
							},
						],
					}),
			},
		});

		await runtime.generate('Please remember the Postgres decision.', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		const entries = await memory.episodic.searchEntries(
			{ resourceId: 'resource-1' },
			'Postgres storage',
			{ queryEmbedding: [1, 0] },
		);
		expect(entries).toHaveLength(1);
		expect(entries[0].content).toBe('User chose Postgres for memory storage.');
		const cursor = await memory.episodic.getCursor({
			observationScopeId: 'thread-1',
		});
		expect(typeof cursor?.lastIndexedObservationId).toBe('string');
		const firstLockCall = episodicLockSpy.mock.calls.at(0);
		if (!firstLockCall) throw new Error('Expected episodic memory lock acquisition');
		const [lockedResourceId, lockOptions] = firstLockCall;
		expect(lockedResourceId).toBe('resource-1');
		expect(typeof lockOptions.holderId).toBe('string');
		expect(typeof lockOptions.ttlMs).toBe('number');
		const observationLockTaskKinds = observationLockSpy.mock.calls.map((call) => String(call[1]));
		expect(observationLockTaskKinds).not.toContain('episodic-indexer');
	});

	it('skips episodic indexing when the episodic task lock is held', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Plain response'));
		const memory = new InMemoryMemory();
		const observationScope = {
			observationScopeId: 'thread-1',
		};
		const [observation] = await memory.appendObservationLogEntries([
			{
				...observationScope,
				marker: 'critical',
				text: 'User chose Postgres for memory storage.',
				createdAt: new Date('2026-05-20T12:00:00Z'),
			},
		]);
		const extract = vi.fn(async () => {
			await Promise.resolve();

			return {
				entries: [
					{
						content: 'User chose Postgres for memory storage.',
						sources: [{ observationId: observation.id, evidence: 'User chose Postgres' }],
					},
				],
			};
		});
		vi.spyOn(memory.episodic.taskLock!, 'acquire').mockResolvedValue(null);

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			episodicMemory: {
				embedder: { specificationVersion: 'v2' } as never,
				extract,
			},
		});

		await runtime.generate('Please remember this.', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		expect(extract).not.toHaveBeenCalled();
		await expect(memory.episodic.getCursor(observationScope)).resolves.toBeNull();
	});

	it('does not inject episodic memory and exposes recall_memory for explicit recall', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Scoped response'));
		const memory = new InMemoryMemory();
		const fakeEmbedder = { specificationVersion: 'v2' } as never;
		await memory.episodic.saveEntryWithSources(
			{
				resourceId: 'resource-1',
				content: 'Earlier session: user chose Postgres for memory storage.',
				embedding: [1, 0],
			},
			[
				{
					observationId: 'obs-resource-1',
					threadId: 'thread-resource-1',
					evidenceText: 'user chose Postgres',
				},
			],
		);
		await memory.episodic.saveEntryWithSources(
			{
				resourceId: 'resource-2',
				content: 'Earlier session: user chose SQLite for memory storage.',
				embedding: [1, 0],
			},
			[
				{
					observationId: 'obs-resource-2',
					threadId: 'thread-resource-2',
					evidenceText: 'user chose SQLite',
				},
			],
		);

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			episodicMemory: { embedder: fakeEmbedder },
		});

		await runtime.generate('What storage did we choose?', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});

		const callArgs = (generateText.mock.calls[0] as [unknown])[0] as {
			instructions: { content: string };
			tools: Record<string, unknown>;
		};
		const systemPrompt = callArgs.instructions?.content ?? '';
		expect(systemPrompt).not.toContain('<episodic_memory>');
		expect(systemPrompt).not.toContain('Postgres');
		expect(systemPrompt).not.toContain('SQLite');
		expect(callArgs.tools).toHaveProperty('recall_memory');
		expect(embed).not.toHaveBeenCalled();
	});

	it('counts recall_memory query embedding tokens', async () => {
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCall('tc-recall', 'recall_memory', { query: 'storage' }),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Recalled response'));
		embed.mockResolvedValue({ embedding: [1, 0], usage: { tokens: 7 } });
		const counter = makeExecutionCounter();
		const memory = new InMemoryMemory();
		await memory.episodic.saveEntryWithSources(
			{
				resourceId: 'resource-1',
				content: 'Earlier session: user chose Postgres for memory storage.',
				embedding: [1, 0],
			},
			[
				{
					observationId: 'obs-resource-1',
					threadId: 'thread-resource-1',
					evidenceText: 'user chose Postgres',
				},
			],
		);

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			episodicMemory: { embedder: { specificationVersion: 'v2' } as never },
		});

		await runtime.generate('What storage did we choose?', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
			executionCounter: counter,
		});

		expect(counter.incrementTokenCount).toHaveBeenCalledWith(7);
	});

	it('does not schedule observation jobs without policy callbacks', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Plain response'));
		const memory = new InMemoryMemory();
		await memory.saveThread({ id: 'thread-1', resourceId: 'resource-1' });

		const runtime = new AgentRuntime({
			name: 'plain-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 10_000,
				observationLogTailLimit: 20,
			},
		});

		await runtime.generate('Please remember this.', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		expect(
			await memory.getActiveObservationLog({
				observationScopeId: 'thread-1',
			}),
		).toEqual([]);
		expect(await memory.getCursor('thread-1')).toBeNull();
	});

	// TODO: Fix this test it's flaky
	// eslint-disable-next-line n8n-local-rules/no-skipped-tests
	it.skip('keeps history resource-filtered while observation-log memory is thread-local', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Remembered response'));
		const memory = new InMemoryMemory();
		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 10_000,
				observationLogTailLimit: 20,
				observe: async ({ transcript }) =>
					await Promise.resolve(
						transcript.includes('resource-one')
							? '* CRITICAL (14:30) Resource one memory.'
							: '* CRITICAL (14:30) Resource two memory.',
					),
				reflect: async () => await Promise.resolve('{"drop":[],"merge":[]}'),
			},
		});

		await runtime.generate('remember resource-one preference', {
			persistence: { threadId: 'shared-thread', resourceId: 'resource-1' },
		});
		await runtime.dispose();
		await runtime.generate('remember resource-two preference', {
			persistence: { threadId: 'shared-thread', resourceId: 'resource-2' },
		});
		await runtime.dispose();

		generateText.mockClear();
		generateText.mockResolvedValue(makeGenerateSuccess('Scoped response'));

		await runtime.generate('what is my memory?', {
			persistence: { threadId: 'shared-thread', resourceId: 'resource-2' },
		});

		const generateTextMock = generateText as MockedFunction<
			(input: {
				instructions: { content: string } | Array<{ content: string }>;
				messages: Array<{
					role: string;
					content: unknown;
				}>;
			}) => unknown
		>;
		const [{ instructions, messages }] = generateTextMock.mock.calls[0];
		const systemText = Array.isArray(instructions)
			? instructions.map((entry) => entry.content).join('')
			: instructions.content;
		expect(systemText).toContain('Resource one memory.');
		expect(systemText).toContain('Resource two memory.');
		expect(JSON.stringify(messages)).not.toContain('remember resource-one preference');

		await runtime.dispose();
	});

	it('loads only unobserved history when an observation cursor exists', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Scoped response'));
		const memory = new InMemoryMemory();
		await memory.saveThread({ id: 'thread-1', resourceId: 'resource-1' });

		const observedAt = new Date('2026-01-01T00:01:00.000Z');
		const unobservedAt = new Date('2026-01-01T00:02:00.000Z');
		await memory.saveMessages({
			threadId: 'thread-1',
			resourceId: 'resource-1',
			messages: [
				{
					id: 'm1',
					createdAt: new Date('2026-01-01T00:00:00.000Z'),
					role: 'user',
					content: [{ type: 'text', text: 'OBSERVED_OLD_MESSAGE' }],
				},
				{
					id: 'm2',
					createdAt: observedAt,
					role: 'assistant',
					content: [{ type: 'text', text: 'Old response' }],
				},
				{
					id: 'm3',
					createdAt: unobservedAt,
					role: 'user',
					content: [{ type: 'text', text: 'UNOBSERVED_RECENT_MESSAGE' }],
				},
			],
		});
		await memory.setCursor({
			observationScopeId: 'thread-1',
			lastObservedMessageId: 'm2',
			lastObservedAt: observedAt,
			updatedAt: observedAt,
		});
		// A cursor is only trusted when an observation log stands in for the
		// pre-cursor messages, so seed one alongside the cursor.
		await memory.appendObservationLogEntries([
			{
				observationScopeId: 'thread-1',
				marker: 'critical',
				text: 'Observed old context.',
				parentId: null,
				createdAt: observedAt,
			},
		]);

		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				observe: async () => await Promise.resolve('* CRITICAL (14:30) Observed.'),
			},
		});

		await runtime.generate('Current turn question', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});

		const generateTextMock = generateText as MockedFunction<
			(input: { messages: unknown[] }) => unknown
		>;
		const [{ messages }] = generateTextMock.mock.calls[0];
		const serialized = JSON.stringify(messages);
		expect(serialized).not.toContain('OBSERVED_OLD_MESSAGE');
		expect(serialized).toContain('UNOBSERVED_RECENT_MESSAGE');

		await runtime.dispose();
	});

	it('emits an error event when an observer background task fails', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Plain response'));
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const error = new Error('observer failed');
		const errorEvents: AgentEventData[] = [];
		bus.on(AgentEvent.Error, (event) => errorEvents.push(event));
		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 10_000,
				observationLogTailLimit: 20,
				observe: async () => await Promise.reject(error),
				reflect: async () => await Promise.resolve('{"drop":[],"merge":[]}'),
			},
		});

		await runtime.generate('please remember this', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		expect(errorEvents).toEqual(
			expect.arrayContaining([expect.objectContaining({ error, source: 'observer' })]),
		);
	});

	it('emits an error event when a reflector background task fails', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Plain response'));
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const error = new Error('reflector failed');
		const errorEvents: AgentEventData[] = [];
		bus.on(AgentEvent.Error, (event) => errorEvents.push(event));
		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				reflectorThresholdTokens: 1,
				observationLogTailLimit: 20,
				observe: async () =>
					await Promise.resolve(`* CRITICAL (14:30) ${'Large observation. '.repeat(20)}`),
				reflect: async () => await Promise.reject(error),
			},
		});

		await runtime.generate('please remember this', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		expect(errorEvents).toEqual(
			expect.arrayContaining([expect.objectContaining({ error, source: 'reflector' })]),
		);
	});

	it('emits one error event when an episodic indexer background task fails', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess('Plain response'));
		const memory = new InMemoryMemory();
		const bus = new AgentEventBus();
		const error = new Error('episodic extraction failed');
		const errorEvents: AgentEventData[] = [];
		bus.on(AgentEvent.Error, (event) => errorEvents.push(event));
		const runtime = new AgentRuntime({
			name: 'observing-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			memory,
			observationalMemory: {
				observerThresholdTokens: 1,
				observationLogTailLimit: 20,
				observe: async () =>
					await Promise.resolve('* CRITICAL (14:30) User chose Postgres for memory storage.'),
			},
			episodicMemory: {
				embedder: { specificationVersion: 'v2' } as never,
				extract: async () => await Promise.reject(error),
			},
		});

		await runtime.generate('please remember this', {
			persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
		});
		await runtime.dispose();

		expect(errorEvents).toEqual([
			expect.objectContaining({
				error,
				source: 'episodic-memory',
				message: 'Episodic memory indexing task failed',
			}),
		]);
	});
});

// ---------------------------------------------------------------------------
// Mid-run observation (AGENT-191 / AGENT-228)
// ---------------------------------------------------------------------------

describe('AgentRuntime — mid-run observation', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	const OBSERVE_TEXT = '* CRITICAL (14:30) Mid-run observation captured.';
	const PERSISTENCE = { threadId: 'thread-1', resourceId: 'resource-1' };

	type CapturedModelCall = {
		instructions: { content: string } | Array<{ content: string }>;
		messages: Array<{ role: string; content: unknown }>;
	};

	function capturedCall(index: number): CapturedModelCall {
		return generateText.mock.calls[index][0] as CapturedModelCall;
	}

	function flattenInstructions(instructions: CapturedModelCall['instructions']): string {
		return Array.isArray(instructions)
			? instructions.map((entry) => entry.content).join('')
			: instructions.content;
	}

	function makeStepTool(): BuiltTool {
		return {
			name: 'do_step',
			description: 'Perform a step',
			inputSchema: z.object({ step: z.number() }),
			handler: async () => await Promise.resolve({ done: true }),
		};
	}

	function buildMidRunRuntime(
		memory: InMemoryMemory,
		extra?: { tools?: BuiltTool[]; checkpointStorage?: CheckpointStore },
	): AgentRuntime {
		return new AgentRuntime({
			name: 'mid-run-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			memory,
			tools: extra?.tools ?? [makeStepTool()],
			...(extra?.checkpointStorage ? { checkpointStorage: extra.checkpointStorage } : {}),
			observationalMemory: {
				observerThresholdTokens: 1,
				observationLogTailLimit: 20,
				observe: async () => await Promise.resolve(OBSERVE_TEXT),
			},
		});
	}

	it('compacts the live prompt after crossing the budget mid-run', async () => {
		const memory = new InMemoryMemory();
		const runtime = buildMidRunRuntime(memory);
		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-1', 'do_step', { step: 1 }))
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-2', 'do_step', { step: 2 }))
			.mockResolvedValueOnce(makeGenerateSuccess('all done'));

		const result = await runtime.generate('start work', { persistence: PERSISTENCE });
		await runtime.dispose();

		expect(generateText).toHaveBeenCalledTimes(3);
		expect(JSON.stringify(capturedCall(0).messages)).toContain('start work');

		// After the first boundary's compaction the prompt is just the reminder,
		// and the refreshed system prompt carries the observation.
		const second = capturedCall(1);
		expect(second.messages).toEqual([{ role: 'user', content: OBSERVATION_CONTINUATION_REMINDER }]);
		expect(flattenInstructions(second.instructions)).toContain('Mid-run observation captured.');

		// The caller still receives the full response set of the turn.
		expect(result.messages).toHaveLength(3);
		expect(JSON.stringify(result.messages)).toContain('all done');

		const observations = await memory.getActiveObservationLog({
			observationScopeId: 'thread-1',
		});
		expect(observations.length).toBeGreaterThanOrEqual(1);
		expect(await memory.getCursor('thread-1')).not.toBeNull();
	});

	it('re-derives the mask from the cursor when resuming a suspended run', async () => {
		const memory = new InMemoryMemory();
		const checkpointStore = makeClaimingCheckpointStore();
		const tools = [makeStepTool(), makeInterruptibleTool()];

		const runtime = buildMidRunRuntime(memory, { tools, checkpointStorage: checkpointStore });
		generateText
			// Iteration 1 crosses the budget and compacts before iteration 2.
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc-step', 'do_step', { step: 1 }))
			.mockResolvedValueOnce(
				makeGenerateWithToolCall('tc-approve', 'approve', { question: 'go on?' }),
			);

		const first = await runtime.generate('start work', { persistence: PERSISTENCE });
		const suspension = first.pendingSuspend?.[0];
		if (!suspension) throw new Error('Expected the run to suspend on the approve tool');

		const resumed = buildMidRunRuntime(memory, { tools, checkpointStorage: checkpointStore });
		generateText.mockResolvedValueOnce(makeGenerateSuccess('continued'));
		await resumed.resume(
			'generate',
			{ approved: true },
			{ runId: suspension.runId, toolCallId: suspension.toolCallId },
		);
		await resumed.dispose();
		await runtime.dispose();

		const resumeCall = capturedCall(2);
		expect(JSON.stringify(resumeCall.messages)).not.toContain('start work');
		expect(resumeCall.messages[0]).toEqual({
			role: 'user',
			content: OBSERVATION_CONTINUATION_REMINDER,
		});
	});

	it('observes at the resume boundary when the resumed tool results cross the budget', async () => {
		const memory = new InMemoryMemory();
		const checkpointStore = makeClaimingCheckpointStore();
		const tools = [makeStepTool(), makeInterruptibleTool()];

		const runtime = buildMidRunRuntime(memory, { tools, checkpointStorage: checkpointStore });
		// Suspends on the very first iteration, so no boundary check has run and
		// no cursor exists when the run parks.
		generateText.mockResolvedValueOnce(
			makeGenerateWithToolCall('tc-approve', 'approve', { question: 'go on?' }),
		);
		const first = await runtime.generate('start work', { persistence: PERSISTENCE });
		const suspension = first.pendingSuspend?.[0];
		if (!suspension) throw new Error('Expected the run to suspend on the approve tool');
		expect(await memory.getCursor('thread-1')).toBeNull();

		const resumed = buildMidRunRuntime(memory, { tools, checkpointStorage: checkpointStore });
		generateText.mockResolvedValueOnce(makeGenerateSuccess('continued'));
		await resumed.resume(
			'generate',
			{ approved: true },
			{ runId: suspension.runId, toolCallId: suspension.toolCallId },
		);
		await resumed.dispose();
		await runtime.dispose();

		// The first post-resume model call already sees the compacted window:
		// the resumed tool results were observed before reaching the model.
		const resumeCall = capturedCall(1);
		expect(resumeCall.messages).toEqual([
			{ role: 'user', content: OBSERVATION_CONTINUATION_REMINDER },
		]);
		expect(flattenInstructions(resumeCall.instructions)).toContain('Mid-run observation captured.');
		expect(await memory.getCursor('thread-1')).not.toBeNull();
	});
});

// ---------------------------------------------------------------------------
// Runtime telemetry — generate / stream / tool execution
// ---------------------------------------------------------------------------

describe('AgentRuntime — telemetry propagation', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	const baseTelemetry: BuiltTelemetry = {
		enabled: true,
		functionId: 'test-agent',
		metadata: { env: 'test' },
		recordInputs: true,
		recordOutputs: false,
		integrations: [],
		tracer: { startSpan: vi.fn() },
	};

	it('passes telemetry config into generateText as telemetry', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'telemetry-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry: baseTelemetry,
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const telemetry = callArgs.telemetry as Record<string, unknown>;
		expect(telemetry).toBeDefined();
		expect(telemetry.isEnabled).toBe(true);
		expect(telemetry.functionId).toBe('test-agent');
		expect(telemetry.recordInputs).toBe(true);
		expect(telemetry.recordOutputs).toBe(false);
	});

	it('uses updated telemetry config for later runs', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const updatedTelemetry: BuiltTelemetry = {
			...baseTelemetry,
			functionId: 'updated-agent',
			metadata: { env: 'updated' },
		};

		const runtime = new AgentRuntime({
			name: 'telemetry-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry: baseTelemetry,
		});

		runtime.setTelemetry(updatedTelemetry);
		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const telemetry = callArgs.telemetry as Record<string, unknown>;
		expect(telemetry.functionId).toBe('updated-agent');
	});

	it('wraps generate calls in a telemetry root span when the tracer supports active spans', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const span = {
			end: vi.fn(),
			recordException: vi.fn(),
			setStatus: vi.fn(),
		};
		const tracer = {
			startActiveSpan: vi.fn(async (_name: string, _options: unknown, fn: unknown) => {
				if (typeof fn !== 'function') {
					throw new Error('Expected span callback');
				}
				const spanFn = fn as (spanValue: typeof span) => Promise<unknown>;
				return await spanFn(span);
			}),
		};
		const telemetry: BuiltTelemetry = { ...baseTelemetry, tracer };

		const runtime = new AgentRuntime({
			name: 'telemetry-root-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry,
		});

		await runtime.generate('hello');

		expect(tracer.startActiveSpan).toHaveBeenCalledWith(
			'test-agent.generate',
			{
				root: true,
				// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
				attributes: expect.objectContaining<Record<string, string>>({
					'gen_ai.operation.name': 'invoke_agent',
					'gen_ai.agent.name': 'telemetry-root-test',
				}),
			},
			expect.any(Function),
		);
		// A plain (non-LangSmith) tracer must not get langsmith.* attributes —
		// they'd be noise on a generic OTLP backend.
		const [, options] = tracer.startActiveSpan.mock.calls[0];
		const attributes = (options as { attributes: Record<string, unknown> }).attributes;
		expect(Object.keys(attributes).some((key) => key.startsWith('langsmith.'))).toBe(false);
		expect(span.end).toHaveBeenCalledTimes(1);
	});

	it('can suppress the generic runtime root span while keeping native telemetry enabled', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const tracer = {
			startActiveSpan: vi.fn(),
		};
		const telemetry: BuiltTelemetry = {
			...baseTelemetry,
			runtimeRootSpanEnabled: false,
			integrations: [{ onStart: vi.fn() }],
			tracer,
		};

		const runtime = new AgentRuntime({
			name: 'telemetry-root-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry,
		});

		await runtime.generate('hello');

		expect(tracer.startActiveSpan).not.toHaveBeenCalled();

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.telemetry).toEqual(
			expect.objectContaining({
				isEnabled: true,
				functionId: 'test-agent',
				integrations: telemetry.integrations,
			}),
		);
	});

	it('adds a LangSmith tool catalog to telemetry root spans', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());
		const span = {
			end: vi.fn(),
			recordException: vi.fn(),
			setStatus: vi.fn(),
		};
		const tracer = {
			startActiveSpan: vi.fn(async (_name: string, _options: unknown, fn: unknown) => {
				if (typeof fn !== 'function') {
					throw new Error('Expected span callback');
				}
				const spanFn = fn as (spanValue: typeof span) => Promise<unknown>;
				return await spanFn(span);
			}),
		};
		const telemetry: BuiltTelemetry = {
			...baseTelemetry,
			metadata: {
				...baseTelemetry.metadata,
				langsmith_trace_id: 'trace-1',
				langsmith_actor_run_id: 'actor-run-1',
			},
			isLangSmith: true,
			tracer,
		};
		const tool = new ToolBuilder('lookup')
			.description('Lookup records')
			.input(z.object({ query: z.string() }))
			.handler(async () => await Promise.resolve({ ok: true }))
			.build();

		const runtime = new AgentRuntime({
			name: 'telemetry-root-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			tools: [tool],
			telemetry,
		});

		await runtime.generate('hello');

		const rootSpanOptions = tracer.startActiveSpan.mock.calls[0][1] as {
			attributes: Record<string, unknown>;
		};
		const { attributes } = rootSpanOptions;
		expect(attributes).toEqual(
			expect.objectContaining({
				'langsmith.metadata.available_tools': ['lookup'],
			}),
		);
		expect(attributes).not.toHaveProperty('langsmith.trace.id');
		expect(attributes).not.toHaveProperty('langsmith.span.parent_id');
		expect(attributes['gen_ai.prompt']).toEqual(expect.stringContaining('"name":"lookup"'));
		expect(attributes['gen_ai.prompt']).toEqual(
			expect.stringContaining('"description":"Lookup records"'),
		);
		expect(attributes['gen_ai.prompt']).toEqual(expect.stringContaining('"input_schema"'));
	});

	it('passes telemetry config into streamText as telemetry', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const runtime = new AgentRuntime({
			name: 'telemetry-stream-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry: baseTelemetry,
		});

		const { stream } = await runtime.stream('hello');
		await collectChunks(stream);

		const callArgs = streamText.mock.calls[0][0] as Record<string, unknown>;
		const telemetry = callArgs.telemetry as Record<string, unknown>;
		expect(telemetry).toBeDefined();
		expect(telemetry.isEnabled).toBe(true);
		expect(telemetry.functionId).toBe('test-agent');
	});

	// Real parent/child span nesting (via a genuine OTel context manager, not
	// mock-call-order heuristics) is covered by
	// agent-runtime-memory.otel.test.ts's "memory span nesting under the root
	// span" suite, for both generate() and stream().

	it('enables smoothStream by default on streamText', async () => {
		streamText.mockReturnValue(makeStreamSuccess());
		const smoothStreamSpy = vi.spyOn(aiModule, 'smoothStream');

		const runtime = new AgentRuntime({
			name: 'smooth-stream-default-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
		});

		const { stream } = await runtime.stream('hello');
		await collectChunks(stream);

		const callArgs = streamText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.experimental_transform).toEqual(expect.any(Function));
		expect(smoothStreamSpy).toHaveBeenCalledWith({});

		smoothStreamSpy.mockRestore();
	});

	it('omits smoothStream when explicitly disabled', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const runtime = new AgentRuntime({
			name: 'smooth-stream-disabled-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
		});

		const { stream } = await runtime.stream('hello', { smoothStream: false });
		await collectChunks(stream);

		const callArgs = streamText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.experimental_transform).toBeUndefined();
	});

	it('forwards non-default smoothStream options to the AI SDK', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const smoothStreamSpy = vi.spyOn(aiModule, 'smoothStream');

		const runtime = new AgentRuntime({
			name: 'smooth-stream-options-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
		});

		const smoothStreamOptions = { delayInMs: 25, chunking: 'line' as const };
		const { stream } = await runtime.stream('hello', { smoothStream: smoothStreamOptions });

		await collectChunks(stream);

		expect(smoothStreamSpy).toHaveBeenCalledWith(smoothStreamOptions);

		smoothStreamSpy.mockRestore();
	});

	it('inherits telemetry from ExecutionOptions when no own telemetry is set', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'child-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			// No telemetry set on the runtime itself
		});

		await runtime.generate('hello', {
			persistence: { resourceId: 'r1', threadId: 't1' },
			telemetry: baseTelemetry,
		});

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		const telemetry = callArgs.telemetry as Record<string, unknown>;
		expect(telemetry).toBeDefined();
		expect(telemetry.isEnabled).toBe(true);
		// Inherited telemetry uses the child agent's name as functionId
		expect(telemetry.functionId).toBe('child-agent');
	});

	it('passes resolved telemetry to tool handlers via parentTelemetry', async () => {
		let capturedTelemetry: BuiltTelemetry | undefined;
		let capturedToolCallId: string | undefined;

		const spyTool: BuiltTool = new ToolBuilder('spy')
			.description('captures telemetry from context')
			.input(z.object({ x: z.string() }))
			.output(z.object({ ok: z.boolean() }))
			.handler(async (_input, ctx) => {
				capturedTelemetry = ctx.parentTelemetry;
				capturedToolCallId = ctx.toolCallId;
				return await Promise.resolve({ ok: true });
			})
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc1', 'spy', { x: 'test' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'tool-telemetry-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			tools: [spyTool],
			telemetry: baseTelemetry,
		});

		await runtime.generate('test');

		expect(capturedTelemetry).toBe(baseTelemetry);
		expect(capturedToolCallId).toBe('tc1');
	});

	it('emits AI SDK-compatible telemetry spans for local tool execution', async () => {
		const spans: Array<{
			name: string;
			span: {
				end: Mock;
				recordException: Mock;
				setAttributes: Mock;
				setStatus: Mock;
			};
		}> = [];
		const tracer = {
			startActiveSpan: vi.fn(async (name: string, _options: unknown, fn: unknown) => {
				if (typeof fn !== 'function') {
					throw new Error('Expected span callback');
				}
				const span = {
					end: vi.fn(),
					recordException: vi.fn(),
					setAttributes: vi.fn(),
					setStatus: vi.fn(),
				};
				spans.push({ name, span });
				const spanFn = fn as (spanValue: typeof span) => Promise<unknown>;
				return await spanFn(span);
			}),
		};
		const telemetry: BuiltTelemetry = {
			...baseTelemetry,
			recordOutputs: true,
			tracer,
		};
		const spyTool: BuiltTool = new ToolBuilder('spy')
			.description('captures telemetry from context')
			.input(z.object({ x: z.string() }))
			.output(z.object({ ok: z.boolean() }))
			.handler(async () => await Promise.resolve({ ok: true }))
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc1', 'spy', { x: 'test' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'tool-telemetry-test',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			tools: [spyTool],
			telemetry,
		});

		await runtime.generate('test');

		const toolCallSpan = tracer.startActiveSpan.mock.calls.find(
			([name]) => name === 'execute_tool spy',
		);
		expect(toolCallSpan).toBeDefined();
		expect(toolCallSpan?.[1]).toEqual({
			// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
			attributes: expect.objectContaining<Record<string, string>>({
				'operation.name': 'ai.toolCall test-agent',
				'resource.name': 'test-agent',
				'ai.operationId': 'ai.toolCall',
				'ai.telemetry.functionId': 'test-agent',
				'ai.telemetry.metadata.env': 'test',
				'gen_ai.operation.name': 'execute_tool',
				'gen_ai.tool.name': 'spy',
				'gen_ai.tool.call.id': 'tc1',
				'gen_ai.agent.name': 'tool-telemetry-test',
				'gen_ai.tool.call.arguments': '{"x":"test"}',
				'ai.toolCall.name': 'spy',
				'ai.toolCall.id': 'tc1',
				'ai.toolCall.args': '{"x":"test"}',
			}),
		});
		const toolSpan = spans.find((span) => span.name === 'execute_tool spy')?.span;
		expect(toolSpan?.setAttributes).toHaveBeenCalledWith({
			'ai.toolCall.result': '{"ok":true}',
			'gen_ai.tool.call.result': '{"ok":true}',
		});
		expect(toolSpan?.end).toHaveBeenCalledTimes(1);
	});

	it('passes inherited telemetry to tool handlers for sub-agent scenarios', async () => {
		let capturedTelemetry: BuiltTelemetry | undefined;

		const spyTool: BuiltTool = new ToolBuilder('spy')
			.description('captures telemetry from context')
			.input(z.object({ x: z.string() }))
			.output(z.object({ ok: z.boolean() }))
			.handler(async (_input, ctx) => {
				capturedTelemetry = ctx.parentTelemetry;
				return await Promise.resolve({ ok: true });
			})
			.build();

		generateText
			.mockResolvedValueOnce(makeGenerateWithToolCall('tc1', 'spy', { x: 'test' }))
			.mockResolvedValueOnce(makeGenerateSuccess('done'));

		const runtime = new AgentRuntime({
			name: 'sub-agent',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			tools: [spyTool],
			// No own telemetry
		});

		await runtime.generate('test', {
			persistence: { resourceId: 'r1', threadId: 't1' },
			telemetry: baseTelemetry,
		});

		// Should receive inherited telemetry with the child agent's functionId
		expect(capturedTelemetry).toBeDefined();
		expect(capturedTelemetry!.functionId).toBe('sub-agent');
		expect(capturedTelemetry!.tracer).toBe(baseTelemetry.tracer);
	});

	it('does not include telemetry when telemetry is disabled', async () => {
		generateText.mockResolvedValue(makeGenerateSuccess());

		const runtime = new AgentRuntime({
			name: 'disabled-telemetry',
			model: 'openai/gpt-4o-mini',
			instructions: 'test',
			eventBus: new AgentEventBus(),
			telemetry: { ...baseTelemetry, enabled: false },
		});

		await runtime.generate('hello');

		const callArgs = generateText.mock.calls[0][0] as Record<string, unknown>;
		expect(callArgs.telemetry).toBeUndefined();
	});
});

// Cancellation (Feature 1: cancel suspended tool via user message)
// ---------------------------------------------------------------------------

describe('AgentRuntime.resume() with createCancellation() — auto-bypass', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	/** A tool that suspends on first call and returns on resume. */
	function makeSuspendToolForCancel(): BuiltTool {
		return {
			name: 'interactive_tool',
			description: 'A tool that suspends',
			inputSchema: z.object({ prompt: z.string() }),
			suspendSchema: z.object({ prompt: z.string() }),
			resumeSchema: z.object({ answer: z.string() }),
			handler: async (_input: unknown, ctx: unknown) => {
				const { suspend, resumeData } = ctx as InterruptibleToolContext;
				if (!resumeData) {
					return await suspend({ prompt: 'What should I do?' });
				}
				return { result: (resumeData as { answer: string }).answer };
			},
		};
	}

	it('auto-bypass: does NOT call the tool handler on cancellation', async () => {
		const onCancellation = makeCancellationSpy();
		const handlerSpy = vi.fn().mockImplementation(async (_input: unknown, ctx: unknown) => {
			const { suspend, resumeData } = ctx as InterruptibleToolContext;
			if (!resumeData) {
				return await suspend({ prompt: 'What should I do?' });
			}
			return { result: (resumeData as { answer: string }).answer };
		});
		const tool: BuiltTool = {
			name: 'interactive_tool',
			description: 'A tool that suspends',
			inputSchema: z.object({ prompt: z.string() }),
			suspendSchema: z.object({ prompt: z.string() }),
			resumeSchema: z.object({ answer: z.string() }),
			handler: handlerSpy,
			onCancellation,
		};

		const { runtime } = createRuntimeWithTools([tool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'interactive_tool', args: { prompt: 'continue?' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done after cancel'));

		const first = await runtime.generate('start', {});
		const { runId, toolCallId } = first.pendingSuspend![0];

		// Reset call count to check the handler is NOT called on resume
		handlerSpy.mockClear();

		const resumed = await runtime.resume(
			'generate',
			createCancellation('do something else instead'),
			{ runId, toolCallId },
		);

		// Handler should NOT have been called for the resume
		expect(handlerSpy).not.toHaveBeenCalled();
		expect(onCancellation).toHaveBeenCalledWith(
			{ prompt: 'continue?' },
			expect.objectContaining({
				cancellation: { message: 'do something else instead' },
				runId,
				toolCallId,
				suspendPayload: { prompt: 'What should I do?' },
			}),
		);
		// The generation should have continued after cancellation
		expect(resumed.finishReason).toBe('stop');
	});

	it('auto-bypass: injects the steering message and the LLM sees it', async () => {
		const tool = makeSuspendToolForCancel();
		const { runtime } = createRuntimeWithTools([tool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'interactive_tool', args: { prompt: 'continue?' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Understood, doing something else'));

		const first = await runtime.generate('start');
		const { runId, toolCallId } = first.pendingSuspend![0];

		await runtime.resume('generate', createCancellation('pivot to plan B'), { runId, toolCallId });

		// The second generateText call should include the user steering message
		const secondCallMessages = (
			generateText.mock.calls[1][0] as { messages: Array<{ role: string; content: unknown }> }
		).messages;
		const userMessages = secondCallMessages.filter((m) => m.role === 'user');
		const steeringMsg = userMessages.find((m) =>
			JSON.stringify(m.content).includes('pivot to plan B'),
		);
		expect(steeringMsg).toBeDefined();
	});

	it('stream: emits a cancellation tool result then continues', async () => {
		const tool = makeSuspendToolForCancel();
		const { runtime } = createRuntimeWithTools([tool], 1);

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'interactive_tool',
									args: { prompt: 'continue?' },
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-1', toolName: 'interactive_tool', input: { prompt: 'continue?' } },
				]),
			})
			.mockReturnValueOnce(makeStreamSuccess('Redirected'));

		const firstResult = await runtime.stream('start');
		const firstChunks = await collectChunks(firstResult.stream);

		const suspendChunk = firstChunks.find((c) => c.type === 'tool-call-suspended');
		expect(suspendChunk).toBeDefined();
		const { runId, toolCallId } = suspendChunk as Extract<
			StreamChunk,
			{ type: 'tool-call-suspended' }
		>;

		const resumed = await runtime.resume('stream', createCancellation('go another direction'), {
			runId,
			toolCallId,
		});
		const resumedChunks = await collectChunks(resumed.stream);

		const cancellationResult = resumedChunks.find(
			(c) => c.type === 'tool-result' && c.toolCallId === 'tc-1',
		);
		expect(cancellationResult).toEqual(
			expect.objectContaining({
				type: 'tool-result',
				toolCallId: 'tc-1',
				toolName: 'interactive_tool',
				output: '[Tool call cancelled. User said: "go another direction"]',
				canceled: true,
			}),
		);

		// Generation should continue after the cancellation result
		const textChunks = resumedChunks.filter((c) => c.type === 'text-delta');
		expect(textChunks.length).toBeGreaterThan(0);
	});

	it('rejects with an error if the checkpoint is not found', async () => {
		const { runtime } = createRuntimeWithTools([makeSuspendToolForCancel()], 1);
		await expect(
			runtime.resume('generate', createCancellation('no checkpoint'), {
				runId: 'nonexistent',
				toolCallId: 'tc-1',
			}),
		).rejects.toThrow('No suspended run found for runId: nonexistent');
	});
});

describe('AgentRuntime.resume() with createCancellation() — manual handling (handleCancellation)', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	it('calls the tool handler with ctx.cancellation set', async () => {
		const handlerSpy = vi.fn().mockImplementation(async (_input: unknown, ctx: unknown) => {
			const { suspend, resumeData, cancellation } = ctx as InterruptibleToolContext;
			if (cancellation) {
				// Manual cleanup path — return a note for the LLM
				return `Cancelled: ${cancellation.message}`;
			}
			if (!resumeData) {
				return await suspend({ prompt: 'Confirm?' });
			}
			return 'done';
		});
		const tool: BuiltTool = {
			name: 'manual_cancel_tool',
			description: 'A tool with manual cancellation',
			inputSchema: z.object({ value: z.string() }),
			suspendSchema: z.object({ prompt: z.string() }),
			resumeSchema: z.object({ confirmed: z.boolean() }),
			handleCancellation: true,
			handler: handlerSpy,
		};

		const { runtime } = createRuntimeWithTools([tool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-1', toolName: 'manual_cancel_tool', args: { value: 'x' } },
				]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Done after manual cancel'));

		const first = await runtime.generate('test');
		const { runId, toolCallId } = first.pendingSuspend![0];

		handlerSpy.mockClear();

		await runtime.resume('generate', createCancellation('user said stop'), { runId, toolCallId });

		// Handler SHOULD have been called for the resume
		expect(handlerSpy).toHaveBeenCalledTimes(1);
		const callCtx = handlerSpy.mock.calls[0][1] as InterruptibleToolContext;
		expect(callCtx.cancellation).toEqual({ message: 'user said stop' });
		expect(callCtx.resumeData).toBeUndefined();
	});
});

// ---------------------------------------------------------------------------
// toModelOutput error resilience — processToolCall must never re-throw
// ---------------------------------------------------------------------------

describe('AgentRuntime — toModelOutput error resilience', () => {
	beforeEach(() => {
		vi.clearAllMocks();
	});

	function makeTransformErrorTool(name = 'transform_tool'): BuiltTool {
		return {
			name,
			description: 'A tool whose toModelOutput always throws',
			inputSchema: z.object({ value: z.string().optional() }),
			handler: async () => await Promise.resolve({ raw: 'data' }),
			toModelOutput: () => {
				throw new Error('toModelOutput failed');
			},
		};
	}

	function makeSuspendingTransformErrorTool(name = 'suspend_tool'): BuiltTool {
		return {
			name,
			description: 'A suspending tool whose toModelOutput throws on the resumed call',
			inputSchema: z.object({ value: z.string().optional() }),
			suspendSchema: z.object({ reason: z.string() }),
			resumeSchema: z.object({ approved: z.boolean() }),
			handler: async (_input, ctx) => {
				const suspendCtx = ctx as InterruptibleToolContext;
				if (!suspendCtx.resumeData) {
					return await suspendCtx.suspend({ reason: 'needs approval' });
				}
				return { done: true };
			},
			toModelOutput: () => {
				throw new Error('toModelOutput failed on resume');
			},
		};
	}

	it('generate(): toModelOutput throwing is treated as a tool error — loop continues', async () => {
		const tool = makeTransformErrorTool();
		const { runtime } = createRuntimeWithTools([tool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'transform_tool', args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('I see the tool errored'));

		const result = await runtime.generate('run the tool');

		// The agent loop must not crash — it continues and the LLM finishes
		expect(result.finishReason).toBe('stop');
		// Two LLM calls: tool-call turn + follow-up after error
		expect(generateText).toHaveBeenCalledTimes(2);
	});

	it('stream(): toModelOutput throwing surfaces as isError tool-result — loop continues', async () => {
		const tool = makeTransformErrorTool();
		const { runtime } = createRuntimeWithTools([tool], 1);

		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-1',
									toolName: 'transform_tool',
									args: {},
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([{ toolCallId: 'tc-1', toolName: 'transform_tool', input: {} }]),
			})
			.mockReturnValueOnce(makeStreamSuccess('I see the tool errored'));

		const { stream } = await runtime.stream('run the tool');
		const chunks = await collectChunks(stream);

		const finishChunk = chunks.find((c) => c.type === 'finish') as
			| (StreamChunk & { type: 'finish' })
			| undefined;
		expect(finishChunk?.finishReason).toBe('stop');

		const toolResultChunk = chunks.find(
			(c) => c.type === 'tool-result' && c.toolCallId === 'tc-1',
		) as (StreamChunk & { type: 'tool-result' }) | undefined;
		expect(toolResultChunk?.isError).toBe(true);
	});

	it('generate() resume: toModelOutput throwing in resumed tool call is captured — loop continues', async () => {
		const tool = makeSuspendingTransformErrorTool();
		const { runtime } = createRuntimeWithTools([tool], 1);

		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-1', toolName: 'suspend_tool', args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('Handled the resume error'));

		// First call: tool suspends
		const firstResult = await runtime.generate('run the tool');
		expect(firstResult.finishReason).toBe('tool-calls');
		expect(firstResult.pendingSuspend).toHaveLength(1);

		const { runId, toolCallId } = firstResult.pendingSuspend![0];

		// Resume: tool returns a result but toModelOutput throws
		// Bug: without fix this propagates out of iteratePendingToolCallsConcurrent and
		// causes generate() to return with finishReason 'error' instead of 'stop'.
		const resumeResult = await runtime.resume(
			'generate',
			{ approved: true },
			{
				runId,
				toolCallId,
			},
		);

		expect(resumeResult.finishReason).toBe('stop');
		// Two LLM calls: initial + after resume error
		expect(generateText).toHaveBeenCalledTimes(2);
	});

	it('stream() resume: toModelOutput throwing in resumed tool call does not close the stream with error', async () => {
		const tool = makeSuspendingTransformErrorTool();
		const { runtime } = createRuntimeWithTools([tool], 1);

		// First stream: agent calls the tool and suspends
		streamText.mockReturnValueOnce({
			stream: makeChunkStream([{ type: 'text-delta', id: 'text-1', text: 'thinking...' }]),
			finishReason: Promise.resolve('tool-calls'),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
			response: Promise.resolve({
				messages: [
					{
						role: 'assistant',
						content: [
							{
								type: 'tool-call',
								toolCallId: 'tc-1',
								toolName: 'suspend_tool',
								args: {},
							},
						],
					},
				],
			}),
			toolCalls: Promise.resolve([{ toolCallId: 'tc-1', toolName: 'suspend_tool', input: {} }]),
		});

		const firstResult = await runtime.stream('run the tool');
		const firstChunks = await collectChunks(firstResult.stream);

		const suspendChunk = firstChunks.find((c) => c.type === 'tool-call-suspended') as
			| (StreamChunk & { type: 'tool-call-suspended' })
			| undefined;
		expect(suspendChunk).toBeDefined();

		// Resume: toModelOutput throws
		// Bug: without fix this closes the stream with finishReason 'error' via closeStreamWithError.
		streamText.mockReturnValueOnce(makeStreamSuccess('Handled the resume error'));

		const resumed = await runtime.resume(
			'stream',
			{ approved: true },
			{
				runId: suspendChunk!.runId,
				toolCallId: suspendChunk!.toolCallId,
			},
		);
		const resumedChunks = await collectChunks(resumed.stream);

		const finishChunk = resumedChunks.find((c) => c.type === 'finish') as
			| (StreamChunk & { type: 'finish' })
			| undefined;
		expect(finishChunk).toBeDefined();
		// Must NOT be 'error' — the tool error should be contained, not kill the stream
		expect(finishChunk!.finishReason).not.toBe('error');

		// The resumed stream should contain a tool-result with isError: true
		const toolResultChunk = resumedChunks.find(
			(c) => c.type === 'tool-result' && c.toolCallId === 'tc-1',
		) as (StreamChunk & { type: 'tool-result' }) | undefined;
		expect(toolResultChunk).toBeDefined();
		expect(toolResultChunk!.isError).toBe(true);
	});
});

// ---------------------------------------------------------------------------
// empty model responses (e.g. provider safety blocks)
// ---------------------------------------------------------------------------

describe('AgentRuntime — empty model responses', () => {
	beforeEach(() => {
		generateText.mockReset();
		streamText.mockReset();
	});

	function makeEmptyStream(finishReason: string, extraChunks: Array<Record<string, unknown>> = []) {
		return {
			stream: makeChunkStream(extraChunks),
			finishReason: Promise.resolve(finishReason),
			usage: Promise.resolve({ inputTokens: 10, outputTokens: 0, totalTokens: 10 }),
			response: Promise.resolve({ messages: [] }),
			toolCalls: Promise.resolve([]),
		};
	}

	it.each(['other', 'unknown', 'content-filter'])(
		'stream: yields an error chunk when the model returns no output with finish reason "%s"',
		async (finishReason) => {
			streamText.mockReturnValue(makeEmptyStream(finishReason));

			const { runtime } = createRuntime();
			const { stream: readableStream } = await runtime.stream('hello');
			const chunks = await collectChunks(readableStream);

			const errorChunk = chunks.find((c) => c.type === 'error') as
				| (StreamChunk & { type: 'error' })
				| undefined;
			expect(errorChunk).toBeDefined();
			expect(String((errorChunk!.error as Error).message)).toContain('no output');
			expect(String((errorChunk!.error as Error).message)).toContain(finishReason);

			const finishChunk = chunks.find((c) => c.type === 'finish') as
				| (StreamChunk & { type: 'finish' })
				| undefined;
			expect(finishChunk).toBeDefined();
			expect(finishChunk!.finishReason).toBe('error');
		},
	);

	it('stream: includes the Google prompt block reason captured from raw chunks', async () => {
		streamText.mockReturnValue(
			makeEmptyStream('other', [
				{
					type: 'raw',
					rawValue: { promptFeedback: { blockReason: 'PROHIBITED_CONTENT' } },
				},
			]),
		);

		const { runtime } = createRuntime(undefined, 'google/gemini-2.5-flash');
		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		const errorChunk = chunks.find((c) => c.type === 'error') as
			| (StreamChunk & { type: 'error' })
			| undefined;
		expect(errorChunk).toBeDefined();
		expect(String((errorChunk!.error as Error).message)).toContain('PROHIBITED_CONTENT');
	});

	it('stream: requests raw chunks for google models so block reasons are observable', async () => {
		streamText.mockReturnValue(makeStreamSuccess());

		const { runtime } = createRuntime(undefined, 'google/gemini-2.5-flash');
		const { stream: readableStream } = await runtime.stream('hello');
		await collectChunks(readableStream);

		const args = streamText.mock.calls[0][0] as Record<string, unknown>;
		expect(args).toHaveProperty('include.rawChunks', true);
	});

	it('stream: an empty response finishing with "stop" is not treated as an error', async () => {
		streamText.mockReturnValue(makeEmptyStream('stop'));

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		expect(chunks.find((c) => c.type === 'error')).toBeUndefined();
		const finishChunk = chunks.find((c) => c.type === 'finish') as
			| (StreamChunk & { type: 'finish' })
			| undefined;
		expect(finishChunk!.finishReason).not.toBe('error');
	});

	it('generate: returns finishReason "error" with a descriptive error when the model returns no output', async () => {
		generateText.mockResolvedValue({
			finishReason: 'other',
			usage: { inputTokens: 10, outputTokens: 0, totalTokens: 10 },
			response: { messages: [] },
			toolCalls: [],
		});

		const { runtime } = createRuntime();
		const result = await runtime.generate('hello');

		expect(result.finishReason).toBe('error');
		expect(String((result.error as Error).message)).toContain('no output');
	});
});

describe('AgentRuntime — MCP connection failure warnings', () => {
	it('emits a warning chunk per recorded MCP failure before the LLM loop, then completes', async () => {
		streamText.mockReturnValue(makeStreamSuccess('Hello'));

		const bus = new AgentEventBus();
		const runtime = new AgentRuntime({
			name: 'mcp-warn',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			mcpConnectionFailures: [
				{ server: 'dead', error: 'fetch failed' },
				{ server: 'also_dead', error: 'boom' },
			],
		});

		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		const warnings = chunks.filter((c) => c.type === 'warning') as Array<
			StreamChunk & { type: 'warning'; message: string; server?: string; source?: string }
		>;
		expect(warnings).toHaveLength(2);
		expect(warnings[0]).toMatchObject({
			type: 'warning',
			message: 'fetch failed',
			source: 'mcp',
			server: 'dead',
			code: 'mcp_connection_failed',
		});
		expect(warnings[1]).toMatchObject({ type: 'warning', server: 'also_dead', message: 'boom' });

		// The run still completes normally — warnings are non-fatal.
		const finishChunk = chunks.find((c) => c.type === 'finish') as
			| (StreamChunk & { type: 'finish'; finishReason: string })
			| undefined;
		expect(finishChunk).toBeDefined();
		expect(finishChunk!.finishReason).not.toBe('error');
	});

	it('emits no warning chunks when there are no MCP connection failures', async () => {
		streamText.mockReturnValue(makeStreamSuccess('Hello'));

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		const chunks = await collectChunks(readableStream);

		expect(chunks.filter((c) => c.type === 'warning')).toEqual([]);
	});

	it('injects a model-facing mcp-connection-status note into the system message when failures are set', async () => {
		streamText.mockReturnValue(makeStreamSuccess('Hello'));

		const bus = new AgentEventBus();
		const runtime = new AgentRuntime({
			name: 'mcp-note',
			model: 'openai/gpt-4o-mini',
			instructions: 'You are a test assistant.',
			eventBus: bus,
			mcpConnectionFailures: [{ server: 'dead', error: 'fetch failed' }],
		});

		const { stream: readableStream } = await runtime.stream('hello');
		await collectChunks(readableStream);

		const callArgs = streamText.mock.calls.at(-1)![0] as Record<string, unknown>;
		const system = callArgs.instructions;
		const systemText = Array.isArray(system)
			? system.map((e) => String((e as { content: string }).content)).join('')
			: String((system as { content: string }).content);

		expect(systemText).toContain('<mcp-connection-status>');
		expect(systemText).toContain('dead');
		expect(systemText).toContain('fetch failed');
		expect(systemText).toMatch(/If this affects the user's request/i);
	});

	it('omits the mcp-connection-status note from the system message when there are no failures', async () => {
		streamText.mockReturnValue(makeStreamSuccess('Hello'));

		const { runtime } = createRuntime();
		const { stream: readableStream } = await runtime.stream('hello');
		await collectChunks(readableStream);

		const callArgs = streamText.mock.calls.at(-1)![0] as Record<string, unknown>;
		const system = callArgs.instructions;
		const systemText = Array.isArray(system)
			? system.map((e) => String((e as { content: string }).content)).join('')
			: String((system as { content: string }).content);

		expect(systemText).not.toContain('<mcp-connection-status>');
	});
});

describe('AgentRuntime — oversized tool results', () => {
	type TruncationEnvelope = {
		_truncated: true;
		originalCharCount: number;
		estimatedTokenCount: number;
		head: string;
		tail: string;
	};

	beforeEach(() => {
		vi.clearAllMocks();
	});

	function parseEnvelope(value: string): TruncationEnvelope {
		try {
			return JSON.parse(value) as TruncationEnvelope;
		} catch {
			throw new Error('Expected a valid truncation envelope');
		}
	}

	async function expectWithinTokenLimit(value: unknown): Promise<void> {
		const { getEncoding } = await import('@n8n/ai-utilities/tokenizer');
		const encoder = await getEncoding('o200k_base');
		expect(encoder.encode(JSON.stringify(value)).length).toBeLessThanOrEqual(
			MAX_MODEL_TOOL_RESULT_TOKENS,
		);
	}

	function getModelToolResultOutput(callIndex = 1): { type: string; value: unknown } | undefined {
		const call = generateText.mock.calls[callIndex][0] as {
			messages: Array<{
				role: string;
				content: Array<{ type: string; output?: { type: string; value: unknown } }>;
			}>;
		};
		const toolMessage = call.messages.find((message) => message.role === 'tool');
		return toolMessage?.content.find((part) => part.type === 'tool-result')?.output;
	}

	function getModelToolResult(callIndex = 1): unknown {
		return getModelToolResultOutput(callIndex)?.value;
	}

	it('leaves a small transformed result unchanged', async () => {
		const tool: BuiltTool = {
			name: 'small_result',
			description: 'Return a small result',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve({ raw: 'value' }),
			toModelOutput: () => ({ summary: 'small' }),
		};
		const { runtime } = createRuntimeWithTools([tool], 1);
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-small', toolName: tool.name, args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess());

		await runtime.generate('run');

		expect(getModelToolResult()).toEqual({ summary: 'small' });
	});

	it('preserves oversized content media without offloading it', async () => {
		const rawOutput = { ok: true };
		const modelOutput = {
			type: 'content' as const,
			value: [
				{
					type: 'file-data' as const,
					data: 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'.repeat(8_000),
					mediaType: 'application/pdf',
				},
			],
		};
		const tool: BuiltTool = {
			name: 'file_result',
			description: 'Return a file',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve(rawOutput),
			toModelOutput: () => modelOutput,
		};
		const { runtime } = createRuntimeWithTools([tool], 1);
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-file', toolName: tool.name, args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess());

		const result = await runtime.generate('run');

		expect(getModelToolResultOutput()).toEqual(modelOutput);
		expect(result.toolCalls?.[0]?.output).toEqual(rawOutput);

		const { runtime: streamRuntime } = createRuntimeWithTools([tool], 1);
		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-file-stream',
									toolName: tool.name,
									args: {},
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-file-stream', toolName: tool.name, input: {} },
				]),
			})
			.mockReturnValueOnce(makeStreamSuccess());

		const { stream } = await streamRuntime.stream('run');
		const chunks = await collectChunks(stream);
		const streamResult = chunks.find(
			(chunk) => chunk.type === 'tool-result' && chunk.toolCallId === 'tc-file-stream',
		) as (StreamChunk & { type: 'tool-result' }) | undefined;

		expect(streamResult?.output).toEqual(modelOutput);
	});

	it('bounds an oversized transformed result while preserving raw output', async () => {
		const rawOutput = {
			value: `RAW_HEAD${'r '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}RAW_TAIL`,
		};
		const transformedOutput = `TRANSFORM_HEAD${'t '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}TRANSFORM_TAIL`;
		const tool: BuiltTool = {
			name: 'large_result',
			description: 'Return a large result',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve(rawOutput),
			toModelOutput: () => transformedOutput,
		};
		const { runtime } = createRuntimeWithTools([tool], 1);
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-large', toolName: tool.name, args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess());

		const result = await runtime.generate('run');
		const envelope = getModelToolResult() as TruncationEnvelope;

		await expectWithinTokenLimit(envelope);
		expect(envelope).toMatchObject({
			_truncated: true,
			originalCharCount: JSON.stringify(transformedOutput).length,
		});
		expect(envelope.head).toContain('TRANSFORM_HEAD');
		expect(envelope.tail).toContain('TRANSFORM_TAIL');
		expect(result.toolCalls?.[0]?.output).toEqual(rawOutput);

		const { runtime: streamRuntime } = createRuntimeWithTools([tool], 1);
		streamText
			.mockReturnValueOnce({
				stream: makeChunkStream([]),
				finishReason: Promise.resolve('tool-calls'),
				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
				response: Promise.resolve({
					messages: [
						{
							role: 'assistant',
							content: [
								{
									type: 'tool-call',
									toolCallId: 'tc-large-stream',
									toolName: tool.name,
									args: {},
								},
							],
						},
					],
				}),
				toolCalls: Promise.resolve([
					{ toolCallId: 'tc-large-stream', toolName: tool.name, input: {} },
				]),
			})
			.mockReturnValueOnce(makeStreamSuccess());

		const { stream } = await streamRuntime.stream('run');
		const chunks = await collectChunks(stream);
		const streamResult = chunks.find(
			(chunk) => chunk.type === 'tool-result' && chunk.toolCallId === 'tc-large-stream',
		) as (StreamChunk & { type: 'tool-result' }) | undefined;

		expect(streamResult?.output).toEqual(envelope);
	});

	it('bounds an oversized error without changing its rejected state', async () => {
		const tool: BuiltTool = {
			name: 'large_error',
			description: 'Throw a large error',
			inputSchema: z.object({}),
			handler: () => {
				throw new Error(
					`ERROR_HEAD${'e '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}ERROR_TAIL`,
				);
			},
		};
		const { runtime } = createRuntimeWithTools([tool], 1);
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-error', toolName: tool.name, args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess('recovered'));

		const result = await runtime.generate('run');
		const toolCall = result.messages
			.flatMap((message) => ('content' in message ? message.content : []))
			.find(
				(content): content is ContentToolCall =>
					content.type === 'tool-call' && content.toolCallId === 'tc-error',
			);
		const envelope = parseEnvelope(toolCall?.state === 'rejected' ? toolCall.error : '');

		expect(result.finishReason).toBe('stop');
		expect(toolCall?.state).toBe('rejected');
		await expectWithinTokenLimit(envelope);
		expect(envelope.head).toContain('ERROR_HEAD');
		expect(envelope.tail).toContain('ERROR_TAIL');
	});

	it('bounds aggregate toMessage text while preserving file content', async () => {
		const fileData = Buffer.from('file').toString('base64');
		const textBlockTokenCount = Math.floor(MAX_MODEL_TOOL_RESULT_TOKENS * 0.6);
		const tool: BuiltTool = {
			name: 'large_message',
			description: 'Return a large message',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve({ ok: true }),
			toMessage: () => ({
				role: 'assistant',
				content: [
					{ type: 'text', text: `MESSAGE_HEAD${'m '.repeat(textBlockTokenCount)}` },
					{ type: 'file', mediaType: 'text/plain', data: fileData },
					{ type: 'text', text: `${'n '.repeat(textBlockTokenCount)}MESSAGE_TAIL` },
				],
			}),
		};
		const { runtime } = createRuntimeWithTools([tool], 1);
		generateText
			.mockResolvedValueOnce(
				makeGenerateWithToolCalls([{ toolCallId: 'tc-message', toolName: tool.name, args: {} }]),
			)
			.mockResolvedValueOnce(makeGenerateSuccess());

		const result = await runtime.generate('run');
		const message = result.messages.find(
			(candidate) =>
				'content' in candidate &&
				candidate.content.some((content) => content.type === 'file' && content.data === fileData),
		);
		const textBlocks =
			message && 'content' in message
				? message.content.filter((content) => content.type === 'text')
				: [];
		const fileBlock =
			message && 'content' in message
				? message.content.find((content) => content.type === 'file')
				: undefined;
		const envelope = parseEnvelope(textBlocks[0]?.text ?? '');

		expect(textBlocks).toHaveLength(1);
		await expectWithinTokenLimit(envelope);
		expect(envelope.head).toContain('MESSAGE_HEAD');
		expect(envelope.tail).toContain('MESSAGE_TAIL');
		expect(fileBlock).toMatchObject({ type: 'file', mediaType: 'text/plain', data: fileData });
		expect(message).toMatchObject({ origin: { kind: 'tool', toolName: 'large_message' } });
	});

	describe('Workspace-backed oversized tool results', () => {
		type OffloadedEnvelope = {
			_offloaded: true;
			path: string;
			originalCharCount: number;
			estimatedTokenCount: number;
			requiredAction: {
				toolName: 'workspace_read_tool_result';
				input: { path: string; view: 'describe' };
			};
		};

		function parseOffloadedEnvelope(value: string | undefined): OffloadedEnvelope {
			try {
				return JSON.parse(value ?? '') as OffloadedEnvelope;
			} catch {
				throw new Error('Expected a valid offloaded result envelope');
			}
		}

		type ModelMessages = Array<{
			role: string;
			content: Array<{
				type: string;
				text?: string;
				output?: { type: string; value: unknown };
			}>;
		}>;

		function toolResultsFromModelMessages(messages: ModelMessages): unknown[] {
			return messages
				.filter((message) => message.role === 'tool')
				.flatMap((message) => message.content)
				.filter((part) => part.type === 'tool-result')
				.map((part) => part.output?.value);
		}

		function modelToolResults(): unknown[] {
			const call = generateText.mock.calls[1][0] as { messages: ModelMessages };
			return toolResultsFromModelMessages(call.messages);
		}

		function createWorkspaceAgent(filesystem: InMemoryFilesystem, tools: BuiltTool[]) {
			return new Agent('workspace-result-test')
				.model('openai/gpt-4o-mini')
				.instructions('Test')
				.tool(tools)
				.toolCallConcurrency(Infinity)
				.workspace(new Workspace({ filesystem }));
		}

		function createRunScopedRuntime(
			filesystem: InMemoryFilesystem,
			tools: BuiltTool[],
			runId: string,
			eventBus?: AgentEventBus,
			checkpointStorage: 'memory' | CheckpointStore = 'memory',
		): AgentRuntime {
			return new AgentRuntime({
				name: 'workspace-result-test',
				model: 'openai/gpt-4o-mini',
				instructions: 'Test',
				tools,
				runId,
				checkpointStorage,
				workspaceFilesystem: filesystem,
				toolCallConcurrency: Infinity,
				...(eventBus ? { eventBus } : {}),
			});
		}

		async function seedUnaffectedWorkspaceFiles(filesystem: InMemoryFilesystem): Promise<void> {
			await filesystem.writeFile(`${getToolResultRunDirectory('concurrent-run')}/keep.json`, '{}', {
				recursive: true,
			});
		}

		async function expectOnlyCurrentRunRemoved(
			filesystem: InMemoryFilesystem,
			runId: string,
		): Promise<void> {
			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(false);
			await expect(filesystem.exists(getToolResultRunDirectory('concurrent-run'))).resolves.toBe(
				true,
			);
		}

		const largeResultOutput = {
			value: 'x '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000),
		};
		const largeResultTool: BuiltTool = {
			name: 'large_result',
			description: 'large result',
			inputSchema: z.object({}),
			handler: async () => await Promise.resolve(largeResultOutput),
		};

		it('offloads oversized content text while preserving media parts', async () => {
			const filesystem = new InMemoryFilesystem();
			const textParts = [
				{
					type: 'text' as const,
					text: `HEAD${'a '.repeat(Math.floor(MAX_MODEL_TOOL_RESULT_TOKENS * 0.6))}`,
				},
				{
					type: 'text' as const,
					text: `${'b '.repeat(Math.floor(MAX_MODEL_TOOL_RESULT_TOKENS * 0.6))}TAIL`,
				},
			];
			const filePart = {
				type: 'file-data' as const,
				data: 'base64-pdf',
				mediaType: 'application/pdf',
			};
			const tool: BuiltTool = {
				name: 'mixed_content',
				description: 'Return text and a file',
				inputSchema: z.object({}),
				handler: async () => await Promise.resolve({ ok: true }),
				toModelOutput: () => ({
					type: 'content',
					value: [textParts[0], filePart, textParts[1]],
				}),
			};
			const agent = createWorkspaceAgent(filesystem, [tool]);
			let modelContent: Array<{ type: string; text?: string }> = [];
			let storedResult: string | undefined;
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([{ toolCallId: 'tc-mixed', toolName: tool.name, args: {} }]),
				)
				.mockImplementationOnce(async ({ messages }: { messages: ModelMessages }) => {
					modelContent = toolResultsFromModelMessages(messages)[0] as typeof modelContent;
					const envelope = parseOffloadedEnvelope(
						modelContent.find((part) => part.type === 'text')?.text,
					);
					storedResult = String(await filesystem.readFile(envelope.path, { encoding: 'utf8' }));
					return await Promise.resolve(makeGenerateSuccess());
				});

			await agent.generate('run');

			expect(modelContent).toEqual([{ type: 'text', text: expect.any(String) }, filePart]);
			expect(storedResult).toBe(JSON.stringify(textParts));
		});

		it('stores complete concurrent transformed results without changing raw outputs', async () => {
			const filesystem = new InMemoryFilesystem();
			const transformedOutputs = {
				first: { value: `FIRST${'a '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}` },
				second: { value: `SECOND${'b '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}` },
			};
			const tools = (Object.keys(transformedOutputs) as Array<keyof typeof transformedOutputs>).map(
				(name): BuiltTool => ({
					name,
					description: name,
					inputSchema: z.object({}),
					handler: async () => await Promise.resolve({ raw: name }),
					toModelOutput: () => transformedOutputs[name],
				}),
			);
			const agent = createWorkspaceAgent(filesystem, tools);
			const storedResultsSeenByNextModel: string[] = [];
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([
						{ toolCallId: 'tc/first', toolName: 'first', args: {} },
						{ toolCallId: 'tc/second', toolName: 'second', args: {} },
					]),
				)
				.mockImplementationOnce(async ({ messages }: { messages: ModelMessages }) => {
					for (const envelope of toolResultsFromModelMessages(messages) as OffloadedEnvelope[]) {
						storedResultsSeenByNextModel.push(
							String(await filesystem.readFile(envelope.path, { encoding: 'utf8' })),
						);
					}
					return await Promise.resolve(makeGenerateSuccess());
				});

			const result = await agent.generate('run', {
				persistence: { threadId: 'slack:C123:123.456', resourceId: 'user-1' },
			});
			const envelopes = modelToolResults() as OffloadedEnvelope[];

			expect(envelopes).toHaveLength(2);
			expect(new Set(envelopes.map(({ path }) => path)).size).toBe(2);
			expect(envelopes.every(({ _offloaded }) => _offloaded)).toBe(true);
			expect(envelopes.map(({ requiredAction }) => requiredAction)).toEqual(
				envelopes.map(({ path }) => ({
					toolName: 'workspace_read_tool_result',
					input: { path, view: 'describe' },
				})),
			);
			expect(storedResultsSeenByNextModel).toEqual([
				JSON.stringify(transformedOutputs.first),
				JSON.stringify(transformedOutputs.second),
			]);
			expect(result.toolCalls?.map(({ output }) => output)).toEqual([
				{ raw: 'first' },
				{ raw: 'second' },
			]);
			await expect(filesystem.exists(getToolResultRunDirectory(result.runId))).resolves.toBe(false);
		});

		it('retains current run results when checkpoint deletion fails', async () => {
			const filesystem = new InMemoryFilesystem();
			const checkpointStorage: CheckpointStore = {
				save: vi.fn(async () => await Promise.resolve()),
				load: vi.fn().mockResolvedValue(undefined),
				delete: vi.fn().mockRejectedValue(new Error('checkpoint delete failed')),
			};
			const runId = 'checkpoint-delete-failed-run';
			const runtime = createRunScopedRuntime(
				filesystem,
				[largeResultTool],
				runId,
				undefined,
				checkpointStorage,
			);
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([
						{ toolCallId: 'tc-large', toolName: largeResultTool.name, args: {} },
					]),
				)
				.mockResolvedValueOnce(makeGenerateSuccess());

			await runtime.generate('run');

			expect(checkpointStorage.delete).toHaveBeenCalledWith(runId);
			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(true);
		});

		it('removes only the current run results after a terminal model error', async () => {
			const filesystem = new InMemoryFilesystem();
			const runId = 'failed-run';
			await seedUnaffectedWorkspaceFiles(filesystem);
			const runtime = createRunScopedRuntime(filesystem, [largeResultTool], runId);
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([
						{ toolCallId: 'tc-large', toolName: largeResultTool.name, args: {} },
					]),
				)
				.mockRejectedValueOnce(new Error('model unavailable'));

			await runtime.generate('run');

			await expectOnlyCurrentRunRemoved(filesystem, runId);
		});

		it('removes only the current run results after cancellation', async () => {
			const filesystem = new InMemoryFilesystem();
			const eventBus = new AgentEventBus();
			const runId = 'cancelled-run';
			eventBus.on(AgentEvent.ToolExecutionEnd, () => eventBus.abort());
			const runtime = createRunScopedRuntime(filesystem, [largeResultTool], runId, eventBus);
			generateText.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-large', toolName: largeResultTool.name, args: {} },
				]),
			);

			await runtime.generate('run');

			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(false);
		});

		it('completes terminal runs when result cleanup hangs', async () => {
			const filesystem = new InMemoryFilesystem();
			const runId = 'hung-cleanup-run';
			vi.spyOn(filesystem, 'rmdir').mockReturnValue(new Promise(() => undefined));
			const runtime = createRunScopedRuntime(filesystem, [largeResultTool], runId);
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([
						{ toolCallId: 'tc-large', toolName: largeResultTool.name, args: {} },
					]),
				)
				.mockResolvedValueOnce(makeGenerateSuccess());

			await expect(runtime.generate('run')).resolves.toMatchObject({ finishReason: 'stop' });
		});

		it('performs no filesystem calls during cleanup when the run never offloaded', async () => {
			const filesystem = new InMemoryFilesystem();
			const existsSpy = vi.spyOn(filesystem, 'exists');
			const rmdirSpy = vi.spyOn(filesystem, 'rmdir');
			const smallResultTool: BuiltTool = {
				name: 'small_result',
				description: 'small result',
				inputSchema: z.object({}),
				handler: async () => await Promise.resolve({ ok: true }),
			};
			const runtime = createRunScopedRuntime(filesystem, [smallResultTool], 'no-offload-run');
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([
						{ toolCallId: 'tc-small', toolName: smallResultTool.name, args: {} },
					]),
				)
				.mockResolvedValueOnce(makeGenerateSuccess());

			await runtime.generate('run');

			expect(existsSpy).not.toHaveBeenCalled();
			expect(rmdirSpy).not.toHaveBeenCalled();
		});

		it('retains offloaded results across suspension and re-suspension, then removes them', async () => {
			const filesystem = new InMemoryFilesystem();
			const checkpointStorage = makeClaimingCheckpointStore();
			const approvalTool = makeSuspendingTool('approval', async (_input, ctx) => {
				if (ctx.resumeData) return { approved: true };
				return await ctx.suspend({ reason: 'approve' });
			});
			const runId = 'suspended-run';
			const runtime = createRunScopedRuntime(
				filesystem,
				[largeResultTool, approvalTool],
				runId,
				undefined,
				checkpointStorage,
			);
			generateText.mockResolvedValueOnce(
				makeGenerateWithToolCalls([
					{ toolCallId: 'tc-large', toolName: largeResultTool.name, args: {} },
					{ toolCallId: 'tc-approval-1', toolName: approvalTool.name, args: {} },
					{ toolCallId: 'tc-approval-2', toolName: approvalTool.name, args: {} },
				]),
			);

			await runtime.generate('run', {
				persistence: { threadId: 'thread-1', resourceId: 'resource-1' },
			});

			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(true);

			const second = await runtime.resume(
				'generate',
				{ approved: true },
				{ runId, toolCallId: 'tc-approval-1' },
			);
			expect(second.pendingSuspend?.map(({ toolCallId }) => toolCallId)).toEqual(['tc-approval-2']);
			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(true);

			let resultReadByResumedModel: string | undefined;
			generateText.mockImplementationOnce(async ({ messages }: { messages: ModelMessages }) => {
				const [resumedEnvelope] = toolResultsFromModelMessages(messages) as OffloadedEnvelope[];
				resultReadByResumedModel = String(
					await filesystem.readFile(resumedEnvelope.path, { encoding: 'utf8' }),
				);
				return await Promise.resolve(makeGenerateSuccess('done'));
			});
			await runtime.resume('generate', { approved: true }, { runId, toolCallId: 'tc-approval-2' });

			expect(resultReadByResumedModel).toBe(JSON.stringify(largeResultOutput));
			await expect(filesystem.exists(getToolResultRunDirectory(runId))).resolves.toBe(false);
		});

		it('stores oversized errors without changing the rejected tool-call state', async () => {
			const filesystem = new InMemoryFilesystem();
			const tool: BuiltTool = {
				name: 'large_error',
				description: 'large error',
				inputSchema: z.object({}),
				handler: () => {
					throw new Error(`ERROR_HEAD${'e '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}`);
				},
			};
			const agent = createWorkspaceAgent(filesystem, [tool]);
			let storedError: string | undefined;
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([{ toolCallId: 'tc-error', toolName: tool.name, args: {} }]),
				)
				.mockImplementationOnce(async ({ messages }: { messages: ModelMessages }) => {
					const [errorEnvelope] = toolResultsFromModelMessages(messages);
					const envelope = parseOffloadedEnvelope(
						typeof errorEnvelope === 'string' ? errorEnvelope : undefined,
					);
					storedError = String(await filesystem.readFile(envelope.path, { encoding: 'utf8' }));
					return await Promise.resolve(makeGenerateSuccess());
				});

			const result = await agent.generate('run', {
				persistence: { threadId: 'thread-1', resourceId: 'user-1' },
			});
			const toolCall = result.messages
				.flatMap((message) => ('content' in message ? message.content : []))
				.find(
					(content): content is ContentToolCall =>
						content.type === 'tool-call' && content.toolCallId === 'tc-error',
				);

			expect(toolCall?.state).toBe('rejected');
			expect(storedError).toContain('ERROR_HEAD');
		});

		it('stores oversized custom-message text while preserving file content', async () => {
			const filesystem = new InMemoryFilesystem();
			const fileData = Buffer.from('file').toString('base64');
			const messageText = `MESSAGE${'m '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}`;
			const tool: BuiltTool = {
				name: 'large_message',
				description: 'large message',
				inputSchema: z.object({}),
				handler: async () => await Promise.resolve({ ok: true }),
				toMessage: async () =>
					await Promise.resolve({
						role: 'assistant',
						content: [
							{ type: 'text', text: messageText },
							{ type: 'file', mediaType: 'text/plain', data: fileData },
						],
					}),
			};
			const agent = createWorkspaceAgent(filesystem, [tool]);
			let storedMessage: string | undefined;
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([{ toolCallId: 'tc-message', toolName: tool.name, args: {} }]),
				)
				.mockImplementationOnce(async ({ messages }: { messages: ModelMessages }) => {
					const envelopeText = messages
						.filter((message) => message.role === 'assistant')
						.flatMap((message) => message.content)
						.find((part) => part.type === 'text')?.text;
					const envelope = parseOffloadedEnvelope(envelopeText);
					storedMessage = String(await filesystem.readFile(envelope.path, { encoding: 'utf8' }));
					return await Promise.resolve(makeGenerateSuccess());
				});

			const result = await agent.generate('run', {
				persistence: { threadId: 'thread-1', resourceId: 'user-1' },
			});
			const message = result.messages.find(
				(candidate) =>
					'content' in candidate &&
					candidate.content.some((content) => content.type === 'file' && content.data === fileData),
			);

			expect(message && 'content' in message ? message.content : []).toContainEqual({
				type: 'file',
				mediaType: 'text/plain',
				data: fileData,
			});
			expect(storedMessage).toBe(JSON.stringify([{ type: 'text', text: messageText }]));
		});

		it('falls back to bounded truncation when storing the result fails', async () => {
			const filesystem = new InMemoryFilesystem();
			vi.spyOn(filesystem, 'writeFile').mockRejectedValue(new Error('storage unavailable'));
			const tool: BuiltTool = {
				name: 'large_result',
				description: 'large result',
				inputSchema: z.object({}),
				handler: async () =>
					await Promise.resolve({
						value: `HEAD${'x '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}TAIL`,
					}),
			};
			const agent = createWorkspaceAgent(filesystem, [tool]);
			generateText
				.mockResolvedValueOnce(
					makeGenerateWithToolCalls([{ toolCallId: 'tc-large', toolName: tool.name, args: {} }]),
				)
				.mockResolvedValueOnce(makeGenerateSuccess());

			await agent.generate('run');

			expect(modelToolResults()[0]).toMatchObject({ _truncated: true });
		});
	});
});

// ---------------------------------------------------------------------------
// model stream stall handling
// ---------------------------------------------------------------------------

describe('AgentRuntime — model stream stall handling', () => {
	beforeEach(() => {
		streamText.mockReset();
	});

	/**
	 * streamText response that emits `chunks` then goes silent — a dead
	 * connection. Always leads with the SDK's synthetic `start` lifecycle chunk,
	 * which arrives before any provider byte and must not count as content.
	 */
	function makeStalledStream(chunks: Array<Record<string, unknown>> = []) {
		chunks = [{ type: 'start' }, ...chunks];
		return {
			stream: (async function* () {
				for (const c of chunks) yield await Promise.resolve(c);
				await new Promise(() => {});
			})(),
			finishReason: new Promise(() => {}),
			usage: new Promise(() => {}),
			response: new Promise(() => {}),
			toolCalls: new Promise(() => {}),
		};
	}

	it('silently retries a turn that stalls before any content and completes', async () => {
		streamText
			.mockReturnValueOnce(makeStalledStream())
			.mockReturnValueOnce(makeStreamSuccess('Recovered'));
		const { runtime } = createRuntime();

		const result = await runtime.stream('hi', {
			modelStreamIdleTimeoutMs: 50,
			modelStreamFirstOutputTimeoutMs: 50,
		});
		const chunks = await collectChunks(result.stream);

		expect(streamText).toHaveBeenCalledTimes(2);
		const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
			| (StreamChunk & { type: 'finish'; finishReason: string })
			| undefined;
		expect(finish?.finishReason).toBe('stop');
		const text = chunks
			.filter(
				(c): c is StreamChunk & { type: 'text-delta'; delta: string } => c.type === 'text-delta',
			)
			.map((c) => c.delta)
			.join('');
		expect(text).toBe('Recovered');
		expect(runtime.getState().status).toBe('success');
	});

	it('fails without retry when the stalled turn already streamed content', async () => {
		streamText
			.mockReturnValueOnce(
				makeStalledStream([{ type: 'text-delta', id: 'text-1', text: 'partial' }]),
			)
			// A wrongly-taken retry completes loudly instead of hanging the test.
			.mockReturnValue(makeStreamSuccess('should-not-run'));
		const { runtime } = createRuntime();

		const result = await runtime.stream('hi', {
			modelStreamIdleTimeoutMs: 50,
			modelStreamFirstOutputTimeoutMs: 50,
		});
		const chunks = await collectChunks(result.stream);

		expect(streamText).toHaveBeenCalledTimes(1);
		const errorChunk = chunks.find((c) => c.type === 'error') as
			| (StreamChunk & { type: 'error'; error: unknown })
			| undefined;
		expect(String(errorChunk?.error)).toContain('stalled');
		expect(runtime.getState().status).toBe('failed');
	});

	it('surfaces the stall error when the retry stalls too', async () => {
		streamText.mockReturnValueOnce(makeStalledStream()).mockReturnValueOnce(makeStalledStream());
		const { runtime } = createRuntime();

		const result = await runtime.stream('hi', {
			modelStreamIdleTimeoutMs: 50,
			modelStreamFirstOutputTimeoutMs: 50,
		});
		const chunks = await collectChunks(result.stream);

		expect(streamText).toHaveBeenCalledTimes(2);
		const errorChunk = chunks.find((c) => c.type === 'error') as
			| (StreamChunk & { type: 'error'; error: unknown })
			| undefined;
		expect(String(errorChunk?.error)).toContain('stalled');
		expect(runtime.getState().status).toBe('failed');
	});
});
