@empyria/restate 0.2.2 → 0.4.5

package/AGENTS.md

@@ -4,10 +4,9 @@ Restate.dev helpers for **Empyria**, a nanoservice framework built primarily on
 an Admin API client, a dynamic-dispatch service, cron-driven workflow scheduling, and
 schema validation for workflow handlers.
 
-Targets Restate server `1.7.x` and `@restatedev/restate-sdk` `1.6.x` — check both when
-touching version pins; the SDK's own version numbering runs well ahead of the server's
-(e.g. `1.16.x` was latest on npm when `1.6.x` was what this repo targets), so "latest"
-is not the right default here.
+Targets Restate server `1.7.x` and `@restatedev/restate-sdk` `1.17.x` (pinned in `package.json`) —
+check both when touching version pins; the SDK's own version numbering runs well ahead of the
+server's, so "latest" is not the right default here.
 
 ## Runtime
 
@@ -32,14 +31,31 @@ SDK packages (`ai`, `@ai-sdk/*`) are **optional** peer dependencies (and devDepe
 the tests): a consumer that only uses Admin/Cron/Pubsub must be able to import the package
 root without them installed. Never import `lib/agent/` from a root-exported module.
 
-Agent tests use `test/agent/fakes.js`'s `fakeCtx`, which mimics the journal's JSON round trip
+Agent tests use `fakeCtx` from `lib/testing/Fakes.js` (exported as `@empyria/restate/testing`; `test/agent/fakes.js` re-exports it), which mimics the journal's JSON round trip
 and a bounded `ctx.run` turning an exhausted retry into a `TerminalError` — keep it in step
 with the real SDK behaviour when touching either.
 
+## Keep this package generic
+
+Only things that are generic to Restate or to LLM agents belong here. Anything specific to one
+consumer (its folder layout, its service names, its prompts, its UI) stays in that consumer. A consumer
+that finds itself writing a generic helper should add it here with tests, release it, then bump its pin.
+
+Journal discipline is enforced by the library, not left to each agent: step metadata is small numbers
+only (never prompt or reply text), and tool arguments/results have byte caps (`Tools.js`). Keep both
+true when touching `Loop.js` or `Tools.js`.
+
+Block tools take `key(args, meta)` and `input(args, meta)` hooks (`meta.idempotencyKey` is unique per tool call); both must be
+total, because an exception there is a non-terminal error that Restate would retry forever.
+
+`defineAgent`'s `tools` may be a function of the request; it is evaluated outside `ctx.run` on every
+replay, so it must stay pure. `test/agent/fakes.js`'s result helpers use the AI SDK's v4 result shape
+(`finishReason: {unified, raw}`); the older plain-string form makes `generateText` report no finish reason.
+
 ## No official Admin API client
 
 Verified directly against the published `@restatedev/restate-sdk`/`-clients` package
-internals (both the `1.6.x` line this repo targets and the latest available at the time):
+internals (both the line this repo targets and the latest available at the time):
 zero Admin API exports. Only the ingress/invocation client and the service-authoring API
 (`restate.service`/`object`/`workflow`) are official. `lib/Admin.js` fetches the Admin API's
 REST endpoints directly — that's not a stopgap, it's the only way to do this today.

package/README.md

@@ -4,7 +4,7 @@
 primarily on Bun: an Admin API client, a generic dynamic-dispatch service, cron-driven
 workflow scheduling, and Restate-aware input/output schema validation.
 
-Targets Restate server `1.7.x` and `@restatedev/restate-sdk` `1.6.x`.
+Targets Restate server `1.7.x` and `@restatedev/restate-sdk` `1.17.x`.
 
 ## Requirements
 
@@ -106,6 +106,41 @@ export const triage = defineAgent({
 | [lib/agent/Errors.js](./lib/agent/Errors.js) | `toRestateLLMError` (AI SDK error → terminal / `RetryableError` / transient) and the default retry policies for LLM steps, tool steps and the agent invocation.                                                                                                                                             |
 | [lib/agent/Model.js](./lib/agent/Model.js)   | `createModel` — AI SDK model for an OpenAI-compatible endpoint (e.g. Bifrost).                                                                                                                                                                                                                              |
 
+### Shaping an agent
+
+`defineAgent` options beyond the basics:
+
+| Option                   | Does                                                                                                                                                                                                                                       |
+| ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `tools` as a function    | `tools: (input) => ({...})` builds the tools from the validated request, for tools that depend on it. It runs outside any `ctx.run` on every execution and replay, so it must be pure. An invalid map fails that request terminally (500). |
+| `respond(result, input)` | Shapes what `ask` returns from the loop result `{text, output?, steps, totalTokens, trace}`. Pair it with `responseSchema`, which validates the shaped value before Restate records it.                                                    |
+| `label(input)`           | A terminal failure then reads `[label] message`, so a failed run says what it was working on. Code and metadata are kept; cancellations, timeouts and retryable errors are untouched.                                                      |
+| `limits.maxInputBytes`   | Largest request (JSON bytes, default 64 KB); a larger one is refused terminally (413).                                                                                                                                                     |
+
+**Keying block tools.** A block tool calls another Restate building block. Its `key(args, meta)` and `input(args, meta)`
+hooks decide the target key and the parameter from the model's arguments and a per-call idempotency key
+(`meta.idempotencyKey`, unique for every tool call). A key that is the same for several calls (e.g. `lock:<resource>`)
+makes a virtual object run those calls one at a time and queue the rest durably; a key built from the idempotency key
+makes every call independent. `input` adds what the model must not control, such as the caller or the allowed folder.
+Both hooks must not throw: validate before the run, not inside them.
+
+**Size caps on tool calls.** Tool arguments and results are journaled and replayed, so each tool has
+`maxInputBytes` (default 16 KB) and `maxOutputBytes` (default 32 KB). Oversized arguments are not run and
+the model is told why; an oversized result is replaced by an error the model sees. For an `execute` tool the
+check happens inside its own `ctx.run`, so the big value is never recorded.
+
+**What a run records.** Each `llm-step-N` journals a small `meta` (`model`, `finishReason`, `inputTokens`,
+`outputTokens`, `durationMs`), never the prompt or the reply. The loop returns the same facts as `trace`
+(one entry per step with the tool names asked for) and logs one `ctx.console` line per step, which stays
+silent on replay.
+
+**Configuration.** `readLlmEnv(process.env)` returns `{LLM_BASE_URL, LLM_API_KEY?, MODEL_ID}` and fails with a
+terminal 500 that names the missing setting and never prints a value.
+
+**Testing agents.** `@empyria/restate/testing` exports `fakeCtx` (a journaling `run`, `genericCall`,
+`awakeable`, state and a capturing `console`) and the model-result helpers `textResult`, `toolCallResult`
+and `sequence`, without loading the AI SDK.
+
 Error policy: a tool that fails terminally (rejected or timed-out approval, a callee's
 `TerminalError`, an execute tool out of retries) goes back to the model as an error result,
 so the agent can adapt. Transient errors are left to Restate to retry, and cancellation

package/lib/agent/Agent.js

@@ -4,7 +4,10 @@ import { withValidation } from '../Validation.js'
 import { DEFAULT_AGENT_RETRY_POLICY } from './Errors.js'
 import { runAgentLoop } from './Loop.js'
 import { createModel } from './Model.js'
-import { assertToolSpecs } from './Tools.js'
+import { assertToolSpecs, jsonBytes } from './Tools.js'
+
+/** Largest request an agent accepts, in JSON bytes: the input is journaled and sent to the model. */
+export const DEFAULT_MAX_INPUT_BYTES = 64 * 1024
 
 const HISTORY = 'history'
 const HISTORY_REF = 'historyRef'
@@ -57,20 +60,35 @@ export const AskOutputSchema = {
  *   body is fetched again before every model call. So it must return the same content
  *   for the whole invocation (skills shipped with the deployment do), or the invocation
  *   fails terminally rather than send the model changed instructions.
- * @property {Record<string, import('./Tools.js').ToolSpec>} [tools]
+ * @property {Record<string, import('./Tools.js').ToolSpec> | ((input: any) => Record<string, import('./Tools.js').ToolSpec> | Promise<Record<string, import('./Tools.js').ToolSpec>>)} [tools]
+ *   A fixed map, or a function of the validated request for tools that depend on it (for example
+ *   confined to a folder the request names). The function runs outside any `ctx.run`, on every
+ *   execution and replay: it must be pure and deterministic, and it must not do I/O. The returned
+ *   map is checked like a fixed one; an invalid map fails the request terminally (500).
  * @property {object} [input] Input JSON Schema. Defaults to {@link AskInputSchema}.
  * @property {(input: any) => string} [prompt] Builds the user prompt from the validated
  *   input. Defaults to `input.prompt`.
  * @property {object} [output] Output JSON Schema. When given, the model must answer through
  *   the `final_answer` tool and `ask` returns its arguments. Defaults to
  *   {@link AskOutputSchema} (`{text, steps}`).
+ * @property {(result: {text: string, output?: unknown, steps: number, totalTokens: number, trace: Array<object>}, input: any) => any} [respond]
+ *   Shapes what `ask` returns from the loop's result (`output` is set when an `output` schema is
+ *   given; `trace` is the per-step summary). Without it, `ask` returns `output`, or `{text, steps}`.
+ * @property {object} [responseSchema] JSON Schema of what `respond` returns, checked before Restate
+ *   records it (a mismatch is terminal, 500). Only used together with `respond`.
+ * @property {(input: any) => string} [label] A short name for this request, for example the item it
+ *   works on. A terminal failure inside the agent then reads `[label] message`, so a failed run says
+ *   what it was working on. The error's code and metadata are kept; cancellations and timeouts are
+ *   left alone.
  * @property {'none'|'session'} [memory] `none` (default): a stateless service. `session`: a
  *   virtual object keyed by session ID, so turns of one session run one at a time and see
  *   the conversation so far.
  * @property {HistoryStore} [historyStore] Session memory only. Without it, the conversation
  *   is kept in object state, trimmed to `limits.maxHistoryTurns` — fine for short
  *   sessions; long ones should use a store.
- * @property {{maxSteps?: number, maxTokens?: number, maxHistoryTurns?: number}} [limits]
+ * @property {{maxSteps?: number, maxTokens?: number, maxHistoryTurns?: number, maxInputBytes?: number}} [limits]
+ *   `maxInputBytes` caps the request (JSON bytes, default {@link DEFAULT_MAX_INPUT_BYTES}); a larger
+ *   one is refused terminally (413).
  * @property {import('@restatedev/restate-sdk').RunOptions<any>} [llmRetry] Retry policy of
  *   each LLM call step. Defaults to `DEFAULT_LLM_RETRY`.
  * @property {import('@restatedev/restate-sdk').RetryPolicy} [retryPolicy] Invocation-level
@@ -109,6 +127,9 @@ export function defineAgent(spec) {
 		input = AskInputSchema,
 		prompt = (validatedInput) => validatedInput.prompt,
 		output,
+		respond,
+		responseSchema,
+		label,
 		memory = 'none',
 		historyStore,
 		limits = {},
@@ -126,17 +147,51 @@ export function defineAgent(spec) {
 	if (historyStore && memory !== 'session') {
 		throw new TypeError(`Agent '${name}': historyStore needs memory: 'session'`)
 	}
-	assertToolSpecs(tools)
+	if (typeof tools !== 'function') assertToolSpecs(tools)
+	if (respond !== undefined && typeof respond !== 'function') {
+		throw new TypeError(`Agent '${name}': respond must be a function`)
+	}
+	if (label !== undefined && typeof label !== 'function') {
+		throw new TypeError(`Agent '${name}': label must be a function`)
+	}
+
+	const {
+		maxSteps,
+		maxTokens,
+		maxHistoryTurns = 20,
+		maxInputBytes = DEFAULT_MAX_INPUT_BYTES,
+	} = limits
 
-	const { maxSteps, maxTokens, maxHistoryTurns = 20 } = limits
+	const outputSchema = respond ? (responseSchema ?? {}) : (output ?? AskOutputSchema)
 
-	const ask = withValidation(input, output ?? AskOutputSchema, async (ctx, validatedInput) => {
+	const ask = withValidation(input, outputSchema, async (ctx, validatedInput) => {
+		const inputBytes = jsonBytes(validatedInput)
+		if (inputBytes > maxInputBytes) {
+			throw withLabel(
+				new TerminalError(
+					`Agent '${name}': the request is ${inputBytes} bytes, over the limit of ${maxInputBytes}`,
+					{ errorCode: 413 },
+				),
+				label,
+				validatedInput,
+			)
+		}
+		try {
+			return await runAsk(ctx, validatedInput)
+		} catch (error) {
+			throw withLabel(error, label, validatedInput)
+		}
+	})
+
+	async function runAsk(ctx, validatedInput) {
 		const userPrompt = prompt(validatedInput)
 		if (typeof userPrompt !== 'string' || userPrompt === '') {
 			throw new TerminalError(`Agent '${name}': the input produced no prompt`, {
 				errorCode: 400,
 			})
 		}
+		const requestTools =
+			typeof tools === 'function' ? await resolveTools(name, tools, validatedInput) : tools
 
 		const context = loadContext
 			? await ctx.run('load-context', async () => {
@@ -158,7 +213,7 @@ export function defineAgent(spec) {
 			model,
 			system: [system, context.systemPrompt].filter(Boolean).join('\n\n'),
 			prompt: userPrompt,
-			tools,
+			tools: requestTools,
 			skills: context.skills,
 			loadSkill: skillLoader({ loadSkill, loadContext, input: validatedInput }),
 			loadHistory: session?.load,
@@ -170,8 +225,9 @@ export function defineAgent(spec) {
 
 		if (session) await session.save(result.messages, maxHistoryTurns)
 
+		if (respond) return respond(result, validatedInput)
 		return output ? result.output : { text: result.text, steps: result.steps }
-	})
+	}
 
 	const options = { retryPolicy }
 
@@ -235,6 +291,34 @@ function skillLoader({ loadSkill, loadContext, input }) {
 	}
 }
 
+/** Resolves a request-dependent `tools` function, failing the request (not retrying) on a bad map. */
+async function resolveTools(name, tools, input) {
+	try {
+		const resolved = await tools(input)
+		assertToolSpecs(resolved ?? {})
+		return resolved ?? {}
+	} catch (error) {
+		throw new TerminalError(
+			`Agent '${name}': invalid tools for this request: ${error.message}`,
+			{
+				errorCode: 500,
+			},
+		)
+	}
+}
+
+/**
+ * Prefixes a plain `TerminalError` with the request's label. Anything else (a retryable error, a
+ * cancellation, a timeout) is returned untouched, so Restate's own handling of it is unchanged.
+ */
+function withLabel(error, label, input) {
+	if (!label || error?.constructor !== TerminalError) return error
+	return new TerminalError(`[${label(input)}] ${error.message}`, {
+		errorCode: error.code,
+		metadata: error.metadata,
+	})
+}
+
 /**
  * @deprecated Use {@link defineAgent} with an explicit `name`. Kept so callers migrating
  *   from `@empyria/restate-llm` only change their import: defines a stateless agent named

package/lib/agent/Loop.js

@@ -37,8 +37,19 @@ const FINAL_ANSWER_NUDGE = `Respond by calling the ${FINAL_ANSWER_TOOL} tool wit
  * }} params `skills` is the index only (name + description); `loadSkill` fetches a body
  *   when the model asks for it. With `outputSchema`, the model must answer through the
  *   `final_answer` tool and the loop returns that tool's input as `output`.
- * @returns {Promise<{text: string, output?: unknown, steps: number, totalTokens: number, messages: Array<object>}>}
+ * @returns {Promise<{text: string, output?: unknown, steps: number, totalTokens: number, trace: Array<StepTrace>, messages: Array<object>}>}
  *   `messages` are the new messages of this turn only (prompt, responses, tool results).
+ *   `trace` has one small entry per LLM step. The same facts are recorded in the journal with each
+ *   `llm-step-N` (as `meta`) and logged through `ctx.console`, which stays silent on replay.
+ *
+ * @typedef {Object} StepTrace What one LLM step did, with nothing from the prompt or the reply.
+ * @property {number} step
+ * @property {string} [model] The model ID as the provider reports it.
+ * @property {string} [finishReason] `stop`, `tool-calls`, `length`, ...
+ * @property {number} [inputTokens]
+ * @property {number} [outputTokens]
+ * @property {number} [durationMs] Wall-clock time of the model call.
+ * @property {string[]} tools Names of the tools the model asked for in this step.
  */
 export async function runAgentLoop(
 	ctx,
@@ -61,6 +72,7 @@ export async function runAgentLoop(
 	const messages = [{ role: 'user', content: prompt }]
 	let totalTokens = 0
 	const fetchSkill = cachedSkillFetcher(loadSkill)
+	const trace = []
 
 	for (let step = 0; step < maxSteps; step++) {
 		// `ctx.run` JSON-serialises the callback's result for the journal and hands back that
@@ -69,6 +81,7 @@ export async function runAgentLoop(
 		const result = await ctx.run(
 			`llm-step-${step}`,
 			async () => {
+				const startedAt = Date.now()
 				try {
 					// `maxRetries: 0`: `ctx.run`'s RunOptions are the single retry authority,
 					// instead of AI SDK's own retries stacking invisibly underneath them.
@@ -94,6 +107,14 @@ export async function runAgentLoop(
 						})),
 						responseMessages: r.response.messages,
 						totalTokens: r.usage?.totalTokens ?? 0,
+						// A few numbers: enough to see whether the step worked, and what it cost.
+						meta: {
+							model: model.modelId,
+							finishReason: r.finishReason,
+							inputTokens: r.usage?.inputTokens,
+							outputTokens: r.usage?.outputTokens,
+							durationMs: Date.now() - startedAt,
+						},
 					}
 				} catch (error) {
 					throw toRestateLLMError(error)
@@ -102,6 +123,13 @@ export async function runAgentLoop(
 			retry,
 		)
 
+		const entry = { step, ...result.meta, tools: result.toolCalls.map((call) => call.toolName) }
+		trace.push(entry)
+		// `ctx.console` is replay-aware: one line per step, never repeated when Restate replays.
+		ctx.console?.log(
+			`llm-step-${step} model=${entry.model ?? '?'} finish=${entry.finishReason ?? '?'} in=${entry.inputTokens ?? '?'} out=${entry.outputTokens ?? '?'} ${entry.durationMs ?? '?'}ms tools=[${entry.tools.join(',')}]`,
+		)
+
 		messages.push(...result.responseMessages)
 		totalTokens += result.totalTokens
 		if (maxTokens !== undefined && totalTokens > maxTokens) {
@@ -113,7 +141,7 @@ export async function runAgentLoop(
 
 		if (result.toolCalls.length === 0) {
 			if (!outputSchema) {
-				return { text: result.text, steps: step + 1, totalTokens, messages }
+				return { text: result.text, steps: step + 1, totalTokens, trace, messages }
 			}
 			messages.push({ role: 'user', content: FINAL_ANSWER_NUDGE })
 			continue
@@ -126,6 +154,7 @@ export async function runAgentLoop(
 					output: call.input,
 					steps: step + 1,
 					totalTokens,
+					trace,
 					messages,
 				}
 			}

package/lib/agent/Model.js

@@ -1,5 +1,45 @@
+import { TerminalError } from '@restatedev/restate-sdk'
+import { createValidator } from '@empyria/common'
 import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
 
+/** JSON Schema of the environment variables that describe an LLM connection. */
+export const LlmEnvSchema = {
+	type: 'object',
+	required: ['LLM_BASE_URL', 'MODEL_ID'],
+	properties: {
+		LLM_BASE_URL: { type: 'string', minLength: 1 },
+		LLM_API_KEY: { type: 'string' },
+		MODEL_ID: { type: 'string', minLength: 1 },
+	},
+}
+
+const validateLlmEnv = createValidator(LlmEnvSchema)
+
+/**
+ * Reads the LLM connection from the environment (or any object) and checks it. A missing setting is
+ * a deployment mistake that no retry can fix, so it fails as a `TerminalError` (code 500) that names
+ * the problem and never prints a value, in particular not the API key.
+ * @param {Record<string, string|undefined>} [source] Defaults to `process.env`.
+ * @returns {{LLM_BASE_URL: string, LLM_API_KEY?: string, MODEL_ID: string}} Only the three LLM
+ *   settings, ready for {@link createModel}.
+ * @throws {TerminalError}
+ */
+export function readLlmEnv(source = process.env) {
+	const picked = {}
+	for (const key of ['LLM_BASE_URL', 'LLM_API_KEY', 'MODEL_ID']) {
+		if (typeof source[key] === 'string' && source[key] !== '') picked[key] = source[key]
+	}
+	try {
+		return validateLlmEnv(picked)
+	} catch {
+		const missing = LlmEnvSchema.required.filter((key) => !(key in picked))
+		throw new TerminalError(
+			`LLM is not configured: set ${missing.join(' and ') || 'LLM_BASE_URL and MODEL_ID'} (and usually LLM_API_KEY)`,
+			{ errorCode: 500 },
+		)
+	}
+}
+
 /**
  * AI SDK Core model, pointed at an OpenAI-compatible endpoint — Bifrost
  * (llm-system's gateway) by default, fully overridable via env. `LLM_BASE_URL`

package/lib/agent/Tools.js

@@ -12,6 +12,19 @@ export const FINAL_ANSWER_TOOL = 'final_answer'
 
 const TOOL_NAME = /^[a-zA-Z0-9_-]{1,64}$/
 
+/**
+ * Default size caps for a tool call, in bytes of its JSON form. Tool arguments are recorded in the
+ * journal as part of the model's response and tool results are recorded when the tool finishes, and
+ * Restate replays all of it on every retry: only small values belong there. Override per tool with
+ * `maxInputBytes` / `maxOutputBytes`; a tool that must handle bigger data should keep it outside the
+ * journal and pass a reference instead.
+ */
+export const DEFAULT_TOOL_MAX_INPUT_BYTES = 16 * 1024
+export const DEFAULT_TOOL_MAX_OUTPUT_BYTES = 32 * 1024
+
+/** Size of a value's JSON form, in bytes. */
+export const jsonBytes = (value) => Buffer.byteLength(JSON.stringify(value) ?? '')
+
 /**
  * @typedef {Object} ApprovalSpec Human (or system) approval required before a tool runs.
  * @property {number} timeout Milliseconds to wait for a decision. A timeout counts as a
@@ -27,8 +40,18 @@ const TOOL_NAME = /^[a-zA-Z0-9_-]{1,64}$/
  * @property {object} inputSchema JSON Schema of the tool's arguments, shown to the model.
  * @property {string|{name: string}} block Target service/object/workflow, by name or definition.
  * @property {string} handler Target handler.
- * @property {(args: any) => string} [key] Virtual object / workflow key, derived from the arguments.
+ * @property {(args: any, meta: {idempotencyKey: string}) => string} [key] Virtual object / workflow key, derived from
+ *   the model's arguments and a per-call idempotency key (`<invocation id>:<step>.<index>`, unique for every tool
+ *   call). A key that is the same for several calls makes those calls run one at a time; a key built from
+ *   `meta.idempotencyKey` makes every call its own key, so they run in parallel. Must not throw.
+ * @property {(args: any, meta: {idempotencyKey: string}) => any} [input] Builds the parameter sent to the target
+ *   handler from the model's arguments, for adding what the model must not control (who the caller is, which
+ *   folder is allowed) or the idempotency key. Default: the arguments as they are. Must not throw.
  * @property {ApprovalSpec} [approval]
+ * @property {number} [maxInputBytes] Largest accepted arguments (JSON bytes). Default
+ *   {@link DEFAULT_TOOL_MAX_INPUT_BYTES}. A larger call is not made; the model is told why.
+ * @property {number} [maxOutputBytes] Largest result handed to the model (JSON bytes). Default
+ *   {@link DEFAULT_TOOL_MAX_OUTPUT_BYTES}. A larger result is replaced by an error the model sees.
  *
  * @typedef {Object} ExecuteToolSpec A tool executed as a direct side effect inside its own `ctx.run`.
  * @property {string} description
@@ -39,6 +62,9 @@ const TOOL_NAME = /^[a-zA-Z0-9_-]{1,64}$/
  * @property {import('@restatedev/restate-sdk').RunOptions<any>} [retry] Defaults to
  *   {@link DEFAULT_TOOL_RETRY}.
  * @property {ApprovalSpec} [approval]
+ * @property {number} [maxInputBytes] As on {@link BlockToolSpec}.
+ * @property {number} [maxOutputBytes] As on {@link BlockToolSpec}. Checked inside the tool's own
+ *   `ctx.run`, so an oversized result is never recorded in the journal.
  *
  * @typedef {BlockToolSpec|ExecuteToolSpec} ToolSpec
  */
@@ -68,6 +94,16 @@ export function assertToolSpecs(tools) {
 		if (isBlock && typeof spec.handler !== 'string') {
 			throw new TypeError(`Block tool '${name}' needs a 'handler'`)
 		}
+		for (const hook of ['key', 'input']) {
+			if (spec[hook] !== undefined && typeof spec[hook] !== 'function') {
+				throw new TypeError(`Tool '${name}': ${hook} must be a function`)
+			}
+		}
+		for (const cap of ['maxInputBytes', 'maxOutputBytes']) {
+			if (spec[cap] !== undefined && !(Number.isInteger(spec[cap]) && spec[cap] > 0)) {
+				throw new TypeError(`Tool '${name}': ${cap} must be a positive integer`)
+			}
+		}
 		if (spec.approval) {
 			const { timeout, notify } = spec.approval
 			if (!(timeout > 0) || !notify?.block || !notify?.handler) {
@@ -142,6 +178,16 @@ export async function executeTool(ctx, call, { tools, loadSkill, step, index })
 		const spec = tools[call.toolName]
 		if (!spec) return { type: 'error-text', value: `Unknown tool "${call.toolName}".` }
 
+		const maxIn = spec.maxInputBytes ?? DEFAULT_TOOL_MAX_INPUT_BYTES
+		const inBytes = jsonBytes(call.input)
+		if (inBytes > maxIn) {
+			return {
+				type: 'error-text',
+				value: `Tool "${call.toolName}" was not run: its arguments are ${inBytes} bytes, over the limit of ${maxIn}. Send less.`,
+			}
+		}
+		const maxOut = spec.maxOutputBytes ?? DEFAULT_TOOL_MAX_OUTPUT_BYTES
+
 		if (spec.approval) {
 			const decision = await requestApproval(ctx, call, spec.approval)
 			if (!decision.approved) {
@@ -152,24 +198,27 @@ export async function executeTool(ctx, call, { tools, loadSkill, step, index })
 			}
 		}
 
+		const meta = { idempotencyKey: `${ctx.request().id}:${stepId}` }
 		const value = spec.execute
 			? await ctx.run(
 					`tool:${call.toolName}-${stepId}`,
-					async () =>
-						(await spec.execute(call.input, {
-							idempotencyKey: `${ctx.request().id}:${stepId}`,
-						})) ?? null,
+					async () => {
+						const result = (await spec.execute(call.input, meta)) ?? null
+						assertOutputSize(call.toolName, result, maxOut)
+						return result
+					},
 					spec.retry ?? DEFAULT_TOOL_RETRY,
 				)
 			: await ctx.genericCall({
 					service: targetName(spec.block),
 					method: spec.handler,
-					key: spec.key?.(call.input),
-					parameter: call.input,
+					key: spec.key?.(call.input, meta),
+					parameter: spec.input ? spec.input(call.input, meta) : call.input,
 					inputSerde: serde.json,
 					outputSerde: serde.json,
 				})
 
+		if (!spec.execute) assertOutputSize(call.toolName, value, maxOut)
 		return { type: 'json', value: value ?? null }
 	} catch (error) {
 		if (error instanceof CancelledError || !(error instanceof TerminalError)) throw error
@@ -177,6 +226,17 @@ export async function executeTool(ctx, call, { tools, loadSkill, step, index })
 	}
 }
 
+/** @throws {TerminalError} when a tool result is too big to hand to the model or record. */
+function assertOutputSize(toolName, value, max) {
+	const bytes = jsonBytes(value)
+	if (bytes > max) {
+		throw new TerminalError(
+			`Tool "${toolName}" returned ${bytes} bytes, over the limit of ${max}. Ask for less, or a part of it.`,
+			{ errorCode: 413 },
+		)
+	}
+}
+
 /**
  * Asks for approval through an awakeable and waits for the decision, durably.
  * @param {import('@restatedev/restate-sdk').Context} ctx

package/lib/testing/Fakes.js

@@ -0,0 +1,92 @@
+import { TerminalError } from '@restatedev/restate-sdk'
+
+/**
+ * Test doubles for agent code, shipped as `@empyria/restate/testing`. No dependency on the AI SDK.
+ *
+ * A stand-in for Restate's `Context` / `ObjectContext`. Per this repo's testing convention,
+ * handlers are unit-tested against a mocked `ctx`, never a live Restate server.
+ *
+ * `run` mimics the parts of the real one the agent depends on: the result goes through a
+ * JSON round trip (as it does through the journal), and when a bounded `ctx.run` gives up
+ * it throws a `TerminalError` wrapping the original message. One attempt stands in for
+ * "all retries exhausted".
+ */
+export function fakeCtx({ key, state = {}, onCall, onApproval } = {}) {
+	const steps = []
+	const calls = []
+	const sends = []
+	const logs = []
+	return {
+		key,
+		logs,
+		console: { log: (...args) => logs.push(args.join(' ')) },
+		state,
+		steps,
+		calls,
+		sends,
+		run: async (name, fn, options) => {
+			const step = { name, options }
+			steps.push(step)
+			let value
+			try {
+				value = await fn()
+			} catch (error) {
+				if (error instanceof TerminalError || !options?.maxRetryAttempts) throw error
+				throw new TerminalError(error.message)
+			}
+			step.result = value === undefined ? undefined : JSON.parse(JSON.stringify(value))
+			return step.result
+		},
+		request: () => ({ id: 'inv-1' }),
+		genericCall: async (call) => {
+			calls.push(call)
+			return onCall ? onCall(call) : null
+		},
+		genericSend: (send) => {
+			sends.push(send)
+		},
+		awakeable: () => ({
+			id: 'awk-1',
+			promise: { orTimeout: async (timeout) => onApproval?.(timeout) },
+		}),
+		get: async (name) => state[name],
+		set: (name, value) => {
+			state[name] = value
+		},
+		clear: (name) => {
+			delete state[name]
+		},
+	}
+}
+
+const usage = { inputTokens: { total: 1 }, outputTokens: { total: 1 } }
+
+export function textResult(text) {
+	return {
+		content: [{ type: 'text', text }],
+		finishReason: { unified: 'stop', raw: 'stop' },
+		usage,
+		warnings: [],
+	}
+}
+
+/** @param {Array<{toolName: string, input: object}>} calls */
+export function toolCallResult(calls) {
+	return {
+		content: calls.map(({ toolName, input }, i) => ({
+			type: 'tool-call',
+			toolCallId: `call_${i}`,
+			toolName,
+			input: JSON.stringify(input),
+		})),
+		finishReason: { unified: 'tool-calls', raw: 'tool_calls' },
+		usage,
+		warnings: [],
+	}
+}
+
+/** `doGenerate` that returns `results` one after another. */
+export function sequence(...results) {
+	let i = 0
+	return async () => results[Math.min(i++, results.length - 1)]
+}

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@empyria/restate",
-	"version": "0.2.2",
+	"version": "0.4.5",
 	"description": "Restate.dev helpers for the Empyria nanoservice framework",
 	"license": "MIT",
 	"author": "Imre Fazekas <imre.fazekas@icloud.com>",
@@ -13,6 +13,7 @@
 	"exports": {
 		".": "./index.js",
 		"./agent": "./agent.js",
+		"./testing": "./testing.js",
 		"./lib/*": "./lib/*.js",
 		"./package.json": "./package.json"
 	},

package/test/Testing.test.js

@@ -0,0 +1,16 @@
+import { describe, test, expect } from 'bun:test'
+import * as testing from '../testing.js'
+
+describe('@empyria/restate/testing', () => {
+	test('exports the fakes agent tests need, without loading the AI SDK', () => {
+		expect(Object.keys(testing).sort()).toEqual([
+			'fakeCtx',
+			'sequence',
+			'textResult',
+			'toolCallResult',
+		])
+		const ctx = testing.fakeCtx()
+		ctx.console.log('a', 'b')
+		expect(ctx.logs).toEqual(['a b'])
+	})
+})

package/test/agent/Agent.test.js

@@ -193,3 +193,164 @@ describe('defineAgent: session memory', () => {
 		expect(journaled).not.toContain('first question')
 	})
 })
+
+describe('defineAgent: tools that depend on the request', () => {
+	const ctxInput = {
+		type: 'object',
+		properties: { prompt: { type: 'string', minLength: 1 }, folder: { type: 'string' } },
+		required: ['prompt', 'folder'],
+	}
+
+	test('a tools function receives the validated request and its tools are used', async () => {
+		const seen = []
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'where', input: {} }]),
+				textResult('ok'),
+			),
+		})
+		const agent = defineAgent({
+			name: 'A',
+			model,
+			input: ctxInput,
+			tools: (request) => {
+				seen.push(request.folder)
+				return {
+					where: {
+						description: 'Where am I',
+						inputSchema: { type: 'object' },
+						execute: async () => ({ folder: request.folder }),
+					},
+				}
+			},
+		})
+		const ctx = fakeCtx()
+		await agent.service.ask(ctx, { prompt: 'hi', folder: '/f' })
+		expect(seen.length).toBeGreaterThan(0)
+		expect(new Set(seen)).toEqual(new Set(['/f']))
+		expect(ctx.steps.find((s) => s.name.startsWith('tool:where')).result).toEqual({
+			folder: '/f',
+		})
+	})
+
+	test('an invalid tools map for one request fails that request terminally (500)', async () => {
+		const agent = defineAgent({
+			name: 'A',
+			model: okModel(),
+			tools: () => ({ load_skill: { description: 'x', inputSchema: {} } }),
+		})
+		const error = await agent.service.ask(fakeCtx(), { prompt: 'hi' }).catch((e) => e)
+		expect(error).toBeInstanceOf(TerminalError)
+		expect(error.code).toBe(500)
+	})
+})
+
+describe('defineAgent: respond and responseSchema', () => {
+	test('respond shapes the return value from the loop result', async () => {
+		const agent = defineAgent({
+			name: 'A',
+			model: okModel(),
+			respond: (r, input) => ({
+				answer: r.text,
+				steps: r.steps,
+				trace: r.trace.length,
+				echoed: input.prompt,
+			}),
+			responseSchema: { type: 'object', required: ['answer'] },
+		})
+		expect(await agent.service.ask(fakeCtx(), { prompt: 'hi' })).toEqual({
+			answer: 'ok',
+			steps: 1,
+			trace: 1,
+			echoed: 'hi',
+		})
+	})
+
+	test('a response that breaks responseSchema is terminal (500)', async () => {
+		const agent = defineAgent({
+			name: 'A',
+			model: okModel(),
+			respond: () => ({ nope: true }),
+			responseSchema: { type: 'object', required: ['answer'] },
+		})
+		const error = await agent.service.ask(fakeCtx(), { prompt: 'hi' }).catch((e) => e)
+		expect(error).toBeInstanceOf(TerminalError)
+		expect(error.code).toBe(500)
+	})
+
+	test('without respond the default result is unchanged', async () => {
+		expect(
+			await defineAgent({ name: 'A', model: okModel() }).service.ask(fakeCtx(), {
+				prompt: 'hi',
+			}),
+		).toEqual({ text: 'ok', steps: 1 })
+	})
+
+	test('rejects respond or label that is not a function', () => {
+		expect(() => defineAgent({ name: 'A', model: okModel(), respond: 1 })).toThrow(TypeError)
+		expect(() => defineAgent({ name: 'A', model: okModel(), label: 'x' })).toThrow(TypeError)
+	})
+})
+
+describe('defineAgent: label', () => {
+	test('prefixes a terminal failure with the label and keeps its code', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: toolCallResult([{ toolName: 'nope', input: {} }]),
+		})
+		const agent = defineAgent({
+			name: 'A',
+			model,
+			limits: { maxSteps: 2 },
+			label: (input) => `job ${input.prompt}`,
+		})
+		const error = await agent.service.ask(fakeCtx(), { prompt: 'J1' }).catch((e) => e)
+		expect(error.constructor).toBe(TerminalError)
+		expect(error.message).toStartWith('[job J1] ')
+		expect(error.message).toContain('maxSteps')
+	})
+
+	test('leaves cancellations and non-terminal errors untouched', async () => {
+		const { CancelledError } = await import('@restatedev/restate-sdk')
+		const cancelled = new CancelledError()
+		const plain = new Error('transient')
+		for (const thrown of [cancelled, plain]) {
+			const model = new MockLanguageModelV4({
+				doGenerate: () => {
+					throw thrown
+				},
+			})
+			const agent = defineAgent({ name: 'A', model, label: () => 'L' })
+			const ctx = fakeCtx()
+			ctx.run = async () => {
+				throw thrown
+			}
+			expect(await agent.service.ask(ctx, { prompt: 'hi' }).catch((e) => e)).toBe(thrown)
+		}
+	})
+})
+
+describe('defineAgent: request size cap', () => {
+	test('refuses a request over limits.maxInputBytes (413), labelled when a label is set', async () => {
+		const agent = defineAgent({
+			name: 'A',
+			model: okModel(),
+			limits: { maxInputBytes: 50 },
+			label: () => 'L',
+		})
+		const error = await agent.service
+			.ask(fakeCtx(), { prompt: 'x'.repeat(200) })
+			.catch((e) => e)
+		expect(error).toBeInstanceOf(TerminalError)
+		expect(error.code).toBe(413)
+		expect(error.message).toStartWith('[L] ')
+	})
+
+	test('has a generous default', async () => {
+		const agent = defineAgent({ name: 'A', model: okModel() })
+		expect((await agent.service.ask(fakeCtx(), { prompt: 'x'.repeat(10_000) })).text).toBe('ok')
+		const error = await agent.service
+			.ask(fakeCtx(), { prompt: 'x'.repeat(70_000) })
+			.catch((e) => e)
+		expect(error.code).toBe(413)
+	})
+})

package/test/agent/Loop.test.js

@@ -171,3 +171,50 @@ describe('runAgentLoop', () => {
 		expect(JSON.stringify(result.messages)).not.toContain('EARLIER TURN')
 	})
 })
+
+describe('runAgentLoop: step metadata and trace', () => {
+	test('journals small per-step facts, returns a trace and logs one line per step', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'echo', input: { v: 'a' } }]),
+				textResult('done'),
+			),
+		})
+
+		const result = await runAgentLoop(ctx, {
+			model,
+			system: '',
+			prompt: 'go',
+			tools: { echo: echoTool },
+		})
+
+		const meta = ctx.steps[0].result.meta
+		expect(meta).toMatchObject({ model: 'mock-model-id', finishReason: 'tool-calls' })
+		expect(meta.durationMs).toBeGreaterThanOrEqual(0)
+		expect(result.trace).toHaveLength(2)
+		expect(result.trace[0]).toMatchObject({
+			step: 0,
+			tools: ['echo'],
+			finishReason: 'tool-calls',
+		})
+		expect(result.trace[1]).toMatchObject({ step: 1, tools: [], finishReason: 'stop' })
+		expect(ctx.logs).toHaveLength(2)
+		expect(ctx.logs[0]).toContain('llm-step-0')
+		expect(ctx.logs[0]).toContain('tools=[echo]')
+	})
+
+	test('the journaled meta carries no prompt or reply text', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({ doGenerate: textResult('secret reply') })
+		await runAgentLoop(ctx, { model, system: 'secret system', prompt: 'secret prompt' })
+		const meta = JSON.stringify(ctx.steps[0].result.meta)
+		expect(meta).not.toContain('secret')
+	})
+
+	test('works with a context that has no console', async () => {
+		const ctx = { ...fakeCtx(), console: undefined }
+		const model = new MockLanguageModelV4({ doGenerate: textResult('hi') })
+		expect((await runAgentLoop(ctx, { model, system: '', prompt: 'x' })).text).toBe('hi')
+	})
+})

package/test/agent/Model.test.js

@@ -0,0 +1,42 @@
+import { describe, test, expect } from 'bun:test'
+import { TerminalError } from '@restatedev/restate-sdk'
+import { createModel, readLlmEnv } from '../../lib/agent/Model.js'
+
+const ok = { LLM_BASE_URL: 'https://llm.example/v1', MODEL_ID: 'm', LLM_API_KEY: 'sk-secret-value' }
+
+describe('readLlmEnv', () => {
+	test('returns just the three LLM settings', () => {
+		expect(readLlmEnv({ ...ok, PATH: '/bin', OTHER: 'x' })).toEqual(ok)
+	})
+
+	test('the key is optional', () => {
+		const { LLM_API_KEY: _, ...noKey } = ok
+		expect(readLlmEnv(noKey)).toEqual(noKey)
+	})
+
+	test('empty strings count as missing', () => {
+		expect(() => readLlmEnv({ ...ok, MODEL_ID: '' })).toThrow(TerminalError)
+	})
+
+	test('a missing setting is a terminal 500 that names it and never prints a value', () => {
+		const error = (() => {
+			try {
+				readLlmEnv({
+					LLM_BASE_URL: 'https://llm.example/v1',
+					LLM_API_KEY: 'sk-secret-value',
+				})
+			} catch (e) {
+				return e
+			}
+		})()
+		expect(error).toBeInstanceOf(TerminalError)
+		expect(error.code).toBe(500)
+		expect(error.message).toContain('MODEL_ID')
+		expect(error.message).not.toContain('sk-secret-value')
+		expect(error.message).not.toContain('llm.example')
+	})
+
+	test('the result feeds createModel', () => {
+		expect(createModel(readLlmEnv(ok)).modelId).toBe('m')
+	})
+})

package/test/agent/Tools.test.js

@@ -275,3 +275,123 @@ describe('executeTool', () => {
 		})
 	})
 })
+
+describe('tool size caps', () => {
+	const big = (n) => 'x'.repeat(n)
+	const base = { description: 'd', inputSchema: schema }
+
+	test('assertToolSpecs rejects non-positive caps', () => {
+		for (const cap of [0, -1, 1.5, '10']) {
+			expect(() =>
+				assertToolSpecs({ t: { ...base, execute: async () => 1, maxOutputBytes: cap } }),
+			).toThrow(TypeError)
+			expect(() =>
+				assertToolSpecs({ t: { ...base, execute: async () => 1, maxInputBytes: cap } }),
+			).toThrow(TypeError)
+		}
+	})
+
+	test('oversized arguments: the tool is not run and the model is told', async () => {
+		const ctx = fakeCtx()
+		let ran = false
+		const tools = { t: { ...base, execute: async () => ((ran = true), 1), maxInputBytes: 100 } }
+		const output = await executeTool(ctx, call('t', { v: big(500) }), { tools, ...at })
+		expect(ran).toBe(false)
+		expect(ctx.steps).toHaveLength(0)
+		expect(output.type).toBe('error-text')
+		expect(output.value).toContain('over the limit of 100')
+	})
+
+	test('default input cap applies when the tool sets none', async () => {
+		const tools = { t: { ...base, execute: async () => 1 } }
+		const output = await executeTool(fakeCtx(), call('t', { v: big(20_000) }), { tools, ...at })
+		expect(output.type).toBe('error-text')
+	})
+
+	test('oversized result of an execute tool: error to the model, nothing big journaled', async () => {
+		const ctx = fakeCtx()
+		const tools = {
+			t: { ...base, execute: async () => ({ v: big(500) }), maxOutputBytes: 100 },
+		}
+		const output = await executeTool(ctx, call('t'), { tools, ...at })
+		expect(output.type).toBe('error-text')
+		expect(output.value).toContain('over the limit of 100')
+		expect(JSON.stringify(ctx.steps)).not.toContain('xxxxxxxxxx')
+	})
+
+	test('oversized result of a block tool never reaches the model', async () => {
+		const ctx = fakeCtx({ onCall: () => ({ v: big(500) }) })
+		const tools = { t: { ...base, block: 'Svc', handler: 'h', maxOutputBytes: 100 } }
+		const output = await executeTool(ctx, call('t'), { tools, ...at })
+		expect(output.type).toBe('error-text')
+	})
+
+	test('results within the cap pass through', async () => {
+		const tools = { t: { ...base, execute: async () => ({ ok: true }), maxOutputBytes: 100 } }
+		expect(await executeTool(fakeCtx(), call('t'), { tools, ...at })).toEqual({
+			type: 'json',
+			value: { ok: true },
+		})
+	})
+})
+
+describe('block tool key and input hooks', () => {
+	const base = { description: 'd', inputSchema: schema, block: 'Runner', handler: 'run' }
+
+	test('key and input receive the arguments and a per-call idempotency key', async () => {
+		const ctx = fakeCtx({ onCall: () => ({ ok: true }) })
+		const seen = []
+		const tools = {
+			t: {
+				...base,
+				key: (args, meta) => (seen.push(['key', args, meta]), `lock:${args.folder}`),
+				input: (args, meta) => (
+					seen.push(['input', args, meta]),
+					{ ...args, idempotencyKey: meta.idempotencyKey, who: 'caller' }
+				),
+			},
+		}
+
+		await executeTool(ctx, call('t', { folder: '/a' }), { tools, ...at })
+
+		expect(ctx.calls[0].key).toBe('lock:/a')
+		expect(ctx.calls[0].parameter).toEqual({
+			folder: '/a',
+			idempotencyKey: 'inv-1:2.1',
+			who: 'caller',
+		})
+		expect(seen.map(([name]) => name)).toEqual(['key', 'input'])
+		expect(seen[0][2]).toEqual({ idempotencyKey: 'inv-1:2.1' })
+	})
+
+	test('a key built from the idempotency key is different for every call', async () => {
+		const ctx = fakeCtx({ onCall: () => null })
+		const tools = { t: { ...base, key: (_, meta) => `idem:${meta.idempotencyKey}` } }
+		await executeTool(ctx, call('t'), { tools, step: 0, index: 0 })
+		await executeTool(ctx, call('t'), { tools, step: 0, index: 1 })
+		await executeTool(ctx, call('t'), { tools, step: 1, index: 0 })
+		expect(new Set(ctx.calls.map((c) => c.key)).size).toBe(3)
+	})
+
+	test('a key shared by several calls is the same key', async () => {
+		const ctx = fakeCtx({ onCall: () => null })
+		const tools = { t: { ...base, key: (args) => `skill:${args.skill}` } }
+		await executeTool(ctx, call('t', { skill: 'a' }), { tools, step: 0, index: 0 })
+		await executeTool(ctx, call('t', { skill: 'a' }), { tools, step: 1, index: 0 })
+		expect(ctx.calls[0].key).toBe(ctx.calls[1].key)
+	})
+
+	test('without hooks, behaviour is unchanged: arguments are the parameter, no key', async () => {
+		const ctx = fakeCtx({ onCall: () => null })
+		await executeTool(ctx, call('t', { x: 1 }), { tools: { t: base }, ...at })
+		expect(ctx.calls[0]).toMatchObject({ key: undefined, parameter: { x: 1 } })
+	})
+
+	test('assertToolSpecs rejects hooks that are not functions', () => {
+		expect(() => assertToolSpecs({ t: { ...base, key: 'x' } })).toThrow(TypeError)
+		expect(() => assertToolSpecs({ t: { ...base, input: {} } })).toThrow(TypeError)
+		expect(() =>
+			assertToolSpecs({ t: { ...base, key: () => 'k', input: (a) => a } }),
+		).not.toThrow()
+	})
+})

package/test/agent/fakes.js

@@ -1,84 +1,3 @@
-import { TerminalError } from '@restatedev/restate-sdk'
-
-/**
- * A stand-in for Restate's `Context` / `ObjectContext`. Per this repo's testing convention,
- * handlers are unit-tested against a mocked `ctx`, never a live Restate server.
- *
- * `run` mimics the parts of the real one the agent depends on: the result goes through a
- * JSON round trip (as it does through the journal), and when a bounded `ctx.run` gives up
- * it throws a `TerminalError` wrapping the original message. One attempt stands in for
- * "all retries exhausted".
- */
-export function fakeCtx({ key, state = {}, onCall, onApproval } = {}) {
-	const steps = []
-	const calls = []
-	const sends = []
-	return {
-		key,
-		state,
-		steps,
-		calls,
-		sends,
-		run: async (name, fn, options) => {
-			const step = { name, options }
-			steps.push(step)
-			let value
-			try {
-				value = await fn()
-			} catch (error) {
-				if (error instanceof TerminalError || !options?.maxRetryAttempts) throw error
-				throw new TerminalError(error.message)
-			}
-			step.result = value === undefined ? undefined : JSON.parse(JSON.stringify(value))
-			return step.result
-		},
-		request: () => ({ id: 'inv-1' }),
-		genericCall: async (call) => {
-			calls.push(call)
-			return onCall ? onCall(call) : null
-		},
-		genericSend: (send) => {
-			sends.push(send)
-		},
-		awakeable: () => ({
-			id: 'awk-1',
-			promise: { orTimeout: async (timeout) => onApproval?.(timeout) },
-		}),
-		get: async (name) => state[name],
-		set: (name, value) => {
-			state[name] = value
-		},
-		clear: (name) => {
-			delete state[name]
-		},
-	}
-}
-
-const usage = { inputTokens: { total: 1 }, outputTokens: { total: 1 } }
-
-export function textResult(text) {
-	return { content: [{ type: 'text', text }], finishReason: 'stop', usage, warnings: [] }
-}
-
-/** @param {Array<{toolName: string, input: object}>} calls */
-export function toolCallResult(calls) {
-	return {
-		content: calls.map(({ toolName, input }, i) => ({
-			type: 'tool-call',
-			toolCallId: `call_${i}`,
-			toolName,
-			input: JSON.stringify(input),
-		})),
-		finishReason: 'tool-calls',
-		usage,
-		warnings: [],
-	}
-}
-
-/** `doGenerate` that returns `results` one after another. */
-export function sequence(...results) {
-	let i = 0
-	return async () => results[Math.min(i++, results.length - 1)]
-}
+export * from '../../lib/testing/Fakes.js'
 
 export const env = { LLM_BASE_URL: 'https://unused.example', MODEL_ID: 'unused' }

package/testing.js

@@ -0,0 +1 @@
+export * from './lib/testing/Fakes.js'