@empyria/restate 0.4.11 → 0.4.13

package/.github/dependabot.yml

@@ -5,7 +5,7 @@
 
 version: 2
 updates:
-  - package-ecosystem: "" # See documentation for possible values
-    directory: "/" # Location of package manifests
-    schedule:
-      interval: "daily"
+    - package-ecosystem: '' # See documentation for possible values
+      directory: '/' # Location of package manifests
+      schedule:
+          interval: 'daily'

package/AGENTS.md

@@ -74,6 +74,11 @@ Attempted to assign to readonly property`). Use `mock.module('@restatedev/restat
 any mocking) if you need to restore it — reading it back off the live `clients` binding later
 may observe an already-mocked value. See `test/Admin.test.js`.
 
+`setupRestate`'s shutdown decision is `shouldDeleteOnShutdown` (unit tested); `keepDeploymentOnShutdown` and
+`forceRegistration` exist because invocations are pinned to their deployment: Restate 1.7 keeps the deployment ID when the
+same address registers again, but only `force: true` updates its handler list. Check both against a live server when
+touching them.
+
 `setupRestate` (in `lib/Admin.js`) is not unit tested: it binds real HTTP/2 + health-check
 ports, installs process-wide `SIGTERM`/`SIGINT` handlers, and its returned `forceClose` calls
 `process.exit()` — none of which are safe to exercise in-process in a test run.

package/README.md

@@ -124,7 +124,15 @@ Cron expressions are `croner` patterns (five fields, or six with seconds for tes
 
 ### An endpoint that does not register itself
 
-`setupRestate` serves, registers the deployment, and deletes it again on SIGINT/SIGTERM. When something else owns
+`setupRestate` serves, registers the deployment, and deletes it again on SIGINT/SIGTERM. Invocations in flight are
+pinned to their deployment, so deleting it strands them: a restarted process cannot finish them and they can only be
+killed. For a process restarted at the same address, pass `keepDeploymentOnShutdown: true` and `forceRegistration: true`:
+Restate keeps the deployment ID when the same address registers again (checked against Restate 1.7), running
+invocations resume on the new process, and `force` updates the services and handlers in place (a plain
+re-registration keeps the old handler list). Only compatible changes are safe this way: an invocation of a removed or
+renamed handler cannot resume.
+
+When something else owns
 registration (a supervisor that registers a process's address once and restarts the process), use
 `serveEndpoint({port, services})` instead: it only serves. Invocations already running on a deployment are pinned to its
 address, so the deployment must survive the process stopping and come back at the same address.
@@ -133,12 +141,13 @@ address, so the deployment must survive the process stopping and come back at th
 
 `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).                                                                                                                                                     |
+| 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).                                                                                                                            |
+| `output` as a function   | `output: (input) => schema \| undefined` chooses the answer's JSON Schema per request (same rules as `tools`: pure, no I/O); `undefined` means a free-text answer for that request. With any `output`, the model's answer is checked against the schema: a wrong one is handed back to the model to correct, and a model that never gets it right ends at `maxSteps`. |
+| `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

package/lib/Admin.js

@@ -581,7 +581,15 @@ export function defineObject(def) {
 /**
  * Starts the HTTP/2 Restate endpoint, a health-check HTTP server, registers
  * the deployment with the admin API, and wires SIGTERM/SIGINT for graceful
- * shutdown (deregistering the deployment before exiting).
+ * shutdown (deregistering the deployment before exiting, unless `keepDeploymentOnShutdown`).
+ *
+ * Invocations in flight are pinned to their deployment. Deleting it on shutdown means a restarted process cannot
+ * finish them, and they cannot even be cancelled (only killed). For a process that is restarted at the same address
+ * (a dev loop, a service that is redeployed in place), pass `keepDeploymentOnShutdown: true` and
+ * `forceRegistration: true`: Restate keeps the deployment ID when the same address registers again, so running
+ * invocations resume on the new process, and `force` updates the deployment's services and handlers in place (a plain
+ * re-registration keeps the old handler list). Only compatible changes are safe that way: an invocation of a handler
+ * that was removed or renamed cannot resume.
  * @param {Object} params
  * @param {string} params.restateAdminURL
  * @param {string} params.host
@@ -590,6 +598,9 @@ export function defineObject(def) {
  * @param {Parameters<typeof restate.createEndpointHandler>[0]['services']} [params.services]
  * @param {() => void} [params.shutdownCallbackFn]
  * @param {(deploymentId: string) => void} [params.registrationCallbackFn]
+ * @param {boolean} [params.keepDeploymentOnShutdown=false] Leave the deployment registered on SIGINT/SIGTERM.
+ * @param {boolean} [params.forceRegistration=false] Register with `force: true`, updating an existing deployment at
+ *   the same address in place.
  * @returns {Promise<RestateServer>}
  */
 export async function setupRestate({
@@ -600,6 +611,8 @@ export async function setupRestate({
 	services = [],
 	shutdownCallbackFn,
 	registrationCallbackFn,
+	keepDeploymentOnShutdown = false,
+	forceRegistration = false,
 }) {
 	const handler = restate.createEndpointHandler({ services })
 
@@ -647,8 +660,12 @@ export async function setupRestate({
 		isHealthy = false
 
 		try {
-			if (deploymentId && deploymentId !== 'NA') {
+			if (shouldDeleteOnShutdown({ deploymentId, keepDeploymentOnShutdown })) {
 				await deleteDeployment(restateAdminURL, deploymentId)
+			} else if (keepDeploymentOnShutdown && deploymentId) {
+				console.log(
+					`Keeping deployment ${deploymentId}: invocations in flight resume when it comes back`,
+				)
 			}
 		} catch (err) {
 			console.error(err)
@@ -688,7 +705,11 @@ export async function setupRestate({
 
 			for (let attempt = 1; attempt <= attempts; attempt++) {
 				try {
-					const result = await registerDeployment({ restateAdminURL, serviceURL })
+					const result = await registerDeployment({
+						restateAdminURL,
+						serviceURL,
+						force: forceRegistration,
+					})
 
 					try {
 						registrationCallbackFn?.(result.id)
@@ -751,15 +772,27 @@ export async function setupRestate({
 }
 
 /**
- * POSTs this process's HTTP endpoint to the Restate admin API so its services are discoverable.
- * @param {{restateAdminURL: string, serviceURL: string}} params
+ * Whether `setupRestate`'s shutdown removes the deployment: only a real one (registered, not `'NA'`), and not when the
+ * caller keeps it for the next start.
+ * @param {{deploymentId: string|undefined, keepDeploymentOnShutdown?: boolean}} params
+ * @returns {boolean}
+ */
+export function shouldDeleteOnShutdown({ deploymentId, keepDeploymentOnShutdown = false }) {
+	return Boolean(deploymentId) && deploymentId !== 'NA' && !keepDeploymentOnShutdown
+}
+
+/**
+ * POSTs this process's HTTP endpoint to the Restate admin API so its services are discoverable. Registering an
+ * address that is already registered returns the same deployment; with `force` its services and handlers are
+ * updated in place (without it, Restate keeps the ones it knew).
+ * @param {{restateAdminURL: string, serviceURL: string, force?: boolean}} params
  * @returns {Promise<RegisterDeploymentResponse>}
  */
-export async function registerDeployment({ restateAdminURL, serviceURL }) {
+export async function registerDeployment({ restateAdminURL, serviceURL, force = false }) {
 	const uri = serviceURL
 
 	/** @type {RegisterDeploymentRequest} */
-	const body = { uri, force: null }
+	const body = { uri, force: force ? true : null }
 
 	console.log(`Registering: ${uri} at: ${restateAdminURL}/deployments`)
 
@@ -872,7 +905,8 @@ export function createRestateAdmin({ restateAdminURL, restateURL }) {
 		pollInvocation: ({ invocationId, intervalMs }) =>
 			pollInvocation({ restateURL, invocationId, intervalMs }),
 
-		getInvocationOutput: ({ invocationId }) => getInvocationOutput({ restateURL, invocationId }),
+		getInvocationOutput: ({ invocationId }) =>
+			getInvocationOutput({ restateURL, invocationId }),
 
 		sendMessage: ({ name, message, key, payload, idempotencyKey }) =>
 			sendMessage({ restateURL, name, message, key, payload, idempotencyKey }),

package/lib/agent/Agent.js

@@ -68,9 +68,11 @@ export const AskOutputSchema = {
  * @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 {object | ((input: any) => object | undefined)} [output] Output JSON Schema. When given, the model
+ *   must answer through the `final_answer` tool; its arguments are checked against the schema (a wrong answer is
+ *   handed back to the model to correct) and `ask` returns them. A function chooses the schema per request, like
+ *   `tools`: it runs on every execution and replay, so it must be pure and do no I/O; `undefined` means a free-text
+ *   answer for that request. 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}`.
@@ -162,7 +164,12 @@ export function defineAgent(spec) {
 		maxInputBytes = DEFAULT_MAX_INPUT_BYTES,
 	} = limits
 
-	const outputSchema = respond ? (responseSchema ?? {}) : (output ?? AskOutputSchema)
+	// A per-request `output` cannot be known here; the loop checks the answer against the schema it was given.
+	const outputSchema = respond
+		? (responseSchema ?? {})
+		: typeof output === 'function'
+			? {}
+			: (output ?? AskOutputSchema)
 
 	const ask = withValidation(input, outputSchema, async (ctx, validatedInput) => {
 		const inputBytes = jsonBytes(validatedInput)
@@ -193,6 +200,8 @@ export function defineAgent(spec) {
 		const requestTools =
 			typeof tools === 'function' ? await resolveTools(name, tools, validatedInput) : tools
 
+		const requestOutput = typeof output === 'function' ? output(validatedInput) : output
+
 		const context = loadContext
 			? await ctx.run('load-context', async () => {
 					const { systemPrompt = '', skills = [] } =
@@ -217,7 +226,7 @@ export function defineAgent(spec) {
 			skills: context.skills,
 			loadSkill: skillLoader({ loadSkill, loadContext, input: validatedInput }),
 			loadHistory: session?.load,
-			outputSchema: output,
+			outputSchema: requestOutput,
 			maxSteps,
 			maxTokens,
 			retry: llmRetry,
@@ -226,7 +235,7 @@ 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 }
+		return requestOutput ? result.output : { text: result.text, steps: result.steps }
 	}
 
 	const options = { retryPolicy }

package/lib/agent/Loop.js

@@ -1,11 +1,16 @@
 import { generateText } from 'ai'
 import { TerminalError } from '@restatedev/restate-sdk'
+import { createValidator } from '@empyria/common'
 import { DEFAULT_LLM_RETRY, toRestateLLMError } from './Errors.js'
 import { cachedSkillFetcher, expandSkillRefs } from './Skills.js'
 import { FINAL_ANSWER_TOOL, buildToolDefs, executeTool } from './Tools.js'
 
 const FINAL_ANSWER_NUDGE = `Respond by calling the ${FINAL_ANSWER_TOOL} tool with your answer.`
 
+/** The text handed back to the model when its final answer does not fit the schema (kept short). */
+const rejection = (message) =>
+	`Your answer was not accepted because it does not match the required format: ${String(message).slice(0, 1500)}\nCall ${FINAL_ANSWER_TOOL} again with a corrected answer.`
+
 /**
  * Drives an AI SDK Core tool-calling loop ONE round at a time, from inside a Restate
  * handler — deliberately NOT `generateText`'s own multi-step loop wrapped in a single
@@ -36,7 +41,8 @@ const FINAL_ANSWER_NUDGE = `Respond by calling the ${FINAL_ANSWER_TOOL} tool wit
  *   retry?: import('@restatedev/restate-sdk').RunOptions<any>,
  * }} 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`.
+ *   `final_answer` tool and the loop returns that tool's input as `output`, after checking it against the
+ *   schema: an answer that does not fit is handed back to the model to correct, not accepted.
  * @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
@@ -72,6 +78,7 @@ export async function runAgentLoop(
 	const messages = [{ role: 'user', content: prompt }]
 	let totalTokens = 0
 	const fetchSkill = cachedSkillFetcher(loadSkill)
+	const validateOutput = outputSchema ? createValidator(outputSchema) : undefined
 	const trace = []
 
 	for (let step = 0; step < maxSteps; step++) {
@@ -149,14 +156,26 @@ export async function runAgentLoop(
 
 		for (const [index, call] of result.toolCalls.entries()) {
 			if (outputSchema && call.toolName === FINAL_ANSWER_TOOL) {
-				return {
-					text: result.text,
-					output: call.input,
-					steps: step + 1,
-					totalTokens,
-					trace,
-					messages,
+				let output
+				try {
+					output = validateOutput(call.input)
+				} catch (error) {
+					// The model's arguments are not checked by the provider. Tell it what is wrong and let it
+					// answer again; this uses a step, so a model that cannot get it right ends at maxSteps.
+					messages.push({
+						role: 'tool',
+						content: [
+							{
+								type: 'tool-result',
+								toolCallId: call.toolCallId,
+								toolName: call.toolName,
+								output: { type: 'error-text', value: rejection(error.message) },
+							},
+						],
+					})
+					continue
 				}
+				return { text: result.text, output, steps: step + 1, totalTokens, trace, messages }
 			}
 			const output = await executeTool(ctx, call, { tools, loadSkill, step, index })
 			messages.push({

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@empyria/restate",
-	"version": "0.4.11",
+	"version": "0.4.13",
 	"description": "Restate.dev helpers for the Empyria nanoservice framework",
 	"license": "MIT",
 	"author": "Imre Fazekas <imre.fazekas@icloud.com>",

package/test/Admin.test.js

@@ -20,6 +20,7 @@ import {
 	sendMessageWithDiscovery,
 	sendMessageAsyncWithDiscovery,
 	registerDeployment,
+	shouldDeleteOnShutdown,
 	sendMessageAsynSafe,
 	defineWorkflow,
 	defineService,
@@ -521,6 +522,42 @@ describe('registerDeployment', () => {
 			registerDeployment({ restateAdminURL: 'http://admin', serviceURL: 'http://svc' }),
 		).rejects.toThrow(/bad uri/)
 	})
+
+	test('sends force: null by default and force: true when asked, to the deployments endpoint', async () => {
+		const sent = []
+		globalThis.fetch = mock(async (url, init) => {
+			sent.push({ url, body: JSON.parse(init.body) })
+			return jsonResponse({ id: 'dep-1', services: [] })
+		})
+		await registerDeployment({ restateAdminURL: 'http://admin', serviceURL: 'http://svc:9080' })
+		await registerDeployment({
+			restateAdminURL: 'http://admin',
+			serviceURL: 'http://svc:9080',
+			force: true,
+		})
+		expect(sent).toEqual([
+			{ url: 'http://admin/deployments', body: { uri: 'http://svc:9080', force: null } },
+			{ url: 'http://admin/deployments', body: { uri: 'http://svc:9080', force: true } },
+		])
+	})
+})
+
+describe('shouldDeleteOnShutdown', () => {
+	test('deletes a registered deployment by default', () => {
+		expect(shouldDeleteOnShutdown({ deploymentId: 'dep-1' })).toBe(true)
+	})
+
+	test('keeps it when asked, so invocations in flight can resume after a restart', () => {
+		expect(
+			shouldDeleteOnShutdown({ deploymentId: 'dep-1', keepDeploymentOnShutdown: true }),
+		).toBe(false)
+	})
+
+	test('never deletes without a real deployment', () => {
+		for (const deploymentId of [undefined, '', 'NA']) {
+			expect(shouldDeleteOnShutdown({ deploymentId })).toBe(false)
+		}
+	})
 })
 
 // `clients` is a live ESM namespace binding (read-only — Object.assign/property

package/test/agent/Agent.test.js

@@ -111,6 +111,50 @@ describe('defineAgent: ask', () => {
 		await expect(answer({ score: 'high' })).rejects.toThrow(TerminalError)
 	})
 
+	test('output can be chosen per request; undefined means a free-text answer', async () => {
+		const schemas = {
+			score: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+			name: { type: 'object', properties: { name: { type: 'string' } }, required: ['name'] },
+		}
+		const agent = (model) =>
+			defineAgent({
+				name: 'A',
+				model,
+				input: { type: 'object', properties: { want: { type: 'string' } } },
+				prompt: () => 'go',
+				output: ({ want }) => schemas[want],
+			})
+		const answering = (input) =>
+			new MockLanguageModelV4({
+				doGenerate: toolCallResult([{ toolName: 'final_answer', input }]),
+			})
+
+		expect(
+			await agent(answering({ score: 3 })).service.ask(fakeCtx(), { want: 'score' }),
+		).toEqual({ score: 3 })
+		expect(
+			await agent(answering({ name: 'x' })).service.ask(fakeCtx(), { want: 'name' }),
+		).toEqual({
+			name: 'x',
+		})
+		// A wrong answer for the request's schema is corrected, not accepted.
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'final_answer', input: { name: 5 } }]),
+				toolCallResult([{ toolName: 'final_answer', input: { name: 'ok' } }]),
+			),
+		})
+		expect(await agent(model).service.ask(fakeCtx(), { want: 'name' })).toEqual({ name: 'ok' })
+		// No schema for this request: the plain answer.
+		expect(await agent(okModel()).service.ask(fakeCtx(), { want: 'none' })).toMatchObject({
+			steps: 1,
+		})
+	})
+
 	test('a custom input schema with a prompt builder', async () => {
 		const model = okModel()
 		const agent = defineAgent({

package/test/agent/Loop.test.js

@@ -154,6 +154,51 @@ describe('runAgentLoop', () => {
 		expect(JSON.stringify(model.doGenerateCalls[1].prompt)).toContain('final_answer')
 	})
 
+	test('an answer that does not fit the schema goes back to the model, which corrects it', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'final_answer', input: { score: 'high' } }]),
+				toolCallResult([{ toolName: 'final_answer', input: { score: 7 } }]),
+			),
+		})
+
+		const result = await runAgentLoop(fakeCtx(), {
+			model,
+			system: '',
+			prompt: 'rate it',
+			outputSchema: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+		})
+
+		expect(result).toMatchObject({ output: { score: 7 }, steps: 2 })
+		const second = JSON.stringify(model.doGenerateCalls[1].prompt)
+		expect(second).toContain('does not match the required format')
+		expect(second).toContain('final_answer')
+	})
+
+	test('a model that never gets the answer right ends at maxSteps, terminally', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: toolCallResult([{ toolName: 'final_answer', input: { score: 'high' } }]),
+		})
+		const caught = await runAgentLoop(fakeCtx(), {
+			model,
+			system: '',
+			prompt: 'rate it',
+			maxSteps: 3,
+			outputSchema: {
+				type: 'object',
+				properties: { score: { type: 'number' } },
+				required: ['score'],
+			},
+		}).catch((e) => e)
+		expect(caught).toBeInstanceOf(TerminalError)
+		expect(caught.message).toContain('maxSteps')
+		expect(model.doGenerateCalls).toHaveLength(3)
+	})
+
 	test('prepends loadHistory’s messages to every model call, without journaling them', async () => {
 		const ctx = fakeCtx()
 		const model = new MockLanguageModelV4({ doGenerate: textResult('ok') })