@empyria/restate 0.1.23 → 0.2.0

package/test/agent/Loop.test.js

@@ -0,0 +1,173 @@
+import { describe, test, expect } from 'bun:test'
+import { MockLanguageModelV4 } from 'ai/test'
+import { TerminalError, RetryableError } from '@restatedev/restate-sdk'
+import { APICallError } from '@ai-sdk/provider'
+import { runAgentLoop } from '../../lib/agent/Loop.js'
+import { DEFAULT_LLM_RETRY } from '../../lib/agent/Errors.js'
+import { fakeCtx, textResult, toolCallResult, sequence } from './fakes.js'
+
+const echoTool = {
+	description: 'Echo the input',
+	inputSchema: { type: 'object', properties: { v: { type: 'string' } } },
+	execute: async ({ v }) => ({ echoed: v }),
+}
+
+const apiError = (statusCode, isRetryable, responseHeaders) =>
+	new APICallError({
+		message: `status ${statusCode}`,
+		url: 'https://llm.example/v1/chat/completions',
+		requestBodyValues: {},
+		statusCode,
+		isRetryable,
+		responseHeaders,
+	})
+
+describe('runAgentLoop', () => {
+	test('returns the text answer and records each LLM call as its own step', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({ doGenerate: textResult('hi') })
+
+		const result = await runAgentLoop(ctx, { model, system: '', prompt: 'hello' })
+
+		expect(result).toMatchObject({ text: 'hi', steps: 1 })
+		expect(result.messages[0]).toEqual({ role: 'user', content: 'hello' })
+		expect(ctx.steps.map((s) => s.name)).toEqual(['llm-step-0'])
+	})
+
+	test('applies DEFAULT_LLM_RETRY by default, and a caller override', async () => {
+		const model = new MockLanguageModelV4({ doGenerate: textResult('hi') })
+		const ctx = fakeCtx()
+		await runAgentLoop(ctx, { model, system: '', prompt: 'hello' })
+		expect(ctx.steps[0].options).toBe(DEFAULT_LLM_RETRY)
+
+		const retry = { maxRetryAttempts: 1 }
+		const ctx2 = fakeCtx()
+		await runAgentLoop(ctx2, { model, system: '', prompt: 'hello', retry })
+		expect(ctx2.steps[0].options).toBe(retry)
+	})
+
+	test('a non-retryable model failure is terminal', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: () => {
+				throw apiError(401, false)
+			},
+		})
+		await expect(
+			runAgentLoop(fakeCtx(), { model, system: '', prompt: 'hello' }),
+		).rejects.toThrow(TerminalError)
+	})
+
+	test('a rate-limited model failure with Retry-After becomes a RetryableError', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: () => {
+				throw apiError(429, true, { 'retry-after': '5' })
+			},
+		})
+		let caught
+		try {
+			await runAgentLoop(fakeCtx(), { model, system: '', prompt: 'hello', retry: {} })
+		} catch (error) {
+			caught = error
+		}
+		expect(caught).toBeInstanceOf(RetryableError)
+		expect(caught.retryAfter).toEqual({ seconds: 5 })
+	})
+
+	test('runs a tool, feeds its result back, and continues', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'echo', input: { v: 'x' } }]),
+				textResult('done'),
+			),
+		})
+
+		const result = await runAgentLoop(ctx, {
+			model,
+			system: '',
+			prompt: 'hello',
+			tools: { echo: echoTool },
+		})
+
+		expect(result).toMatchObject({ text: 'done', steps: 2 })
+		expect(ctx.steps.map((s) => s.name)).toEqual(['llm-step-0', 'tool:echo-0.0', 'llm-step-1'])
+		expect(JSON.stringify(model.doGenerateCalls[1].prompt)).toContain('echoed')
+	})
+
+	test('exceeding maxSteps is terminal', async () => {
+		const model = new MockLanguageModelV4({
+			doGenerate: toolCallResult([{ toolName: 'echo', input: { v: 'x' } }]),
+		})
+		await expect(
+			runAgentLoop(fakeCtx(), {
+				model,
+				system: '',
+				prompt: 'hello',
+				tools: { echo: echoTool },
+				maxSteps: 1,
+			}),
+		).rejects.toThrow(/maxSteps/)
+	})
+
+	test('exceeding maxTokens is terminal with code 429', async () => {
+		const model = new MockLanguageModelV4({ doGenerate: textResult('hi') })
+		let caught
+		try {
+			await runAgentLoop(fakeCtx(), { model, system: '', prompt: 'hello', maxTokens: 1 })
+		} catch (error) {
+			caught = error
+		}
+		expect(caught).toBeInstanceOf(TerminalError)
+		expect(caught.code).toBe(429)
+	})
+
+	test('puts only the skill index in the system prompt', async () => {
+		const model = new MockLanguageModelV4({ doGenerate: textResult('ok') })
+		await runAgentLoop(fakeCtx(), {
+			model,
+			system: 'Be terse.',
+			prompt: 'hello',
+			skills: [{ name: 'summarise', description: 'Summarise things' }],
+		})
+		const sent = JSON.stringify(model.doGenerateCalls[0].prompt)
+		expect(sent).toContain('Be terse.')
+		expect(sent).toContain('summarise: Summarise things')
+	})
+
+	test('with an outputSchema, nudges a plain-text answer and returns final_answer’s input', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				textResult('here you go'),
+				toolCallResult([{ toolName: 'final_answer', input: { score: 7 } }]),
+			),
+		})
+
+		const result = await runAgentLoop(ctx, {
+			model,
+			system: '',
+			prompt: 'rate it',
+			outputSchema: { type: 'object', properties: { score: { type: 'number' } } },
+		})
+
+		expect(result).toMatchObject({ output: { score: 7 }, steps: 2 })
+		expect(JSON.stringify(model.doGenerateCalls[1].prompt)).toContain('final_answer')
+	})
+
+	test('prepends loadHistory’s messages to every model call, without journaling them', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({ doGenerate: textResult('ok') })
+		const history = [{ role: 'user', content: 'EARLIER TURN' }]
+
+		const result = await runAgentLoop(ctx, {
+			model,
+			system: '',
+			prompt: 'hello',
+			loadHistory: async () => history,
+		})
+
+		expect(JSON.stringify(model.doGenerateCalls[0].prompt)).toContain('EARLIER TURN')
+		expect(JSON.stringify(ctx.steps[0].result)).not.toContain('EARLIER TURN')
+		expect(JSON.stringify(result.messages)).not.toContain('EARLIER TURN')
+	})
+})

package/test/agent/Skills.test.js

@@ -0,0 +1,147 @@
+import { describe, test, expect } from 'bun:test'
+import { MockLanguageModelV4 } from 'ai/test'
+import { TerminalError } from '@restatedev/restate-sdk'
+import {
+	cachedSkillFetcher,
+	expandSkillRefs,
+	hashSkill,
+	skillRefOutput,
+} from '../../lib/agent/Skills.js'
+import { runAgentLoop } from '../../lib/agent/Loop.js'
+import { fakeCtx, sequence, textResult, toolCallResult } from './fakes.js'
+
+const skillMessage = (body) => ({
+	role: 'tool',
+	content: [
+		{
+			type: 'tool-result',
+			toolCallId: 'call_0',
+			toolName: 'load_skill',
+			output: skillRefOutput({ name: 'summarise', sha256: hashSkill(body) }),
+		},
+	],
+})
+const fetchFrom = (skills) => async (name) => skills[name]
+
+describe('expandSkillRefs', () => {
+	test('replaces a reference with the body when the hash matches', async () => {
+		const [expanded] = await expandSkillRefs([skillMessage('v1')], {
+			fetchSkill: fetchFrom({ summarise: 'v1' }),
+			strict: true,
+		})
+		expect(expanded.content[0].output).toEqual({ type: 'text', value: 'v1' })
+	})
+
+	test('leaves other messages and tool results untouched', async () => {
+		const messages = [
+			{ role: 'user', content: 'hi' },
+			{
+				role: 'tool',
+				content: [
+					{
+						type: 'tool-result',
+						toolCallId: 'c',
+						toolName: 'echo',
+						output: { type: 'json', value: 1 },
+					},
+				],
+			},
+		]
+		expect(
+			await expandSkillRefs(messages, { fetchSkill: fetchFrom({}), strict: true }),
+		).toEqual(messages)
+	})
+
+	test('strict: a changed or missing skill is terminal', async () => {
+		for (const skills of [{ summarise: 'v2' }, {}]) {
+			await expect(
+				expandSkillRefs([skillMessage('v1')], {
+					fetchSkill: fetchFrom(skills),
+					strict: true,
+				}),
+			).rejects.toBeInstanceOf(TerminalError)
+		}
+	})
+
+	test('lenient: a changed skill uses the current body; a missing one becomes a notice', async () => {
+		const [changed] = await expandSkillRefs([skillMessage('v1')], {
+			fetchSkill: fetchFrom({ summarise: 'v2' }),
+			strict: false,
+		})
+		expect(changed.content[0].output).toEqual({ type: 'text', value: 'v2' })
+
+		const [missing] = await expandSkillRefs([skillMessage('v1')], {
+			fetchSkill: undefined,
+			strict: false,
+		})
+		expect(missing.content[0].output).toEqual({
+			type: 'error-text',
+			value: 'Skill "summarise" is no longer available.',
+		})
+	})
+})
+
+describe('cachedSkillFetcher', () => {
+	test('fetches each skill once per execution', async () => {
+		let fetches = 0
+		const fetchSkill = cachedSkillFetcher(async () => {
+			fetches++
+			return 'body'
+		})
+		await fetchSkill('a')
+		await fetchSkill('a')
+		await fetchSkill('b')
+		expect(fetches).toBe(2)
+	})
+
+	test('is undefined without a loader', () => {
+		expect(cachedSkillFetcher(undefined)).toBeUndefined()
+	})
+})
+
+describe('runAgentLoop with skills', () => {
+	const skills = [{ name: 'summarise', description: 'Summarise' }]
+
+	test('the model sees the body; the journal holds only the reference', async () => {
+		const ctx = fakeCtx()
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'load_skill', input: { name: 'summarise' } }]),
+				textResult('done'),
+			),
+		})
+
+		const result = await runAgentLoop(ctx, {
+			model,
+			system: '',
+			prompt: 'hi',
+			skills,
+			loadSkill: async () => 'FULL SKILL BODY',
+		})
+
+		expect(JSON.stringify(model.doGenerateCalls[1].prompt)).toContain('FULL SKILL BODY')
+		expect(JSON.stringify(ctx.steps.map((s) => s.result))).not.toContain('FULL SKILL BODY')
+		expect(JSON.stringify(result.messages)).not.toContain('FULL SKILL BODY')
+	})
+
+	test('a skill that changes mid-invocation fails the step instead of changing the instructions', async () => {
+		let version = 0
+		const model = new MockLanguageModelV4({
+			doGenerate: sequence(
+				toolCallResult([{ toolName: 'load_skill', input: { name: 'summarise' } }]),
+				textResult('done'),
+			),
+		})
+
+		await expect(
+			runAgentLoop(fakeCtx(), {
+				model,
+				system: '',
+				prompt: 'hi',
+				skills,
+				// The journaled hash is taken from v1; the next LLM step fetches v2.
+				loadSkill: async () => `v${++version}`,
+			}),
+		).rejects.toThrow(/changed or disappeared/)
+	})
+})

package/test/agent/Tools.test.js

@@ -0,0 +1,277 @@
+import { describe, test, expect } from 'bun:test'
+import { CancelledError, TerminalError, TimeoutError, serde } from '@restatedev/restate-sdk'
+import { assertToolSpecs, executeTool } from '../../lib/agent/Tools.js'
+import { DEFAULT_TOOL_RETRY } from '../../lib/agent/Errors.js'
+import { hashSkill } from '../../lib/agent/Skills.js'
+import { fakeCtx } from './fakes.js'
+
+const schema = { type: 'object' }
+const call = (toolName, input = {}) => ({ toolCallId: 'call_0', toolName, input })
+const at = { step: 2, index: 1 }
+
+describe('assertToolSpecs', () => {
+	test('accepts block and execute tools', () => {
+		expect(() =>
+			assertToolSpecs({
+				a: { description: 'a', inputSchema: schema, block: 'Svc', handler: 'do' },
+				b: { description: 'b', inputSchema: schema, execute: async () => 1 },
+			}),
+		).not.toThrow()
+	})
+
+	test.each([
+		[
+			'a reserved name',
+			{ final_answer: { description: 'x', inputSchema: schema, execute() {} } },
+		],
+		['an invalid name', { 'a b': { description: 'x', inputSchema: schema, execute() {} } }],
+		['no inputSchema', { a: { description: 'x', execute() {} } }],
+		[
+			'both block and execute',
+			{
+				a: {
+					description: 'x',
+					inputSchema: schema,
+					block: 'S',
+					handler: 'h',
+					execute() {},
+				},
+			},
+		],
+		['a block without handler', { a: { description: 'x', inputSchema: schema, block: 'S' } }],
+		[
+			'approval without notify',
+			{
+				a: {
+					description: 'x',
+					inputSchema: schema,
+					execute() {},
+					approval: { timeout: 10 },
+				},
+			},
+		],
+	])('rejects %s', (_, tools) => {
+		expect(() => assertToolSpecs(tools)).toThrow(TypeError)
+	})
+})
+
+describe('executeTool', () => {
+	test('execute tool: own bounded step, with an idempotency key', async () => {
+		const ctx = fakeCtx()
+		let meta
+		const tools = {
+			write: {
+				description: 'w',
+				inputSchema: schema,
+				execute: async (args, m) => {
+					meta = m
+					return { wrote: args.v }
+				},
+			},
+		}
+
+		const output = await executeTool(ctx, call('write', { v: 1 }), { tools, ...at })
+
+		expect(output).toEqual({ type: 'json', value: { wrote: 1 } })
+		expect(meta).toEqual({ idempotencyKey: 'inv-1:2.1' })
+		expect(ctx.steps[0]).toMatchObject({ name: 'tool:write-2.1', options: DEFAULT_TOOL_RETRY })
+	})
+
+	test('execute tool: exhausted retries reach the model as error-text', async () => {
+		const tools = {
+			flaky: {
+				description: 'f',
+				inputSchema: schema,
+				execute: async () => {
+					throw new Error('connection reset')
+				},
+			},
+		}
+		const output = await executeTool(fakeCtx(), call('flaky'), { tools, ...at })
+		expect(output).toEqual({ type: 'error-text', value: 'connection reset' })
+	})
+
+	test('execute tool without a retry limit: a transient error propagates, so Restate retries', async () => {
+		const tools = {
+			flaky: {
+				description: 'f',
+				inputSchema: schema,
+				retry: {},
+				execute: async () => {
+					throw new Error('connection reset')
+				},
+			},
+		}
+		await expect(executeTool(fakeCtx(), call('flaky'), { tools, ...at })).rejects.toThrow(
+			'connection reset',
+		)
+	})
+
+	test('block tool: a Restate call to the target handler', async () => {
+		const ctx = fakeCtx({ onCall: () => ({ id: 'ISS-1' }) })
+		const tools = {
+			lookup: {
+				description: 'l',
+				inputSchema: schema,
+				block: { name: 'IssuerService' },
+				handler: 'get',
+				key: (args) => args.isin,
+			},
+		}
+
+		const output = await executeTool(ctx, call('lookup', { isin: 'CH01' }), { tools, ...at })
+
+		expect(output).toEqual({ type: 'json', value: { id: 'ISS-1' } })
+		expect(ctx.calls[0]).toEqual({
+			service: 'IssuerService',
+			method: 'get',
+			key: 'CH01',
+			parameter: { isin: 'CH01' },
+			inputSerde: serde.json,
+			outputSerde: serde.json,
+		})
+		expect(ctx.steps).toEqual([])
+	})
+
+	test('block tool: a terminal error from the callee reaches the model as error-text', async () => {
+		const ctx = fakeCtx({
+			onCall: () => {
+				throw new TerminalError('issuer not found')
+			},
+		})
+		const tools = {
+			lookup: {
+				description: 'l',
+				inputSchema: schema,
+				block: 'IssuerService',
+				handler: 'get',
+			},
+		}
+		expect(await executeTool(ctx, call('lookup'), { tools, ...at })).toEqual({
+			type: 'error-text',
+			value: 'issuer not found',
+		})
+	})
+
+	test('cancellation always propagates', async () => {
+		const ctx = fakeCtx({
+			onCall: () => {
+				throw new CancelledError()
+			},
+		})
+		const tools = {
+			lookup: {
+				description: 'l',
+				inputSchema: schema,
+				block: 'IssuerService',
+				handler: 'get',
+			},
+		}
+		await expect(executeTool(ctx, call('lookup'), { tools, ...at })).rejects.toBeInstanceOf(
+			CancelledError,
+		)
+	})
+
+	test('unknown tool reaches the model as error-text', async () => {
+		expect(await executeTool(fakeCtx(), call('nope'), { tools: {}, ...at })).toEqual({
+			type: 'error-text',
+			value: 'Unknown tool "nope".',
+		})
+	})
+
+	test('load_skill journals only a reference and hash, never the body', async () => {
+		const ctx = fakeCtx()
+		const loadSkill = async (name) => (name === 'summarise' ? 'FULL BODY' : undefined)
+		const skillRef = { name: 'summarise', sha256: hashSkill('FULL BODY') }
+
+		expect(
+			await executeTool(ctx, call('load_skill', { name: 'summarise' }), {
+				tools: {},
+				loadSkill,
+				...at,
+			}),
+		).toEqual({ type: 'json', value: { skillRef } })
+		expect(ctx.steps[0]).toMatchObject({ name: 'skill:summarise-2.1', result: skillRef })
+
+		expect(
+			await executeTool(ctx, call('load_skill', { name: 'other' }), {
+				tools: {},
+				loadSkill,
+				...at,
+			}),
+		).toEqual({ type: 'error-text', value: 'No skill named "other" is available.' })
+	})
+
+	describe('approval', () => {
+		const approvalTools = (executed) => ({
+			notifyDesk: {
+				description: 'n',
+				inputSchema: schema,
+				execute: async () => {
+					executed.push(true)
+					return 'sent'
+				},
+				approval: { timeout: 1000, notify: { block: 'ApprovalInbox', handler: 'request' } },
+			},
+		})
+
+		test('notifies the approver with the awakeable ID, then runs the tool once approved', async () => {
+			const executed = []
+			let waited
+			const ctx = fakeCtx({
+				onApproval: (timeout) => {
+					waited = timeout
+					return { approved: true }
+				},
+			})
+
+			const output = await executeTool(ctx, call('notifyDesk', { msg: 'hi' }), {
+				tools: approvalTools(executed),
+				...at,
+			})
+
+			expect(output).toEqual({ type: 'json', value: 'sent' })
+			expect(executed).toHaveLength(1)
+			expect(waited).toBe(1000)
+			expect(ctx.sends[0]).toMatchObject({
+				service: 'ApprovalInbox',
+				method: 'request',
+				parameter: {
+					awakeableId: 'awk-1',
+					tool: 'notifyDesk',
+					input: { msg: 'hi' },
+					invocationId: 'inv-1',
+				},
+			})
+		})
+
+		test('a rejection skips the tool and tells the model why', async () => {
+			const executed = []
+			const ctx = fakeCtx({ onApproval: () => ({ approved: false, reason: 'not today' }) })
+			const output = await executeTool(ctx, call('notifyDesk'), {
+				tools: approvalTools(executed),
+				...at,
+			})
+			expect(output).toEqual({
+				type: 'error-text',
+				value: 'Tool "notifyDesk" was not approved: not today.',
+			})
+			expect(executed).toHaveLength(0)
+		})
+
+		test('a timeout counts as a rejection', async () => {
+			const executed = []
+			const ctx = fakeCtx({
+				onApproval: () => {
+					throw new TimeoutError()
+				},
+			})
+			const output = await executeTool(ctx, call('notifyDesk'), {
+				tools: approvalTools(executed),
+				...at,
+			})
+			expect(output.value).toContain('approval timed out')
+			expect(executed).toHaveLength(0)
+		})
+	})
+})

package/test/agent/fakes.js

@@ -0,0 +1,84 @@
+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 const env = { LLM_BASE_URL: 'https://unused.example', MODEL_ID: 'unused' }