@empyria/restate 0.4.0 → 0.4.5

package/AGENTS.md

@@ -45,6 +45,9 @@ Journal discipline is enforced by the library, not left to each agent: step meta
 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.

package/README.md

@@ -117,6 +117,13 @@ export const triage = defineAgent({
 | `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

package/lib/agent/Tools.js

@@ -40,7 +40,13 @@ export const jsonBytes = (value) => Buffer.byteLength(JSON.stringify(value) ?? '
  * @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.
@@ -88,6 +94,11 @@ 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`)
@@ -187,14 +198,12 @@ 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 () => {
-						const result =
-							(await spec.execute(call.input, {
-								idempotencyKey: `${ctx.request().id}:${stepId}`,
-							})) ?? null
+						const result = (await spec.execute(call.input, meta)) ?? null
 						assertOutputSize(call.toolName, result, maxOut)
 						return result
 					},
@@ -203,8 +212,8 @@ export async function executeTool(ctx, call, { tools, loadSkill, step, index })
 			: 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,
 				})

package/package.json

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

package/test/agent/Tools.test.js

@@ -334,3 +334,64 @@ describe('tool size caps', () => {
 		})
 	})
 })
+
+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()
+	})
+})