experimental-a2 0.20.0 → 0.22.0

package/dist/ai-server.js

@@ -1,11 +1,12 @@
-import { i as A2Error } from "./session-clear-DGI4jTQ5.js";
-import { a as reducer, s as validateSync } from "./reducer-BAXc8JyF.js";
-import { h as installAmbientToolScopeStorage, m as failedPreviousAttempt, p as consumeSchedulerSendFailure } from "./internal-CscSY7K_.js";
+import { i as A2Error } from "./session-clear-wHd5GCpZ.js";
+import { a as reducer, s as validateSync } from "./reducer-CjmvZSAb.js";
+import { _ as installAmbientToolScopeStorage, g as failedPreviousAttempt, h as drainSettledHooks, m as consumeSchedulerSendFailure } from "./internal-CvsrprnU.js";
 import { t as boundedCache } from "./bounded-cache-CPHWF-8h.js";
-import { a as CONSOLE_TELEMETRY, o as setServerFetchHooks, t as createServer } from "./server-qyb5aJJG.js";
+import { a as CONSOLE_TELEMETRY, n as createServer, o as setServerFetchHooks } from "./server-DtFrWxi3.js";
 import { a as unexplainedAdmission, i as steersOn, n as initialControlState, r as steeringEnd, t as STEERED_REASON } from "./ai-control-CvLbW10B.js";
-import { A as interruptUIMessage, C as reduceToolActivity, D as flattenProgressBatches, E as upsertResponse, M as assertMessageBase, N as restoreMessageBase, O as progressBatches, S as projectedMessage, T as upsertMessage, _ as applyControlChanges, a as events, b as findMessage, d as appSlice, f as isAgentContract, g as isUIMessage, h as aiEventId, j as pauseUIMessage, k as progressOwnedBy, m as AI_EVENT_ID_PREFIX, n as agentContract, p as withAIStorage, u as appReducerOf, v as controlCommit, w as retainMessageBase, y as advanceProjection } from "./ai-mMiWul1T.js";
-import { a as gatewaySessionAffinity, c as harnessLog, f as registerHarnessStates, l as harnessStates, o as harnessDriver } from "./ai-harness-Clbg5ei7.js";
+import { A as interruptUIMessage, C as reduceToolActivity, D as flattenProgressBatches, E as upsertResponse, M as assertMessageBase, N as restoreMessageBase, O as progressBatches, S as projectedMessage, T as upsertMessage, _ as applyControlChanges, b as findMessage, d as appSlice, f as isAgentContract, g as isUIMessage, h as aiEventId, j as pauseUIMessage, k as progressOwnedBy, m as AI_EVENT_ID_PREFIX, n as agentContract, o as events, p as withAIStorage, u as appReducerOf, v as controlCommit, w as retainMessageBase, y as advanceProjection } from "./ai-EEfKyc-z.js";
+import { a as gatewaySessionAffinity, c as harnessLog, f as registerHarnessStates, l as harnessStates, o as harnessDriver } from "./ai-harness-dQeJaR1t.js";
+import { t as runtimeConsumer } from "./runtime-consumer-CIeL33wj.js";
 import { AsyncLocalStorage } from "node:async_hooks";
 import { asSchema, convertToModelMessages, isToolUIPart, stepCountIs, streamText, toUIMessageStream } from "ai";
 import { Buffer } from "node:buffer";
@@ -1606,8 +1607,8 @@ const measuredInputTokens = (options) => {
 //#endregion
 //#region src/ai-agent.ts
 /**
-* The server-only implementation of the experimental-a2/ai protocol: the
-* handler table and `agent()`, its assembly as an ordinary A2 server.
+* The server-only implementation of the experimental-a2/ai protocol: its
+* consumers and `agent()`, its assembly as an ordinary A2 server.
 */
 const isAgentIngressEvent = (event) => {
 	if (event.type !== "ai.control.requested") return false;
@@ -1658,8 +1659,8 @@ const withModelSettings = (options) => {
 const instructionTexts = (instructions) => instructions === void 0 ? [] : typeof instructions === "string" ? [instructions] : (Array.isArray(instructions) ? instructions : [instructions]).map((message) => message.content);
 const errorMessage = (error) => error instanceof Error ? error.message : String(error);
 const toolScopeStorage = installAmbientToolScopeStorage(() => new AsyncLocalStorage());
-/** Run `fn` with `ctx` as the ambient tool handler context, for direct tool tests. */
-function runWithHandlerContext(ctx, fn) {
+/** Run `fn` with `ctx` as the ambient tool consumer context, for direct tool tests. */
+function runWithConsumerContext(ctx, fn) {
 	return toolScopeStorage.run({ context: ctx }, fn);
 }
 /** Sessions whose latest step prompt one process keeps for that step's tools. */
@@ -2118,7 +2119,7 @@ const assertNoLoopOptions = (options) => {
 		throw new TypeError(harness === "ai-sdk" ? `${key} belongs to the harness: agent({ harness: { type: 'ai-sdk', ${key} } })` : `${key} is an option of the AI SDK harness; ${harness}() runs its own loop`);
 	}
 };
-function createHandlers(options) {
+function createConsumers(options) {
 	if (options.model === void 0) throw new TypeError("agent requires model");
 	if (options.progress?.maxChunks !== void 0 && (!Number.isInteger(options.progress.maxChunks) || options.progress.maxChunks < 1)) throw new TypeError("progress.maxChunks must be a positive integer");
 	if (options.progress?.maxDelayMs !== void 0 && (!Number.isFinite(options.progress.maxDelayMs) || options.progress.maxDelayMs < 0)) throw new TypeError("progress.maxDelayMs must be a non-negative number");
@@ -2626,6 +2627,7 @@ function createHandlers(options) {
 			};
 			const updates = consumeGeneration({
 				source: harness !== void 0 ? await harness.step({
+					scope: states,
 					agentName: options.contract.name,
 					sessionId: ctx.event.sessionId,
 					requestIndex: ctx.event.index,
@@ -2742,8 +2744,14 @@ function createHandlers(options) {
 			};
 		}
 	};
-	return {
-		"ai.model.metadata.requested": { handler: async (ctx) => {
+	const contract = options.contract;
+	const { toolLane } = options;
+	return [
+		runtimeConsumer({
+			contract,
+			name: "a2.ai.metadata",
+			types: ["ai.model.metadata.requested"]
+		}).run(async (ctx) => {
 			const limits = await readModelLimits({
 				modelId: ctx.event.payload.modelId,
 				signal: ctx.signal
@@ -2757,27 +2765,29 @@ function createHandlers(options) {
 					limits
 				}
 			};
-		} },
-		"ai.control.requested": {
+		}),
+		runtimeConsumer({
+			contract,
+			name: "a2.ai.control",
+			types: [
+				"ai.control.requested",
+				"ai.work.reported",
+				"ai.session.closed"
+			],
 			lane: ({ event }) => {
+				if (event.type !== "ai.control.requested") return "a2.ai.control";
 				const command = event.payload;
 				if (event.id === "") throw new A2Error("INVALID_PAYLOAD", "AI command IDs must be nonempty");
 				if (event.id?.startsWith("a2.ai:")) throw new A2Error("INVALID_PAYLOAD", `AI command IDs cannot use the reserved '${AI_EVENT_ID_PREFIX}' prefix`);
 				if ((command.action === "send" || command.action === "steer") && event.id !== command.message.id) throw new A2Error("INVALID_PAYLOAD", "AI input event ID must equal message.id");
 				return "a2.ai.control";
-			},
-			handler: controlHandler
-		},
-		"ai.work.reported": {
-			lane: "a2.ai.control",
-			handler: controlHandler
-		},
-		"ai.session.closed": {
-			lane: "a2.ai.control",
-			handler: controlHandler
-		},
-		"ai.generation.requested": {
-			lane: "a2.ai.model",
+			}
+		}).run(controlHandler),
+		runtimeConsumer({
+			contract,
+			name: "a2.ai.model",
+			types: ["ai.generation.requested"],
+			lane: () => "a2.ai.model",
 			abortOn: {
 				"ai.control.requested": (event, trigger) => event.payload.action === "stop" && event.payload.turnId === trigger.payload.control?.turnId,
 				"ai.control.committed": (event, trigger) => {
@@ -2785,32 +2795,34 @@ function createHandlers(options) {
 					return active?.turnId !== trigger.payload.control?.turnId || active?.version !== trigger.payload.control?.version;
 				},
 				"ai.session.closed": true
-			},
-			handler: (ctx) => telemetry.span("a2.ai.generate", {
-				"a2.contract": options.contract.name,
-				"a2.session_id": ctx.event.sessionId,
-				"a2.ai.generate.request_id": ctx.event.id,
-				"a2.ai.generate.attempt": ctx.attempt
-			}, async () => {
-				try {
-					return control.settled({
-						ctx,
-						events: await generationHandler(ctx)
-					});
-				} catch (error) {
-					if (control.cancelled({
-						error,
-						signal: ctx.signal
-					})) return control.settled({
-						ctx,
-						events: void 0
-					});
-					throw error;
-				}
-			})
-		},
-		"ai.tool.execution.requested": {
-			...options.toolLane === void 0 ? {} : { lane: options.toolLane },
+			}
+		}).run((ctx) => telemetry.span("a2.ai.generate", {
+			"a2.contract": options.contract.name,
+			"a2.session_id": ctx.event.sessionId,
+			"a2.ai.generate.request_id": ctx.event.id,
+			"a2.ai.generate.attempt": ctx.attempt
+		}, async () => {
+			try {
+				return control.settled({
+					ctx,
+					events: await generationHandler(ctx)
+				});
+			} catch (error) {
+				if (control.cancelled({
+					error,
+					signal: ctx.signal
+				})) return control.settled({
+					ctx,
+					events: void 0
+				});
+				throw error;
+			}
+		})),
+		runtimeConsumer({
+			contract,
+			name: "a2.ai.tools",
+			types: ["ai.tool.execution.requested"],
+			...toolLane === void 0 ? {} : { lane: () => toolLane },
 			abortOn: {
 				"ai.control.requested": (event, trigger) => event.payload.action === "stop" && event.payload.turnId === trigger.payload.turnId,
 				"ai.generation.failed": (event, trigger) => event.payload.generationId === trigger.payload.call.generationId,
@@ -2819,21 +2831,21 @@ function createHandlers(options) {
 					return active?.turnId !== trigger.payload.turnId || active?.version !== trigger.payload.version;
 				},
 				"ai.session.closed": true
-			},
-			handler: (ctx) => telemetry.span("a2.ai.tool", {
-				"a2.contract": options.contract.name,
-				"a2.session_id": ctx.event.sessionId,
-				"a2.ai.tool.request_id": ctx.event.payload.call.requestId,
-				"a2.ai.tool.call_id": ctx.event.payload.call.toolCallId,
-				"a2.ai.tool.name": ctx.event.payload.call.toolName,
-				"a2.ai.tool.attempt": ctx.attempt
-			}, () => handleToolExecution(ctx))
-		}
-	};
+			}
+		}).run((ctx) => telemetry.span("a2.ai.tool", {
+			"a2.contract": options.contract.name,
+			"a2.session_id": ctx.event.sessionId,
+			"a2.ai.tool.request_id": ctx.event.payload.call.requestId,
+			"a2.ai.tool.call_id": ctx.event.payload.call.toolCallId,
+			"a2.ai.tool.name": ctx.event.payload.call.toolName,
+			"a2.ai.tool.attempt": ctx.attempt
+		}, () => handleToolExecution(ctx)))
+	];
 }
 function agent(options) {
 	assertNoLoopOptions(options);
-	const { model, tools, instructions, maxSteps, harness, toolRecovery, progress, toolLane, handlers, ...rest } = options;
+	if (options.toolLane !== void 0 && (typeof options.toolLane !== "string" || options.toolLane.length === 0)) throw new TypeError("toolLane must be a non-empty string");
+	const { model, tools, instructions, maxSteps, harness, toolRecovery, progress, toolLane, ...rest } = options;
 	let contract;
 	let serverOptions;
 	if ("contract" in rest) {
@@ -2852,7 +2864,7 @@ function agent(options) {
 		serverOptions = remaining;
 	}
 	const states = harnessStates();
-	const builtIns = createHandlers({
+	const builtIns = createConsumers({
 		contract,
 		model,
 		harnessStates: states,
@@ -2865,16 +2877,17 @@ function agent(options) {
 		...progress === void 0 ? {} : { progress },
 		...toolLane === void 0 ? {} : { toolLane }
 	});
-	for (const type of Object.keys(handlers ?? {})) if (Object.hasOwn(builtIns, type)) throw new TypeError(`agent handles '${type}' itself`);
 	const server = createServer({
 		...serverOptions,
 		contract,
-		handlers: {
-			...builtIns,
-			...handlers
-		}
+		consumers: [...builtIns, ...serverOptions.consumers ?? []]
 	});
 	registerHarnessStates(server, states);
+	const release = (harness === void 0 || harness.type === "ai-sdk" ? void 0 : harnessDriver(harness))?.release;
+	if (release !== void 0) drainSettledHooks.set(server, (sessionId) => release({
+		scope: states,
+		sessionId
+	}));
 	const applicationEvents = "contract" in rest;
 	setServerFetchHooks(server, { validateIngress: (context) => validateAgentIngress({
 		context,
@@ -2883,6 +2896,6 @@ function agent(options) {
 	return server;
 }
 //#endregion
-export { agent, runWithHandlerContext };
+export { agent, runWithConsumerContext };
 
 //# sourceMappingURL=ai-server.js.map