diff --git a/services/agent-harness/src/client-tools.test.ts b/services/agent-harness/src/client-tools.test.ts new file mode 100644 index 0000000000..a78c8f2b3a --- /dev/null +++ b/services/agent-harness/src/client-tools.test.ts @@ -0,0 +1,1326 @@ +import { createHash } from 'node:crypto'; +import { env } from 'cloudflare:workers'; +import { abortAllDurableObjects, runDurableObjectAlarm, runInDurableObject } from 'cloudflare:test'; +import { eq } from 'drizzle-orm'; +import { drizzle } from 'drizzle-orm/durable-sqlite'; +import { MockLanguageModelV3 } from 'ai/test'; +import { describe, expect, it } from 'vitest'; +import { z } from 'zod'; +import { canonicalizeValidatedInput, type Command } from '@kilocode/agent-harness/commands'; +import { + ConversationSchema, + ExecutionGrantSchema, + RunSchema, + ToolCallSchema, + type Client, + type ToolOutcome, +} from '@kilocode/agent-harness/contracts'; +import type { BridgeReadiness } from '@kilocode/agent-harness/bridge'; +import { createHarnessClient } from '@kilocode/agent-harness/client'; +import { + JournalSnapshotSchema, + completionCommand, + type CommandReply, + type HarnessJournal, + type JournalSnapshot, +} from '@kilocode/agent-harness/journal'; +import { toolDefinitions, type ToolName } from '@kilocode/agent-harness/tools'; +import { admitCommand, type CommandAdapter } from './commands'; +import { createScheduler, SchedulerStateSchema, type SchedulerAdapter } from './scheduler'; +import type { ClientToolAuthorizer, ClientToolCommand } from './client-tools'; +import { insertCall, insertCheckpoint } from './db/records'; +import { openStore, type ConversationStore } from './db/store'; +import { getTestStoreStub, type TestStore } from './db/test-worker'; +import { StoreError } from './db/wake'; +import * as s from './db/sqlite-schema'; + +const bindings = env as { STORE: DurableObjectNamespace }; +const digest = (text: string) => createHash('sha256').update(text).digest('hex'); +const ready: BridgeReadiness = { + available: true, + foreground: true, + connectivity: 'confirmed', + unlock: 'ready', + gesture: 'not_required', +}; +const receipt: ToolOutcome = { status: 'succeeded', output: { permission: 'denied' } }; +const grantReply = (reply: CommandReply) => { + if (reply.status !== 'accepted') throw new Error(JSON.stringify(reply)); + return z.object({ grant: ExecutionGrantSchema, toolCall: ToolCallSchema }).parse(reply.result); +}; +function ledger(state: DurableObjectState, runId: string) { + return SchedulerStateSchema.parse( + drizzle(state.storage) + .select() + .from(s.checkpoints) + .all() + .find(row => row.runId === runId && row.step === 0)?.data + ); +} +function deferred() { + let resolve: () => void = () => { + throw new Error('Missing resolver'); + }; + const promise = new Promise(done => { + resolve = done; + }); + return { promise, resolve }; +} +async function fixture( + names: ToolName[] = ['app.notifications'], + limits?: CommandAdapter['limits'] +) { + const conversation = ConversationSchema.parse({ + id: crypto.randomUUID(), + ownerUserId: 'auth0|owner', + context: { type: 'personal' }, + permissionMode: 'yolo', + }); + const phone: Client = { + id: crypto.randomUUID(), + ownerUserId: conversation.ownerUserId, + kind: 'mobile', + supportedTools: toolDefinitions + .filter(tool => tool.executorKind === 'client') + .map(tool => ({ name: tool.name, version: tool.version })), + revokedAt: null, + }; + const web: Client = { ...phone, id: crypto.randomUUID(), kind: 'browser' }; + const base = { protocolVersion: 1 as const, conversationId: conversation.id, clientId: phone.id }; + const f = { + registration: phone, + readiness: { ...ready }, + storageReady: true, + authorized: true, + clock: Date.now() + 3_600_000, + executions: [] as string[], + }; + const now = () => f.clock; + const stub = () => getTestStoreStub(bindings.STORE, conversation.id); + const use = (work: (store: ConversationStore, state: DurableObjectState) => T | Promise) => + runInDurableObject(stub(), (instance, state) => work(instance.store, state)); + const authority: ClientToolAuthorizer = async command => ({ + conversation, + client: command.clientId === web.id ? web : f.registration, + readiness: f.readiness, + storageReady: f.storageReady, + }); + const commandAdapter: CommandAdapter = { + limits, + authorize: async () => ({ conversation, client: phone, origin: 'user' }), + validateModel: async () => ({ + contextTokens: 32000, + inputUsdPerMillion: 0.1, + outputUsdPerMillion: 0.2, + }), + now, + }; + const model = new MockLanguageModelV3({ + modelId: 'test/model', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'text-start', id: 'text' }); + controller.enqueue({ type: 'text-delta', id: 'text', delta: 'done' }); + controller.enqueue({ type: 'text-end', id: 'text' }); + controller.enqueue({ + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 1, text: 1, reasoning: 0 }, + }, + }); + controller.close(); + }, + }), + }), + }); + const adapter: SchedulerAdapter = { + definitions: toolDefinitions, + model: () => model, + countTokens: () => 10, + system: 'Untrusted tool data.', + now, + authorize: async () => undefined, + policy: async current => ({ + permissionMode: current.permissionMode, + permissionRevision: current.permissionRevision, + expectedPermissionRevision: current.permissionRevision, + authorized: f.authorized, + available: true, + trustedRead: true, + clientReady: true, + questionAnswered: true, + }), + dispatch: async ({ call }) => { + f.executions.push(call.id); + return { status: 'succeeded', output: [] }; + }, + }; + const send = () => + use(async (store, state) => { + const commandId = crypto.randomUUID(); + expect( + await admitCommand( + state, + store, + { + ...base, + type: 'sendMessage', + commandId, + text: 'hello', + modelId: 'test/model', + permissionRevision: store.snapshot()!.conversation.permissionRevision, + }, + commandAdapter + ) + ).toMatchObject({ status: 'accepted' }); + return commandId; + }); + await use(store => store.bindExistingConversation(conversation)); + const runId = await send(); + const calls = await use(async store => { + const run = store.queuedRuns()[0].data; + const calls = names.map(name => { + const definition = toolDefinitions.find(tool => tool.name === name)!; + const input = + name === 'app.openScreen' + ? { screen: 'preferences' } + : name === 'app.setPreference' + ? { name: 'showToolDetails', value: true } + : {}; + return ToolCallSchema.parse({ + id: crypto.randomUUID(), + runId, + name, + definitionVersion: definition.version, + arguments: input, + context: conversation.context, + effect: definition.effect, + executionTarget: { + kind: definition.executorKind, + ...(definition.executorKind === 'client' ? { clientId: phone.id } : {}), + }, + approval: null, + state: 'pending', + result: null, + }); + }); + await store.transition({ wakeAt: now() }, db => { + const checkpointId = crypto.randomUUID(); + insertCheckpoint(db, { + id: checkpointId, + runId, + step: 1, + status: 'complete', + definitionVersions: Object.fromEntries( + calls.map(call => [call.name, call.definitionVersion]) + ), + data: { + kind: 'complete', + attemptId: crypto.randomUUID(), + messageId: crypto.randomUUID(), + createdAt: new Date(now()).toISOString(), + text: '', + finishReason: 'tool-calls', + citations: [], + usage: { inputTokens: 1, outputTokens: 1 }, + calls: calls.map((call, index) => ({ sdkId: `sdk-${index}`, call })), + responseMessages: [ + { + role: 'assistant', + content: calls.map((call, index) => ({ + type: 'tool-call', + toolCallId: `sdk-${index}`, + toolName: call.name, + input: call.arguments, + })), + }, + ], + }, + }); + calls.forEach((call, position) => + insertCall(db, call, { + checkpointId, + position, + inputDigest: digest(canonicalizeValidatedInput(call.arguments)), + policy: {}, + }) + ); + return { events: [{ type: 'run', run: { ...run, state: { status: 'running' } } }] }; + }); + return calls; + }); + const scheduler = (store: ConversationStore, state: DurableObjectState) => + createScheduler(state, store, adapter); + const alarm = () => + use(async (store, state) => { + await state.storage.deleteAlarm(); + await scheduler(store, state).alarm(); + }); + const claim = (clientId = phone.id, toolCallId = calls[0].id): ClientToolCommand => ({ + ...base, + type: 'claimClientTool', + commandId: crypto.randomUUID(), + clientId, + toolCallId, + }); + const perform = (command: unknown) => + use((store, state) => scheduler(store, state).clientTool(command, authority)); + const complete = (request: ReturnType, result = receipt): ClientToolCommand => + completionCommand({ + ...request, + completionCommandId: crypto.randomUUID(), + receipt: result, + }) as ClientToolCommand; + const cancel = (id = runId) => ({ + ...base, + type: 'cancelRun', + runId: id, + commandId: crypto.randomUUID(), + }); + const setMode = (mode: 'ask' | 'yolo') => + use((store, state) => + admitCommand( + state, + store, + { + ...base, + type: 'setPermissionMode', + commandId: crypto.randomUUID(), + permissionMode: mode, + expectedPermissionRevision: store.snapshot()!.conversation.permissionRevision, + acknowledgePendingActions: true, + }, + commandAdapter + ) + ); + await alarm(); + return Object.assign(f, { + base, + conversation, + phone, + web, + calls, + runId, + adapter, + commandAdapter, + authority, + use, + now, + stub, + scheduler, + send, + alarm, + claim, + perform, + complete, + cancel, + setMode, + }); +} + +describe('designated client grants on real SQLite', () => { + it.each(['authorization', 'policy', 'final-authority'] as const)( + 'charges delayed %s checks before refusing an over-budget grant', + async boundary => { + const f = await fixture(undefined, { activeRunMs: 10 }); + await f.use(async (store, state) => { + let reads = 0; + const runtime = createScheduler(state, store, { + ...f.adapter, + authorize: async (...args) => { + await f.adapter.authorize(...args); + if (boundary === 'authorization') f.clock += 100; + }, + policy: async (...args) => { + const result = await f.adapter.policy(...args); + if (boundary === 'policy') f.clock += 100; + return result; + }, + }); + expect( + await runtime.clientTool(f.claim(), async command => { + const result = await f.authority(command); + if (++reads === 2 && boundary === 'final-authority') f.clock += 100; + return result; + }) + ).toMatchObject({ status: 'rejected', error: { code: 'limit_exceeded' } }); + expect(ledger(state, f.runId).reservations.at(-1)).toMatchObject({ + status: 'released', + activeMs: 10, + }); + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + expect(drizzle(state.storage).select().from(s.attempts).all()).toEqual([]); + }); + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'limit_exceeded' }, + }); + } + ); + + it('charges failed checks and retains the original deadline for a later successful claim', async () => { + const f = await fixture(undefined, { activeRunMs: 20, toolAttemptMs: 10 }); + await f.use(async (store, state) => { + const failed = createScheduler(state, store, { + ...f.adapter, + policy: async () => { + f.clock += 6; + throw new StoreError('storage_unavailable', true); + }, + }); + expect(await failed.clientTool(f.claim(), f.authority)).toMatchObject({ status: 'rejected' }); + expect(ledger(state, f.runId).reservations.at(-1)).toMatchObject({ + status: 'released', + activeMs: 6, + }); + const started = f.now(); + const runtime = createScheduler(state, store, { + ...f.adapter, + policy: async (...args) => { + f.clock += 4; + return f.adapter.policy(...args); + }, + }); + const request = grantReply(await runtime.clientTool(f.claim(), f.authority)); + expect(Date.parse(request.grant.expiresAt)).toBe(started + 10); + expect(Date.parse(request.grant.expiresAt) - f.now()).toBe(6); + f.clock += 2; + expect(await runtime.clientTool(f.complete(request), f.authority)).toMatchObject({ + status: 'accepted', + }); + expect(ledger(state, f.runId).reservations.at(-1)).toMatchObject({ + status: 'finished', + activeMs: 6, + }); + }); + }); + + it('keeps lost preparation time reserved across restart without granting an effect', async () => { + const f = await fixture(undefined, { activeRunMs: 10 }); + await f.use(async (store, state) => { + let persisted = false; + const lost: ConversationStore = { + ...store, + transition: async (...args) => { + if (persisted) throw new StoreError('storage_unavailable', true); + await store.transition(...args); + persisted = true; + throw new StoreError('storage_unavailable', true); + }, + }; + expect(await f.scheduler(lost, state).clientTool(f.claim(), f.authority)).toMatchObject({ + status: 'rejected', + error: { code: 'storage_unavailable' }, + }); + expect(ledger(state, f.runId).reservations.at(-1)).toMatchObject({ + status: 'reserved', + activeMs: 10, + }); + }); + await abortAllDurableObjects(); + f.clock += 100; + await f.alarm(); + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'limit_exceeded' }, + }); + await f.use((_store, state) => { + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + expect(drizzle(state.storage).select().from(s.attempts).all()).toEqual([]); + expect(ledger(state, f.runId).reservations.at(-1)?.activeMs).toBe(10); + }); + }); + + it('does not charge time between genuine device waits', async () => { + const f = await fixture(undefined, { activeRunMs: 10 }); + f.readiness = { ...ready, unlock: 'locked' }; + expect(await f.perform(f.claim())).toMatchObject({ + status: 'accepted', + result: { grant: null }, + }); + f.clock += 100000; + f.readiness = { ...ready }; + const request = grantReply(await f.perform(f.claim())); + expect(Date.parse(request.grant.expiresAt) - f.now()).toBe(10); + }); + + it.each(['policy failure', 'budget exhaustion', 'successful checks'] as const)( + 'replays a lost wait acknowledgment during %s without another reservation', + async scenario => { + const f = await fixture(['app.notifications', 'app.openSettings'], { activeRunMs: 20 }); + f.readiness = { ...ready, unlock: 'locked' }; + const command = f.claim(); + const committed = await f.use(async (store, state) => { + const lost: ConversationStore = { + ...store, + transition: async (...args) => { + const reply = await store.transition(...args); + if (args[0].command?.id === command.commandId) + throw new StoreError('storage_unavailable', true); + return reply; + }, + }; + expect(await f.scheduler(lost, state).clientTool(command, f.authority)).toMatchObject({ + status: 'rejected', + error: { code: 'storage_unavailable' }, + }); + const saved = store.getCommand(command.commandId); + expect(saved?.reply).toMatchObject({ + status: 'accepted', + result: { grant: null, decision: 'client' }, + }); + return saved; + }); + await abortAllDurableObjects(); + f.clock += 1000; + f.readiness = { ...ready }; + const policy = f.adapter.policy; + f.adapter.policy = async (...args) => { + f.clock += scenario === 'budget exhaustion' ? 20 : 4; + if (scenario === 'policy failure') throw new StoreError('storage_unavailable', true); + return policy(...args); + }; + if (scenario === 'budget exhaustion') + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'limit_exceeded' }, + }); + const before = await f.use((store, state) => ({ + snapshot: store.snapshot(), + budget: ledger(state, f.runId), + })); + expect(await f.perform(command)).toEqual(committed?.reply); + expect(await f.perform(command)).toEqual(committed?.reply); + await f.use((store, state) => { + expect(store.getCommand(command.commandId)).toEqual(committed); + expect(store.snapshot()).toEqual(before.snapshot); + expect(ledger(state, f.runId)).toEqual(before.budget); + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + expect(drizzle(state.storage).select().from(s.attempts).all()).toEqual([]); + }); + expect(await f.perform({ ...command, toolCallId: f.calls[1].id })).toMatchObject({ + status: 'rejected', + error: { code: 'command_conflict', retryable: false }, + }); + } + ); + + it.each(['storage', 'revoked', 'capability'] as const)( + 'rejects %s loss before replaying a committed wait', + async reason => { + const f = await fixture(); + f.readiness = { ...ready, unlock: 'locked' }; + const command = f.claim(); + const committed = await f.perform(command); + expect(committed).toMatchObject({ status: 'accepted', result: { grant: null } }); + if (reason === 'storage') f.storageReady = false; + else + f.registration = { + ...f.registration, + ...(reason === 'revoked' + ? { revokedAt: new Date(f.now()).toISOString() } + : { supportedTools: [] }), + }; + expect(await f.perform(command)).toMatchObject({ + status: 'rejected', + error: { + code: + reason === 'storage' + ? 'storage_unavailable' + : reason === 'revoked' + ? 'access_revoked' + : 'unavailable_tool', + }, + }); + await f.use((store, state) => { + expect(store.getCommand(command.commandId)?.reply).toEqual(committed); + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'unavailable', grant: null }, + ]); + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + }); + } + ); + + it('races claims, retains a lost completion acknowledgment, and settles one SDK result', async () => { + const f = await fixture(); + const commands = [f.claim(), f.claim()]; + const replies = await Promise.all(commands.map(f.perform)); + const request = grantReply(replies[0]); + expect(grantReply(replies[1])).toEqual(request); + await f.use((store, state) => { + expect(drizzle(state.storage).select().from(s.grants).all()).toHaveLength(1); + expect(drizzle(state.storage).select().from(s.attempts).all()).toMatchObject([ + { + intent: { + grant: request.grant, + inputDigest: digest('{}'), + policy: { decision: 'dispatch', clientReady: true }, + }, + }, + ]); + expect(request.grant).toMatchObject({ + clientId: f.phone.id, + ownerUserId: f.conversation.ownerUserId, + toolCallId: f.calls[0].id, + definitionVersion: '1', + context: f.conversation.context, + }); + expect(Date.parse(request.grant.expiresAt) - f.now()).toBe(30000); + expect(store.snapshot()?.pendingClientActions).toEqual([]); + }); + const command = f.complete(request); + await f.use(async (store, state) => { + const lost = { + ...store, + transition: async (...args: Parameters) => { + await store.transition(...args); + throw new StoreError('storage_unavailable', true); + }, + } satisfies ConversationStore; + expect(await f.scheduler(lost, state).clientTool(command, f.authority)).toMatchObject({ + status: 'rejected', + error: { code: 'storage_unavailable' }, + }); + expect(store.getCommand(command.commandId)?.reply.status).toBe('accepted'); + }); + await abortAllDurableObjects(); + const reply = await f.perform(command); + expect(reply).toMatchObject({ status: 'accepted', result: { result: receipt } }); + await f.alarm(); + expect(await f.perform(command)).toEqual(reply); + expect( + await f.perform({ + ...command, + result: { status: 'succeeded', output: { permission: 'granted' } }, + }) + ).toMatchObject({ status: 'rejected', error: { code: 'command_conflict' } }); + expect( + await f.perform({ + ...command, + commandId: crypto.randomUUID(), + result: { status: 'cancelled' }, + }) + ).toMatchObject({ status: 'rejected', error: { code: 'command_conflict' } }); + await f.use((store, state) => { + expect(store.snapshot()?.activeRun).toBeNull(); + expect(store.callsForRun(f.runId)[0].data.result).toEqual(receipt); + expect(ledger(state, f.runId).resultMessages[f.calls[0].id]).toMatchObject({ + role: 'tool', + content: [{ toolCallId: 'sdk-0', output: { value: { permission: 'denied' } } }], + }); + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength(1); + expect(f.executions).toEqual([]); + }); + }); + + it.each([ + [{ foreground: false }, 'background'], + [{ connectivity: 'offline' }, 'offline'], + [{ connectivity: 'unknown' }, 'offline'], + [{ unlock: 'locked' }, 'locked'], + [{ unlock: 'unknown' }, 'locked'], + [{ gesture: 'required' }, 'gesture'], + [{ available: false }, 'unavailable'], + ] as const)('persists %j without transferring or consuming a grant', async (change, reason) => { + const f = await fixture(); + const later = await f.send(); + f.readiness = { ...ready, ...change }; + const command = f.claim(); + expect(await f.perform(command)).toMatchObject({ status: 'accepted', result: { grant: null } }); + await abortAllDurableObjects(); + await f.use(async (store, state) => { + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason, grant: null, toolCall: { executionTarget: { clientId: f.phone.id } } }, + ]); + expect(store.queuedRuns().map(row => row.id)).toEqual([later]); + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + expect(ledger(state, f.runId).reservations.every(item => item.status === 'released')).toBe( + true + ); + }); + expect(await f.perform(f.claim(f.web.id))).toMatchObject({ + status: 'rejected', + error: { code: 'access_revoked' }, + }); + f.readiness = { ...ready }; + expect(await f.perform(command)).toMatchObject({ result: { grant: null } }); + expect(grantReply(await f.perform(f.claim())).grant.clientId).toBe(f.phone.id); + }); + + it.each(['expired', 'offline', 'storage_unavailable', 'access_revoked'] as const)( + 'keeps a claimed %s action unknown and accepts only its eventual receipt', + async reason => { + const f = await fixture(); + const request = grantReply(await f.perform(f.claim())); + const later = await f.send(); + await abortAllDurableObjects(); + if (reason === 'expired') { + f.clock += 30001; + await f.alarm(); + } else + await f.use((store, state) => + f.scheduler(store, state).clientUnavailable(f.phone.id, reason) + ); + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'outcome_unknown' }, + }); + expect(await f.perform(f.claim(f.web.id))).toMatchObject({ + status: 'rejected', + error: { code: 'access_revoked' }, + }); + await f.use((store, state) => { + expect(store.snapshot()?.activeRun?.state).toMatchObject({ + status: 'waiting', + waiting: { reason: 'reconciliation' }, + }); + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'reconciliation', grant: request.grant }, + ]); + expect(store.queuedRuns().map(row => row.id)).toEqual([later]); + expect(store.callsForRun(f.runId)[0].data).toMatchObject({ + state: 'executing', + result: null, + }); + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength(1); + }); + expect(await f.perform(f.complete(request))).toMatchObject({ + status: 'accepted', + result: { result: receipt }, + }); + await f.alarm(); + expect(f.executions).toEqual([]); + } + ); + + it.each(['storage', 'revoked', 'capability'] as const)( + 'rejects %s identity before replay and preserves the uncertain effect', + async reason => { + const f = await fixture(); + const claim = f.claim(), + request = grantReply(await f.perform(claim)); + if (reason === 'storage') f.storageReady = false; + else + f.registration = { + ...f.registration, + ...(reason === 'revoked' + ? { revokedAt: new Date(f.now()).toISOString() } + : { supportedTools: [{ name: 'app.notifications', version: '2' }] }), + }; + const error = { + code: + reason === 'storage' + ? 'storage_unavailable' + : reason === 'revoked' + ? 'access_revoked' + : 'unavailable_tool', + }; + expect(await f.perform(claim)).toMatchObject({ status: 'rejected', error }); + expect(await f.perform(f.complete(request))).toMatchObject({ status: 'rejected', error }); + await f.use(store => + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'reconciliation', grant: request.grant }, + ]) + ); + } + ); + + it('keeps an incapable designated client waiting and refuses another registration', async () => { + const f = await fixture(); + f.registration = { ...f.phone, supportedTools: [] }; + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'unavailable_tool' }, + }); + await f.use((store, state) => { + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'unavailable', grant: null }, + ]); + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]); + }); + f.registration = { ...f.phone, id: crypto.randomUUID() }; + expect(await f.perform(f.claim(f.registration.id))).toMatchObject({ + status: 'rejected', + error: { code: 'access_revoked' }, + }); + }); + + it.each(['grantId', 'generation', 'clientId', 'result', 'oversized'] as const)( + 'rejects a foreign or invalid completion %s without changing the actual result', + async field => { + const f = await fixture(); + const request = grantReply(await f.perform(f.claim())), + command = f.complete(request); + const changes = { + grantId: { grantId: crypto.randomUUID() }, + generation: { generation: request.grant.generation + 1 }, + clientId: { clientId: f.web.id }, + result: { result: { status: 'succeeded', output: { enabled: true } } }, + oversized: { result: { status: 'outcome_unknown', reason: 'x'.repeat(65536) } }, + }; + expect(await f.perform({ ...command, ...changes[field] })).toMatchObject({ + status: 'rejected', + error: { + code: + field === 'result' + ? 'invalid_output' + : field === 'oversized' + ? 'limit_exceeded' + : 'access_revoked', + }, + }); + await f.use(store => + expect(store.callsForRun(f.runId)[0].data).toMatchObject({ + state: 'executing', + result: null, + }) + ); + expect(await f.perform(f.complete(request))).toMatchObject({ status: 'accepted' }); + } + ); + + it.each(['stop', 'mode', 'identity', 'policy'] as const)( + 'fences %s changed during the policy check before granting execution', + async boundary => { + const f = await fixture(); + await f.use(async (store, state) => { + const entered = deferred(), + release = deferred(); + const runtime = createScheduler(state, store, { + ...f.adapter, + policy: async (...args) => { + const policy = await f.adapter.policy(...args); + entered.resolve(); + await release.promise; + return boundary === 'policy' ? { ...policy, authorized: false } : policy; + }, + }); + const pending = runtime.clientTool(f.claim(), f.authority); + await entered.promise; + if (boundary === 'stop') await admitCommand(state, store, f.cancel(), f.commandAdapter); + if (boundary === 'mode') + await admitCommand( + state, + store, + { + ...f.base, + type: 'setPermissionMode', + commandId: crypto.randomUUID(), + permissionMode: 'ask', + expectedPermissionRevision: 0, + }, + f.commandAdapter + ); + if (boundary === 'identity') + f.registration = { ...f.phone, revokedAt: new Date(f.now()).toISOString() }; + release.resolve(); + const reply = await pending; + expect(reply).toMatchObject({ + status: 'rejected', + error: { + code: + boundary === 'stop' + ? 'cancelled' + : boundary === 'mode' + ? 'stale_revision' + : 'access_revoked', + }, + }); + expect(drizzle(state.storage).select().from(s.attempts).all()).toEqual([]); + }); + } + ); + + it.each(['succeeded', 'failed', 'outcome_unknown'] as const)( + 'keeps the actual %s receipt after Stop and cancels only remaining calls', + async status => { + const f = await fixture(['app.notifications', 'app.openSettings']); + const request = grantReply(await f.perform(f.claim())); + const later = await f.send(); + await f.use((store, state) => admitCommand(state, store, f.cancel(), f.commandAdapter)); + await f.alarm(); + f.clock += 30001; + await f.alarm(); + const result: ToolOutcome = + status === 'succeeded' + ? receipt + : status === 'failed' + ? { + status, + error: { + code: 'unavailable_tool', + message: 'Permission unavailable.', + retryable: false, + }, + } + : { status, reason: 'Receipt unavailable.' }; + expect(await f.perform(f.complete(request, result))).toMatchObject({ + status: 'accepted', + result: { result }, + }); + await f.use(async (store, state) => { + const calls = store.callsForRun(f.runId); + expect(calls[1].data.result).toEqual({ status: 'cancelled' }); + expect(calls[0].data.result).toEqual(status === 'outcome_unknown' ? null : result); + expect(store.queuedRuns().map(row => row.id)).toEqual([later]); + if (status !== 'outcome_unknown') { + const before = store.snapshot(); + expect(await admitCommand(state, store, f.cancel(), f.commandAdapter)).toMatchObject({ + status: 'accepted', + result: { state: { status: 'cancelled' } }, + }); + expect(store.snapshot()).toEqual(before); + } else + expect(store.snapshot()?.activeRun?.state).toMatchObject({ + status: 'waiting', + waiting: { reason: 'reconciliation' }, + }); + }); + } + ); + + it.each(['claim', 'completion'] as const)( + 'rolls back every %s record when the command journal fails', + async operation => { + const f = await fixture(); + const request = operation === 'completion' ? grantReply(await f.perform(f.claim())) : null; + const command = request ? f.complete(request) : f.claim(); + await f.use(async (original, state) => { + const before = original.snapshot(), + calls = original.callsForRun(f.runId), + budget = ledger(state, f.runId); + const store = { + ...original, + transition: (options, write) => + original.transition(options, db => { + const changes = write(db); + if (options.command) + db.insert(s.commands) + .values({ + id: command.commandId, + fingerprint: 'injected conflict', + reply: changes.reply, + sequence: 0, + }) + .run(); + return changes; + }), + } satisfies ConversationStore; + expect(await f.scheduler(store, state).clientTool(command, f.authority)).toMatchObject({ + status: 'rejected', + error: { code: 'storage_unavailable' }, + }); + expect(original.snapshot()).toEqual(before); + expect(original.callsForRun(f.runId)).toEqual(calls); + const afterBudget = ledger(state, f.runId); + expect(afterBudget.currentReservationId).toEqual(budget.currentReservationId); + expect(afterBudget.reservations.slice(0, budget.reservations.length)).toEqual( + budget.reservations + ); + expect(afterBudget.reservations.slice(budget.reservations.length)).toEqual( + request ? [] : [expect.objectContaining({ status: 'released', activeMs: 0 })] + ); + expect(original.getCommand(command.commandId)).toBeNull(); + expect(drizzle(state.storage).select().from(s.grants).all()).toHaveLength(request ? 1 : 0); + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength( + request ? 1 : 0 + ); + }); + expect(await f.perform(command)).toMatchObject({ status: 'accepted' }); + } + ); + + it.each(['claim', 'completion'] as const)( + 'recovers %s before-arm, after-arm, and after-commit without another client request', + async operation => { + for (const boundary of ['before-arm', 'after-arm', 'after-commit'] as const) { + const f = await fixture(); + const request = operation === 'completion' ? grantReply(await f.perform(f.claim())) : null; + const command = request ? f.complete(request) : f.claim(); + await f.use(async (original, state) => { + const beforeAlarm = await state.storage.getAlarm(); + const store = + boundary === 'after-commit' + ? ({ + ...original, + transition: async (...args: Parameters) => { + const reply = await original.transition(...args); + if (args[0].command) throw new StoreError('storage_unavailable', true); + return reply; + }, + } satisfies ConversationStore) + : await openStore(state, { + getAlarm: () => state.storage.getAlarm(), + setAlarm: async deadline => { + if (boundary === 'after-arm') await state.storage.setAlarm(deadline); + throw new Error('Injected alarm failure'); + }, + }); + expect(await f.scheduler(store, state).clientTool(command, f.authority)).toMatchObject({ + status: 'rejected', + error: { code: 'storage_unavailable' }, + }); + expect(original.getCommand(command.commandId)?.reply.status).toBe( + boundary === 'after-commit' ? 'accepted' : undefined + ); + expect(drizzle(state.storage).select().from(s.grants).all()).toHaveLength( + request || boundary === 'after-commit' ? 1 : 0 + ); + expect(await state.storage.getAlarm()).toBe( + boundary === 'before-arm' ? beforeAlarm : f.now() + 1 + ); + }); + await abortAllDurableObjects(); + f.clock += 30001; + await runInDurableObject(f.stub(), (instance, state) => { + instance.alarm = f.scheduler(instance.store, state).alarm; + }); + expect(await runDurableObjectAlarm(f.stub())).toBe( + operation !== 'claim' || boundary !== 'before-arm' + ); + await f.use((store, state) => { + const completed = operation === 'completion' && boundary === 'after-commit'; + expect( + RunSchema.parse( + drizzle(state.storage).select().from(s.runs).where(eq(s.runs.id, f.runId)).get()?.data + ).state + ).toMatchObject( + completed + ? { status: 'completed' } + : { + status: 'waiting', + waiting: { + reason: request || boundary === 'after-commit' ? 'reconciliation' : 'client', + }, + } + ); + expect(store.callsForRun(f.runId)[0].data.result).toEqual(completed ? receipt : null); + expect(f.executions).toEqual([]); + }); + } + } + ); + + it.each(['execution', 'effect', 'receipt', 'send'] as const)( + 'uses the shared receipt contract across failures before and after %s', + async point => { + for (const side of ['before', 'after'] as const) { + const f = await fixture(); + const request = { + ...grantReply(await f.perform(f.claim())), + completionCommandId: crypto.randomUUID(), + }; + await f.use(async (store, state) => { + const scope = { + ownerUserId: f.phone.ownerUserId, + clientId: f.phone.id, + storageGeneration: crypto.randomUUID(), + }; + let durable: JournalSnapshot = { + scope, + revision: 0, + intents: [], + acknowledgments: [], + executions: [], + }; + let fault = true, + effects = 0; + const edge = (at: string, when: string) => { + if (fault && at === point && when === side) { + fault = false; + throw new Error(`Lost ${when} ${at}`); + } + }; + const journal: HarnessJournal = { + read: async () => structuredClone(durable), + compareAndSwap: async (_scope, revision, next) => { + const at = + next.executions.length > durable.executions.length + ? 'execution' + : next.executions[0]?.receipt && !durable.executions[0]?.receipt + ? 'receipt' + : 'journal'; + edge(at, 'before'); + if (durable.revision !== revision) return false; + durable = JournalSnapshotSchema.parse(structuredClone(next)); + edge(at, 'after'); + return true; + }, + }; + const open = () => + createHarnessClient({ + scope, + currentScope: () => scope, + journal, + now: f.now, + digest, + transport: { + send: async (_scope, command: Command) => { + edge('send', 'before'); + expect(durable.executions[0]?.receipt).toEqual(receipt); + const reply = await f.scheduler(store, state).clientTool(command, f.authority); + edge('send', 'after'); + return reply; + }, + }, + bridge: { + readiness: () => ready, + execute: async () => { + edge('effect', 'before'); + expect(durable.executions[0]?.grant.id).toBe(request.grant.id); + effects++; + edge('effect', 'after'); + return receipt; + }, + reconcileReceipt: async () => (effects ? receipt : null), + }, + }); + await open().dispatch(request); + expect(fault).toBe(false); + await open().recover(); + await open().dispatch(request); + await open().recover(); + const uncertain = + (point === 'execution' && side === 'after') || + (point === 'effect' && side === 'before'); + expect(effects).toBe(uncertain ? 0 : 1); + expect(store.callsForRun(f.runId)[0].data.result).toEqual(uncertain ? null : receipt); + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength(1); + if (uncertain) { + f.clock += 30001; + await f.scheduler(store, state).alarm(); + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'reconciliation', grant: request.grant }, + ]); + } + }); + } + } + ); + + it.each(['storage', 'account'] as const)( + 'invalidates local dispatch after %s loss without replacing the server grant', + async loss => { + const f = await fixture(), + request = { + ...grantReply(await f.perform(f.claim())), + completionCommandId: crypto.randomUUID(), + }; + await f.use(async (store, state) => { + const scope = { + ownerUserId: f.phone.ownerUserId, + clientId: f.phone.id, + storageGeneration: crypto.randomUUID(), + }; + let effects = 0; + const client = createHarnessClient({ + scope, + currentScope: () => (loss === 'account' ? null : scope), + now: f.now, + digest, + journal: { + read: async () => null, + compareAndSwap: async () => { + throw new Error('Unsafe storage replacement'); + }, + }, + transport: { + send: async () => { + throw new Error('Unexpected send'); + }, + }, + bridge: { + readiness: () => ready, + execute: async () => { + effects++; + return receipt; + }, + reconcileReceipt: async () => null, + }, + }); + expect(await client.dispatch(request)).toMatchObject({ + status: 'unknown', + error: { code: loss === 'account' ? 'access_revoked' : 'storage_unavailable' }, + }); + expect(await client.dispatch(request)).toMatchObject({ status: 'unknown' }); + await f + .scheduler(store, state) + .clientUnavailable( + f.phone.id, + loss === 'account' ? 'access_revoked' : 'storage_unavailable' + ); + expect(effects).toBe(0); + expect(store.snapshot()?.pendingClientActions).toMatchObject([ + { reason: 'reconciliation', grant: request.grant }, + ]); + }); + } + ); + + it.each(['partial', 'arguments', 'context', 'definitionVersion', 'inputDigest'] as const)( + 'refuses a changed %s checkpoint or call before issuing a grant', + async field => { + const f = await fixture(); + await f.use((store, state) => { + const db = drizzle(state.storage), + call = store.callsForRun(f.runId)[0]; + if (field === 'partial') + db.update(s.checkpoints) + .set({ status: 'partial' }) + .where(eq(s.checkpoints.id, call.checkpointId)) + .run(); + else if (field === 'inputDigest') + db.update(s.calls).set({ inputDigest: 'changed' }).where(eq(s.calls.id, call.id)).run(); + else + db.update(s.calls) + .set({ + data: { + ...call.data, + [field]: + field === 'arguments' + ? { changed: true } + : field === 'context' + ? { type: 'organization', organizationId: crypto.randomUUID() } + : '2', + }, + }) + .where(eq(s.calls.id, call.id)) + .run(); + }); + expect(await f.perform(f.claim())).toMatchObject({ + status: 'rejected', + error: { code: 'invalid_output' }, + }); + await f.use((_store, state) => + expect(drizzle(state.storage).select().from(s.grants).all()).toEqual([]) + ); + } + ); + + it.each(['ownerUserId', 'inputDigest', 'generation', 'expiresAt'] as const)( + 'rejects corrupt persisted grant %s without another dispatch', + async field => { + const f = await fixture(), + claim = f.claim(), + request = grantReply(await f.perform(claim)); + await f.use((_store, state) => + drizzle(state.storage) + .update(s.grants) + .set({ + data: { + ...request.grant, + [field]: field === 'generation' ? request.grant.generation + 1 : 'changed', + }, + }) + .where(eq(s.grants.id, request.grant.id)) + .run() + ); + expect(await f.perform(claim)).toMatchObject({ + status: 'rejected', + error: { code: 'invalid_output' }, + }); + expect(await f.perform(f.complete(request))).toMatchObject({ + status: 'rejected', + error: { code: 'invalid_output' }, + }); + await f.use((store, state) => { + expect(store.callsForRun(f.runId)[0].data.result).toBeNull(); + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength(1); + }); + } + ); + + it.each(['call', 'attempt', 'unsettled'] as const)( + 'rejects a mismatched %s receipt before canonical replay', + async field => { + const f = await fixture(), + request = grantReply(await f.perform(f.claim())); + const command = f.complete(request); + expect(await f.perform(command)).toMatchObject({ status: 'accepted' }); + await f.use((store, state) => { + const db = drizzle(state.storage), + call = store.callsForRun(f.runId)[0]; + if (field === 'attempt') + db.update(s.attempts) + .set({ outcome: { status: 'cancelled' } }) + .where(eq(s.attempts.toolCallId, call.id)) + .run(); + else + db.update(s.calls) + .set( + field === 'call' + ? { data: { ...call.data, result: { status: 'cancelled' } } } + : { state: 'executing', data: { ...call.data, state: 'executing', result: null } } + ) + .where(eq(s.calls.id, call.id)) + .run(); + }); + expect(await f.perform(command)).toMatchObject({ + status: 'rejected', + error: { code: 'invalid_output' }, + }); + await f.use((_store, state) => + expect(drizzle(state.storage).select().from(s.attempts).all()).toHaveLength(1) + ); + } + ); + + it('preserves approval, client-read uncertainty, sequential calls, and backend settlement', async () => { + const f = await fixture(['app.currentScreen', 'kilo.organizations']); + const request = grantReply(await f.perform(f.claim())); + expect(await f.perform(f.claim(f.phone.id, f.calls[1].id))).toMatchObject({ + status: 'rejected', + }); + f.clock += 30001; + await f.alarm(); + await f.use(store => + expect(store.snapshot()?.activeRun?.state).toMatchObject({ + status: 'waiting', + waiting: { reason: 'reconciliation' }, + }) + ); + const result: ToolOutcome = { + status: 'succeeded', + output: { destination: { screen: 'preferences' }, data: {} }, + }; + expect(await f.perform(f.complete(request, result))).toMatchObject({ status: 'accepted' }); + await f.alarm(); + expect(f.executions).toEqual([f.calls[1].id]); + const approval = await fixture(); + await approval.setMode('ask'); + expect(await approval.perform(approval.claim())).toMatchObject({ + status: 'accepted', + result: { grant: null, decision: 'approval' }, + }); + await approval.use(async (store, state) => { + const interaction = store.snapshot()!.unresolvedInteractions[0]; + expect( + await approval.scheduler(store, state).resolveInteraction( + { + ...approval.base, + type: 'resolveInteraction', + commandId: crypto.randomUUID(), + interactionId: interaction.id, + resolution: { kind: 'approve' }, + }, + async () => ({ + conversation: approval.conversation, + client: approval.phone, + origin: 'user', + }) + ) + ).toMatchObject({ status: 'accepted' }); + }); + const granted = grantReply(await approval.perform(approval.claim())); + expect(granted.toolCall.approval?.decision).toBe('approve'); + await approval.setMode('yolo'); + await approval.setMode('ask'); + expect(await approval.perform(approval.complete(granted))).toMatchObject({ + status: 'accepted', + }); + }); +}); diff --git a/services/agent-harness/src/client-tools.ts b/services/agent-harness/src/client-tools.ts new file mode 100644 index 0000000000..cee1ed4b81 --- /dev/null +++ b/services/agent-harness/src/client-tools.ts @@ -0,0 +1,221 @@ +import { createHash } from 'node:crypto'; +import { eq } from 'drizzle-orm'; +import { z } from 'zod'; +import { BridgeReadinessSchema } from '@kilocode/agent-harness/bridge'; +import { + CommandSchema, + canonicalizeValidatedInput, + fingerprintCommand, + type Command, +} from '@kilocode/agent-harness/commands'; +import { + ClientSchema, + ConversationSchema, + ErrorSchema, + ExecutionGrantSchema, + ToolCallSchema, + ToolOutcomeSchema, + type EventEnvelope, + type ExecutionGrant, + type ToolCall, +} from '@kilocode/agent-harness/contracts'; +import type { CommandReply } from '@kilocode/agent-harness/journal'; +import { readConversation, type StoreDatabase } from './db/records'; +import type { ConversationStore } from './db/store'; +import * as s from './db/sqlite-schema'; +import { StoreError } from './db/wake'; +import { RuntimeError, fail } from './limits'; +import { validateStoredCall } from './model-step'; + +export type ClientToolCommand = Extract< + Command, + { type: 'claimClientTool' | 'completeClientTool' } +>; +const AuthoritySchema = z.strictObject({ + conversation: ConversationSchema, + client: ClientSchema, + readiness: BridgeReadinessSchema, + // Require the original durable journal. A replacement journal cannot restore this registration. + storageReady: z.boolean(), +}); +export type ClientToolAuthority = z.infer; +// Resolve identity and current registration from authentication. Readiness is only a gate hint; +// it cannot supply account authority, replace the target, or attest a durable execution receipt. +export type ClientToolAuthorizer = ( + command: ClientToolCommand +) => Promise }>; +export const supportsClientCall = (authority: ClientToolAuthority, call: ToolCall) => + authority.client.supportedTools.some( + tool => tool.name === call.name && tool.version === call.definitionVersion + ); +export const rejectClientCommand = ( + commandId: string, + error: z.infer +): CommandReply => ({ status: 'rejected', commandId, error }); + +export function clientAction( + call: ToolCall, + reason: 'offline' | 'background' | 'locked' | 'gesture' | 'unavailable' | 'reconciliation' | null, + grant: ExecutionGrant | null = null +): EventEnvelope['event'] { + return { + type: 'client_action', + toolCallId: call.id, + action: reason === null ? null : { toolCall: call, grant, reason }, + }; +} + +const IntentSchema = z.strictObject({ + toolCall: ToolCallSchema, + inputDigest: z.string().min(1), + policy: z.json(), + grant: ExecutionGrantSchema, +}); +export function readClientGrant(db: StoreDatabase, call: ToolCall) { + const grants = db.select().from(s.grants).where(eq(s.grants.toolCallId, call.id)).limit(2).all(); + const attempts = db + .select() + .from(s.attempts) + .where(eq(s.attempts.toolCallId, call.id)) + .limit(2) + .all(); + if (!grants.length && !attempts.length) return null; + if (grants.length !== 1 || attempts.length !== 1) + fail('invalid_output', 'The client call has no single dispatch fence.'); + const grant = ExecutionGrantSchema.parse(grants[0].data), + attempt = attempts[0]; + const intent = IntentSchema.parse(attempt.intent), + conversation = readConversation(db); + validateStoredCall(call, intent.toolCall, [], null); + const digest = createHash('sha256') + .update(canonicalizeValidatedInput(call.arguments)) + .digest('hex'); + if ( + call.executionTarget.kind !== 'client' || + grant.clientId !== call.executionTarget.clientId || + grant.toolCallId !== call.id || + grant.conversationId !== conversation.id || + grant.ownerUserId !== conversation.ownerUserId || + grant.definitionVersion !== call.definitionVersion || + grant.inputDigest !== digest || + intent.inputDigest !== digest || + grants[0].id !== grant.id || + grants[0].generation !== grant.generation || + attempt.generation !== grant.generation || + canonicalizeValidatedInput(grant.context) !== canonicalizeValidatedInput(call.context) || + canonicalizeValidatedInput(intent.grant) !== canonicalizeValidatedInput(grant) + ) + fail('invalid_output', 'The persisted client grant does not match its immutable intent.'); + const outcome = attempt.outcome === null ? null : ToolOutcomeSchema.parse(attempt.outcome); + if ( + (call.result !== null && + canonicalizeValidatedInput(call.result) !== canonicalizeValidatedInput(outcome)) || + (call.result === null && outcome !== null && outcome.status !== 'outcome_unknown') + ) + fail('invalid_output', 'The stored completion does not match its dispatch attempt.'); + return { grant, attemptId: attempt.id, outcome }; +} + +type Changes = ReturnType[1]>; +// Only the scheduler supplies transitions. This boundary validates commands and authenticates retries. +export async function clientToolCommand( + store: ConversationStore, + input: unknown, + authorize: ClientToolAuthorizer, + prepare: ( + command: ClientToolCommand, + authority: ClientToolAuthority, + replay: boolean + ) => Promise<{ call: ToolCall; apply: (authority: ClientToolAuthority) => Changes }>, + now: () => number +): Promise { + const envelope = z.object({ commandId: z.uuid(), protocolVersion: z.unknown() }).parse(input); + const parsed = CommandSchema.safeParse(input); + if ( + !parsed.success || + (parsed.data.type !== 'claimClientTool' && parsed.data.type !== 'completeClientTool') + ) + return rejectClientCommand(envelope.commandId, { + code: envelope.protocolVersion === 1 ? 'invalid_input' : 'unsupported_protocol', + message: 'The client tool command is invalid.', + retryable: false, + }); + const command = parsed.data; + let journal: { id: string; fingerprint: string } | undefined; + async function currentAuthority() { + const result = await authorize(command); + if ('error' in result) throw new RuntimeError(ErrorSchema.parse(result.error)); + const authority = AuthoritySchema.parse(result), + current = store.snapshot()?.conversation; + if ( + !current || + authority.client.id !== command.clientId || + authority.client.ownerUserId !== current.ownerUserId || + authority.conversation.id !== current.id || + command.conversationId !== current.id || + authority.conversation.ownerUserId !== current.ownerUserId || + canonicalizeValidatedInput(authority.conversation.context) !== + canonicalizeValidatedInput(current.context) + ) + fail('access_revoked', 'The command has no current client or context authority.'); + return authority; + } + try { + const authority = await currentAuthority(); + const fingerprint = await fingerprintCommand( + { actorUserId: authority.client.ownerUserId, conversationId: command.conversationId }, + command, + text => createHash('sha256').update(text).digest('hex') + ); + // Validate saved calls before replay, but do not prepare another dispatch for a journaled command. + const replay = store.getCommand(command.commandId) !== null; + const prepared = await prepare(command, authority, replay); + const current = await currentAuthority(); + if ( + prepared.call.executionTarget.kind !== 'client' || + prepared.call.executionTarget.clientId !== current.client.id + ) + fail('access_revoked', 'Only the designated client can claim or complete this call.'); + // Revocation and storage loss override replay: never disclose a saved grant to an invalid client. + if ( + current.client.revokedAt === null && + current.storageReady && + supportsClientCall(current, prepared.call) + ) + journal = { id: command.commandId, fingerprint }; + const reply = await store.transition({ command: journal, wakeAt: now() + 1 }, () => + prepared.apply(current) + ); + if (!reply) throw new StoreError('storage_unavailable', true); + return reply; + } catch (error) { + const detail = + error instanceof RuntimeError + ? error.detail + : { + code: + error instanceof StoreError + ? error.code + : error instanceof z.ZodError + ? ('invalid_output' as const) + : ('storage_unavailable' as const), + message: 'The client command could not be committed.', + retryable: + error instanceof StoreError ? error.retryable : !(error instanceof z.ZodError), + }; + const reply = rejectClientCommand(command.commandId, detail); + if (journal && error instanceof RuntimeError) + return store + .transition({ command: journal, wakeAt: null }, () => ({ events: [], reply })) + .then( + stored => stored ?? reply, + () => + rejectClientCommand(command.commandId, { + code: 'storage_unavailable', + message: 'The client rejection could not be committed. Retry the same command.', + retryable: true, + }) + ); + return reply; + } +} diff --git a/services/agent-harness/src/dispatch.ts b/services/agent-harness/src/dispatch.ts index 6a5e60b823..c247a66152 100644 --- a/services/agent-harness/src/dispatch.ts +++ b/services/agent-harness/src/dispatch.ts @@ -1,8 +1,8 @@ import { toolModelMessageSchema } from 'ai'; import { eq } from 'drizzle-orm'; -import type { ToolCall, ToolOutcome } from '@kilocode/agent-harness/contracts'; +import type { ExecutionGrant, ToolCall, ToolOutcome } from '@kilocode/agent-harness/contracts'; import { evaluateDispatch, type DispatchPolicy } from '@kilocode/agent-harness/policy'; -import { compareAndSetCall, insertAttempt, type StoreDatabase } from './db/records'; +import { compareAndSetCall, insertAttempt, insertGrant, type StoreDatabase } from './db/records'; import type { ConversationStore } from './db/store'; import * as s from './db/sqlite-schema'; import { StoreError } from './db/wake'; @@ -16,7 +16,8 @@ export function commitDispatch( proposed: ToolCall, policy: DispatchPolicy, attemptId: string, - generation: number + generation: number, + grant?: ExecutionGrant ) { const decision = evaluateDispatch(stored.data, proposed, policy); db.update(s.calls) @@ -24,7 +25,8 @@ export function commitDispatch( .where(eq(s.calls.id, stored.id)) .run(); if (decision === 'dispatch') { - insertAttempt(db, { id: attemptId, toolCallId: stored.id, generation }); + if (grant) insertGrant(db, grant); + insertAttempt(db, { id: attemptId, toolCallId: stored.id, generation, grantId: grant?.id }); if ( !compareAndSetCall(db, stored.id, stored.revision, { state: 'executing', diff --git a/services/agent-harness/src/scheduler.ts b/services/agent-harness/src/scheduler.ts index d192e70be0..e16b43ab45 100644 --- a/services/agent-harness/src/scheduler.ts +++ b/services/agent-harness/src/scheduler.ts @@ -16,6 +16,16 @@ import { type ToolOutcome, } from '@kilocode/agent-harness/contracts'; import type { DispatchPolicy } from '@kilocode/agent-harness/policy'; +import { bridgeWaitReason } from '@kilocode/agent-harness/bridge'; +import { ExecutionGrantSchema, type ExecutionGrant } from '@kilocode/agent-harness/contracts'; +import { + clientAction, + clientToolCommand, + readClientGrant, + rejectClientCommand, + supportsClientCall, + type ClientToolAuthorizer, +} from './client-tools'; import { commitDispatch, toolResultMessage } from './dispatch'; import { closeInteraction, @@ -321,24 +331,29 @@ export function createScheduler( const step = limits ? readCompleteStep(row.data, adapter.definitions, limits) : CompleteStepSchema.parse(row.data); - return step.calls.map((item, index) => { + return step.calls.flatMap((item, index) => { const stored = calls.find(call => call.id === item.call.id); if (!stored || stored.checkpointId !== row.id) fail('invalid_output', 'The checkpoint has no matching stored call.'); const call = validateStoredCall(stored.data, item.call, adapter.definitions, limits); // A call owns its display message. Up to 32 bounded outputs must not form one oversized event. - return displayMessage( - run, - { - ...step, - kind: 'partial', - messageId: call.id, - createdAt: new Date(Date.parse(step.createdAt) + index + 1).toISOString(), - text: '', - }, - false, - [call] - ); + return [ + displayMessage( + run, + { + ...step, + kind: 'partial', + messageId: call.id, + createdAt: new Date(Date.parse(step.createdAt) + index + 1).toISOString(), + text: '', + }, + false, + [call] + ), + ...(call.executionTarget.kind === 'client' && call.state === 'settled' + ? [clientAction(call, null)] + : []), + ]; }); }); } @@ -415,19 +430,22 @@ export function createScheduler( }; return [ displayMessage(run, display, false, [call]), + ...(call.executionTarget.kind === 'client' + ? [clientAction(call, 'reconciliation', readClientGrant(db, call)?.grant ?? null)] + : []), runEvent(run, { status: 'waiting', waiting: { reason: 'reconciliation', toolCallId: call.id }, }), ]; } + const needsReceipt = (call: ToolCall) => + call.effect !== 'read' || call.executionTarget.kind === 'client'; function stopRun(run: Run, record: SchedulerRecord): EventEnvelope['event'][] { record.data.stopped = true; const reservation = activeReservation(record); const calls = store.callsForRun(run.id); - const mutation = calls.find( - call => call.data.state === 'executing' && call.data.effect !== 'read' - ); + const mutation = calls.find(call => call.data.state === 'executing' && needsReceipt(call.data)); for (const call of calls) { if (call.data.state !== 'settled' && call.id !== mutation?.id) settleCall(call, { status: 'cancelled' }); @@ -441,7 +459,7 @@ export function createScheduler( ...callEvents(run), ]; if (mutation && reservation && reservation.deadline > now()) { - // A supported read can abort. A mutation retains its lease and can report actual late completion. + // Backend reads can abort. Mutations and client calls retain their lease for actual completion. writeScheduler(db, run.id, record); return events; } @@ -510,7 +528,7 @@ export function createScheduler( ) return { events: [] }; if (executing && !reconciling) { - if (executing.data.effect !== 'read') + if (needsReceipt(executing.data)) return { events: unknownOutcome(run, executing.data, 'The dispatch response was lost.'), }; @@ -579,29 +597,7 @@ export function createScheduler( reservation, }; if (pending) { - const checkpoint = checkpointRows.find(row => row.id === pending.checkpointId); - if (!checkpoint || checkpoint.status !== 'complete') - fail('invalid_output', 'A partial cannot authorize dispatch.'); - const complete = readCompleteStep(checkpoint.data, adapter.definitions, admission.limits); - const expected = complete.calls.find(item => item.call.id === pending.id); - if ( - !expected || - expected.call.runId !== run.id || - canonicalizeValidatedInput(expected.call.context) !== - canonicalizeValidatedInput(currentSnapshot.conversation.context) - ) - fail('invalid_output', 'The call is absent from its scoped checkpoint.'); - validateStoredCall(pending.data, expected.call, adapter.definitions, admission.limits); - if ( - pending.inputDigest !== - createHash('sha256') - .update(canonicalizeValidatedInput(pending.data.arguments)) - .digest('hex') || - canonicalizeValidatedInput( - calls.filter(call => call.checkpointId === checkpoint.id).map(call => call.id) - ) !== canonicalizeValidatedInput(complete.calls.map(item => item.call.id)) - ) - fail('invalid_output', 'The stored call digest or order has changed.'); + checkedCall(run, pending, admission); job = { ...common, kind: 'tool', @@ -759,7 +755,7 @@ export function createScheduler( call && (call.data.state === 'executing' || (call.data.approval !== null && !detail.retryable)); if (failedCall) { - if (call.data.state === 'executing' && call.data.effect !== 'read') + if (call.data.state === 'executing' && needsReceipt(call.data)) return { events: unknownOutcome( run, @@ -787,6 +783,89 @@ export function createScheduler( }; }); } + function dispatchCall( + job: Job & { kind: 'tool' }, + policy: DispatchPolicy, + grant?: ExecutionGrant, + // Scheduler-only producers have no client report. Keep unavailable until a claim supplies readiness. + reason: ReturnType = 'unavailable' + ) { + const call = store.callsForRun(job.run.id).find(item => item.id === job.call.id); + const conversation = store.snapshot()?.conversation; + if (!call || !conversation) fail('invalid_output', 'The stored dispatch call is missing.'); + validateStoredCall(call.data, job.call, adapter.definitions, job.admission.limits); + const decision = commitDispatch( + db, + call, + job.call, + { + ...policy, + permissionMode: conversation.permissionMode, + permissionRevision: conversation.permissionRevision, + // Only durable answers and a designated grant release these gates, never adapter hints. + questionAnswered: false, + clientReady: grant !== undefined && reason === null, + }, + job.reservation.id, + job.epoch, + grant + ); + const result = (events: EventEnvelope['event'][]) => ({ decision, events }); + if (decision === 'dispatch') + return result([ + ...closeInteraction(db, store, { ...call.data, state: 'executing' }, 'approve'), + ...callEvents(job.run), + ...(grant ? [clientAction(call.data, null), runEvent(job.run, { status: 'running' })] : []), + ]); + const record = schedulerRecord(db, job.run.id); + // No effect occurred. Release this slot, but retain time spent checking authority. + updateReservation(record, { ...finishReservation(job.reservation, now()), status: 'released' }); + record.data.currentReservationId = null; + writeScheduler(db, job.run.id, record); + if (decision === 'approval' || decision === 'question' || decision === 'client') { + const waiting = { ...call.data, state: 'waiting' as const }; + if (!compareAndSetCall(db, call.id, call.revision, waiting)) + throw new StoreError('command_conflict'); + return result([ + ...(conversation.permissionMode === 'yolo' + ? closeInteraction(db, store, waiting, 'approve') + : []), + ...(decision === 'client' + ? [clientAction(waiting, reason ?? 'unavailable')] + : [waitForInteraction(store, waiting, decision)]), + ...(decision !== 'client' && call.data.executionTarget.kind === 'client' + ? [clientAction(waiting, null)] + : []), + ...callEvents(job.run), + runEvent(job.run, { + status: 'waiting', + waiting: { reason: decision, toolCallId: call.id }, + }), + ]); + } + if (decision === 'stale_revision') return result([runEvent(job.run, { status: 'running' })]); + if (decision === 'already_dispatched') return result([]); + if (decision === 'denied') { + settleCall(call, { status: 'denied' }); + return result(callEvents(job.run)); + } + const error = { + code: + decision === 'access_revoked' + ? ('access_revoked' as const) + : decision === 'unavailable_tool' + ? ('unavailable_tool' as const) + : ('invalid_input' as const), + message: 'The current dispatch authority does not permit this call.', + retryable: false, + }; + settleCall(call, { status: 'failed', error }); + const pendingEvents = failPendingInteraction(job.run, error); + return result([ + ...(pendingEvents.length ? pendingEvents : callEvents(job.run)), + runEvent(job.run, { status: 'failed', error }), + ]); + } async function executeTool(job: Job & { kind: 'tool' }, controller: AbortController) { const reconciliation = job.reconciliation; if (reconciliation) { @@ -822,87 +901,10 @@ export function createScheduler( retryPolicy = false; await store.transition({ wakeAt: now() + 1 }, () => { fence(job); - const call = store.callsForRun(job.run.id).find(item => item.id === job.call.id); - const conversation = store.snapshot()?.conversation; - if (!call || !conversation) fail('invalid_output', 'The stored dispatch call is missing.'); - validateStoredCall(call.data, job.call, adapter.definitions, job.admission.limits); - const decision = commitDispatch( - db, - call, - job.call, - { - ...policy, - permissionMode: conversation.permissionMode, - permissionRevision: conversation.permissionRevision, - // Only durable answers and a14 client grants can release these gates, not adapter hints. - questionAnswered: false, - clientReady: false, - }, - job.reservation.id, - job.epoch - ); - if (decision === 'dispatch') { - dispatched = true; - return { - events: [ - ...closeInteraction(db, store, { ...call.data, state: 'executing' }, 'approve'), - ...callEvents(job.run), - ], - }; - } - const record = schedulerRecord(db, job.run.id); - // No external request occurred. Release this request slot, but retain time spent checking authority. - updateReservation(record, { - ...finishReservation(job.reservation, now()), - status: 'released', - }); - record.data.currentReservationId = null; - writeScheduler(db, job.run.id, record); - if (decision === 'approval' || decision === 'question' || decision === 'client') { - const waiting = { ...call.data, state: 'waiting' as const }; - if (!compareAndSetCall(db, call.id, call.revision, waiting)) - throw new StoreError('command_conflict'); - return { - events: [ - ...(conversation.permissionMode === 'yolo' - ? closeInteraction(db, store, waiting, 'approve') - : []), - ...(decision === 'client' ? [] : [waitForInteraction(store, waiting, decision)]), - ...callEvents(job.run), - runEvent(job.run, { - status: 'waiting', - waiting: { reason: decision, toolCallId: call.id }, - }), - ], - }; - } - if (decision === 'stale_revision') { - retryPolicy = true; - return { events: [runEvent(job.run, { status: 'running' })] }; - } - if (decision === 'already_dispatched') return { events: [] }; - if (decision === 'denied') { - settleCall(call, { status: 'denied' }); - return { events: callEvents(job.run) }; - } - const error = { - code: - decision === 'access_revoked' - ? ('access_revoked' as const) - : decision === 'unavailable_tool' - ? ('unavailable_tool' as const) - : ('invalid_input' as const), - message: 'The current dispatch authority does not permit this call.', - retryable: false, - }; - settleCall(call, { status: 'failed', error }); - const pendingEvents = failPendingInteraction(job.run, error); - return { - events: [ - ...(pendingEvents.length ? pendingEvents : callEvents(job.run)), - runEvent(job.run, { status: 'failed', error }), - ], - }; + const result = dispatchCall(job, policy); + dispatched = result.decision === 'dispatch'; + retryPolicy = result.decision === 'stale_revision'; + return { events: result.events }; }); if (!dispatched) return !retryPolicy; fence(job); @@ -948,26 +950,430 @@ export function createScheduler( updateReservation(record, finishReservation(job.reservation, now())); record.data.currentReservationId = null; writeScheduler(db, run.id, record); - if (outcome.status === 'outcome_unknown') { - if (outcome.providerReference) - db.update(s.attempts) - .set({ providerReference: outcome.providerReference }) - .where(eq(s.attempts.id, attemptId)) - .run(); + return { events: settleToolOutcome(run, record, call, attemptId, outcome, stopping) }; + }); + return committed; + } + function settleToolOutcome( + run: Run, + record: SchedulerRecord, + call: ReturnType[number], + attemptId: string, + outcome: ToolOutcome, + stopping: boolean + ): EventEnvelope['event'][] { + if (outcome.status === 'outcome_unknown') { + if (outcome.providerReference) + db.update(s.attempts) + .set({ providerReference: outcome.providerReference }) + .where(eq(s.attempts.id, attemptId)) + .run(); + return stopping + ? stopRun(run, record) + : unknownOutcome(run, call.data, outcome.reason, outcome.providerReference); + } + db.update(s.attempts) + .set({ outcome: jsonValue(outcome) }) + .where(eq(s.attempts.id, attemptId)) + .run(); + settleCall(call, outcome); + return stopping + ? stopRun(run, record) + : [ + ...callEvents(run), + ...(run.state.status === 'waiting' ? [runEvent(run, { status: 'running' })] : []), + ]; + } + function checkedCall( + run: Run, + call: ReturnType[number], + admission: Admission + ) { + const checkpoint = db + .select() + .from(s.checkpoints) + .where(eq(s.checkpoints.id, call.checkpointId)) + .get(); + if (!checkpoint || checkpoint.status !== 'complete') + fail('invalid_output', 'A partial cannot authorize dispatch.'); + const complete = readCompleteStep(checkpoint.data, adapter.definitions, admission.limits); + const expected = complete.calls.find(item => item.call.id === call.id); + if ( + !expected || + expected.call.runId !== run.id || + canonicalizeValidatedInput(expected.call.context) !== + canonicalizeValidatedInput(store.snapshot()?.conversation.context) || + (expected.call.executionTarget.kind === 'client' && + expected.call.executionTarget.clientId !== run.originClientId) + ) + fail('invalid_output', 'The call is absent from its scoped checkpoint.'); + validateStoredCall(call.data, expected.call, adapter.definitions, admission.limits); + if ( + call.inputDigest !== + createHash('sha256') + .update(canonicalizeValidatedInput(call.data.arguments)) + .digest('hex') || + canonicalizeValidatedInput( + store + .callsForRun(run.id) + .filter(item => item.checkpointId === checkpoint.id) + .map(item => item.id) + ) !== canonicalizeValidatedInput(complete.calls.map(item => item.call.id)) + ) + fail('invalid_output', 'The stored call digest or order has changed.'); + return checkpoint.step; + } + function releaseClientReservation(record: SchedulerRecord, call: ToolCall, interrupted: boolean) { + const reservation = activeReservation(record); + if (reservation?.toolCallId === call.id) { + updateReservation( + record, + interrupted + ? { ...reservation, status: 'interrupted' } + : finishReservation(reservation, now()) + ); + record.data.currentReservationId = null; + } + record.data.epoch++; + writeScheduler(db, call.runId, record); + } + function unavailableClientCall( + run: Run, + record: SchedulerRecord, + call: ReturnType[number], + code: 'offline' | 'access_revoked' | 'storage_unavailable' | 'unavailable_tool' + ): EventEnvelope['event'][] { + if (call.data.state === 'settled') return []; + if (run.state.status === 'stopping') record.data.stopped = true; + releaseClientReservation(record, call.data, call.data.state === 'executing'); + if (record.data.stopped) return stopRun(run, record); + if (call.data.state === 'executing') + return unknownOutcome( + run, + call.data, + `The designated client requires receipt reconciliation: ${code}.` + ); + const waiting = { ...call.data, state: 'waiting' as const }; + if (!compareAndSetCall(db, call.id, call.revision, waiting)) + throw new StoreError('command_conflict'); + return [ + clientAction(waiting, code === 'offline' ? 'offline' : 'unavailable'), + ...callEvents(run), + runEvent(run, { status: 'waiting', waiting: { reason: 'client', toolCallId: call.id } }), + ]; + } + // Trusted registration/lifecycle notifications never authorize an effect or change its target. + async function clientUnavailable( + clientId: string, + code: 'offline' | 'access_revoked' | 'storage_unavailable' + ) { + z.uuid().parse(clientId); + await store.transition({ wakeAt: now() + 1 }, () => { + const run = store.snapshot()?.activeRun; + const call = run && store.callsForRun(run.id).find(item => item.data.state !== 'settled'); + if ( + !run || + !call || + call.data.executionTarget.kind !== 'client' || + call.data.executionTarget.clientId !== clientId + ) + return { events: [] }; + return { events: unavailableClientCall(run, schedulerRecord(db, run.id), call, code) }; + }); + await maintainAlarm(); + } + async function clientTool(input: unknown, authorize: ClientToolAuthorizer) { + const preparation: { budget?: { runId: string; reservation: Reservation } } = {}; + const reply = await clientToolCommand( + store, + input, + authorize, + async (command, authority, replay) => { + const row = db.select().from(s.calls).where(eq(s.calls.id, command.toolCallId)).get(); + if (!row) fail('invalid_input', 'The client call does not exist.'); + const run = storedRun(db, row.runId), + admission = admissionForRun(store, run); + const original = store.callsForRun(run.id).find(item => item.id === row.id); + if ( + !original || + original.data.executionTarget.kind !== 'client' || + original.data.executionTarget.clientId !== authority.client.id + ) + fail('access_revoked', 'Only the designated client can claim or complete this call.'); + // Validate persisted authority before the command journal can return a canonical replay. + checkedCall(run, original, admission); + if (!readClientGrant(db, original.data) && original.data.state === 'executing') + fail('invalid_output', 'The executing client call has no grant.'); + let policy: DispatchPolicy | undefined; + if ( + !replay && + command.type === 'claimClientTool' && + ['pending', 'waiting'].includes(original.data.state) && + authority.client.revokedAt === null && + authority.storageReady && + supportsClientCall(authority, original.data) + ) { + await store.transition({ wakeAt: now() + 1 }, () => { + const currentRun = storedRun(db, run.id); + const call = store.callsForRun(run.id).find(item => item.id === original.id); + if (!call) fail('invalid_output', 'The stored client call is missing.'); + // Another claim can finish while this command enters the gate. Replay its grant below. + if (call.data.state === 'executing' || call.data.state === 'settled') + return { events: [] }; + if ( + store.snapshot()?.activeRun?.id !== run.id || + store.callsForRun(run.id).find(item => item.data.state !== 'settled')?.id !== call.id + ) + fail('command_conflict', 'This call does not own the active run.', true); + const record = schedulerRecord(db, run.id); + if (record.data.stopped || !['running', 'waiting'].includes(currentRun.state.status)) + fail('cancelled', 'Stop precedes this client claim.'); + if (activeReservation(record)) + fail('command_conflict', 'Another scheduler transition owns the current call.', true); + const step = checkedCall(currentRun, call, admission); + const reservation = reserve( + admission, + record.data.reservations, + { kind: 'tool', step, toolCallId: call.id, webRequest: false }, + now() + ); + // A lost preparation keeps its full time reservation, but grants no execution authority. + // Device waits remain idle; only a committed grant takes the active scheduler lease. + record.data.reservations.push(reservation); + writeScheduler(db, run.id, record); + preparation.budget = { runId: run.id, reservation }; + return { events: [] }; + }); + const budget = preparation.budget; + if (budget) { + const signal = AbortSignal.timeout(Math.max(1, budget.reservation.deadline - now())); + const checkDeadline = () => { + if (signal.aborted || now() >= budget.reservation.deadline) + fail('limit_exceeded', 'The client claim exceeded its execution deadline.'); + }; + try { + checkDeadline(); + await abortable(signal, () => adapter.authorize(authority.conversation, run, signal)); + checkDeadline(); + policy = await abortable(signal, () => + adapter.policy( + store.snapshot()?.conversation ?? authority.conversation, + run, + original.data, + signal + ) + ); + checkDeadline(); + } catch (error) { + checkDeadline(); + throw error; + } + } + } return { - events: stopping - ? stopRun(run, record) - : unknownOutcome(run, call.data, outcome.reason, outcome.providerReference), + call: original.data, + apply: currentAuthority => { + const run = storedRun(db, original.runId), + record = schedulerRecord(db, run.id); + const call = store.callsForRun(run.id).find(item => item.id === original.id); + if (!call) fail('invalid_output', 'The stored client call is missing.'); + checkedCall(run, call, admission); + if ( + call.data.state !== 'settled' && + (store.snapshot()?.activeRun?.id !== run.id || + store.callsForRun(run.id).find(item => item.data.state !== 'settled')?.id !== + call.id) + ) + fail('command_conflict', 'This call does not own the active run.'); + const accept = (result: unknown, events: EventEnvelope['event'][] = []) => ({ + events, + reply: { + status: 'accepted' as const, + commandId: command.commandId, + result: jsonValue(result), + }, + }); + const unavailable = + currentAuthority.client.revokedAt !== null + ? 'access_revoked' + : !currentAuthority.storageReady + ? 'storage_unavailable' + : !supportsClientCall(currentAuthority, call.data) + ? 'unavailable_tool' + : null; + if (unavailable) + return { + events: unavailableClientCall(run, record, call, unavailable), + reply: rejectClientCommand(command.commandId, { + code: unavailable, + message: 'The designated client cannot safely dispatch or report a result.', + retryable: unavailable === 'storage_unavailable', + }), + }; + const claimed = readClientGrant(db, call.data); + if (command.type === 'completeClientTool') { + if ( + !claimed || + claimed.grant.id !== command.grantId || + claimed.grant.generation !== command.generation + ) + fail('access_revoked', 'The result does not belong to the current client grant.'); + // Expiry ends dispatch permission, not reconciliation of this original grant's receipt. + const outcome = validateOutcome( + command.result, + call.data, + adapter.definitions, + admission.limits + ); + if (call.data.result) { + if ( + canonicalizeValidatedInput(call.data.result) !== + canonicalizeValidatedInput(outcome) + ) + fail('command_conflict', 'The client completion is immutable.'); + return accept({ toolCall: call.data, result: call.data.result }); + } + if (call.data.state !== 'executing') + fail('access_revoked', 'The grant has no dispatched call.'); + releaseClientReservation(record, call.data, false); + const events = settleToolOutcome( + run, + record, + call, + claimed.attemptId, + outcome, + run.state.status === 'stopping' || record.data.stopped + ); + return accept( + { + toolCall: store.callsForRun(run.id).find(item => item.id === call.id)?.data, + result: outcome, + }, + events + ); + } + if (call.data.state === 'settled') + fail('cancelled', 'The client call is already settled.'); + if (claimed) { + if (call.data.state !== 'executing') + fail('invalid_output', 'The grant has no executing call.'); + if ( + Date.parse(claimed.grant.expiresAt) <= now() || + claimed.outcome !== null || + record.data.stopped || + run.state.status !== 'running' || + bridgeWaitReason(currentAuthority.readiness) + ) + return { + events: unavailableClientCall(run, record, call, 'offline'), + reply: rejectClientCommand(command.commandId, { + code: 'outcome_unknown', + message: 'Reconcile the existing receipt; do not execute again.', + retryable: false, + }), + }; + return accept({ grant: claimed.grant, toolCall: call.data }); + } + if (record.data.stopped || !['running', 'waiting'].includes(run.state.status)) + fail('cancelled', 'Stop precedes this client claim.'); + if (activeReservation(record)) + fail('command_conflict', 'Another scheduler transition owns the current call.', true); + if (!policy) fail('access_revoked', 'The client claim has no current dispatch policy.'); + const reservation = record.data.reservations.find( + item => item.id === preparation.budget?.reservation.id + ); + if (!reservation || reservation.status !== 'reserved') + fail('command_conflict', 'The client claim has no current time reservation.', true); + if (now() >= reservation.deadline) + fail('limit_exceeded', 'The client claim exceeded its execution deadline.'); + record.data.epoch++; + record.data.currentReservationId = reservation.id; + writeScheduler(db, run.id, record); + const conversation = store.snapshot()?.conversation; + if (!conversation) fail('invalid_input', 'The conversation is missing.'); + const grant = ExecutionGrantSchema.parse({ + id: crypto.randomUUID(), + conversationId: conversation.id, + ownerUserId: conversation.ownerUserId, + clientId: command.clientId, + toolCallId: call.id, + context: call.data.context, + definitionVersion: call.data.definitionVersion, + inputDigest: call.inputDigest, + generation: record.data.epoch, + expiresAt: new Date(reservation.deadline).toISOString(), + }); + const result = dispatchCall( + { + kind: 'tool', + run, + conversation, + admission, + epoch: record.data.epoch, + reservation, + call: call.data, + }, + policy, + grant, + bridgeWaitReason(currentAuthority.readiness) + ); + if ( + result.decision === 'stale_revision' || + result.decision === 'access_revoked' || + result.decision === 'unavailable_tool' || + result.decision === 'new_call_required' + ) + return { + events: result.events, + reply: rejectClientCommand(command.commandId, { + code: result.decision === 'new_call_required' ? 'invalid_input' : result.decision, + message: 'The current policy does not permit this client claim.', + retryable: result.decision === 'stale_revision', + }), + }; + return accept( + { + grant: result.decision === 'dispatch' ? grant : null, + toolCall: store.callsForRun(run.id).find(item => item.id === call.id)?.data, + decision: result.decision, + }, + result.events + ); + }, }; + }, + now + ); + const budget = preparation.budget; + if (budget) { + try { + await store.transition({ wakeAt: now() + 1 }, () => { + const record = schedulerRecord(db, budget.runId); + const reservation = record.data.reservations.find( + item => item.id === budget.reservation.id + ); + // Dispatch owns its lease. Failed checks and losing claims release only their own preparation. + if ( + reservation?.status === 'reserved' && + record.data.currentReservationId !== reservation.id + ) { + updateReservation(record, { + ...finishReservation(reservation, now()), + status: 'released', + }); + writeScheduler(db, budget.runId, record); + } + return { events: [] }; + }); + } catch { + return rejectClientCommand(reply.commandId, { + code: 'storage_unavailable', + message: 'The client claim budget could not be settled. Retry the same command.', + retryable: true, + }); } - db.update(s.attempts) - .set({ outcome: jsonValue(outcome) }) - .where(eq(s.attempts.id, attemptId)) - .run(); - settleCall(call, outcome); - return { events: stopping ? stopRun(run, record) : callEvents(run) }; - }); - return committed; + } + return reply; } async function execute(job: Job) { const controller = new AbortController(); @@ -1137,5 +1543,5 @@ export function createScheduler( if (job) await execute(job); await maintainAlarm(); } - return { alarm, interrupt, resolveInteraction, reconcile }; + return { alarm, interrupt, resolveInteraction, reconcile, clientTool, clientUnavailable }; }