diff --git a/apps/web/src/lib/agent-harness/operation-maintenance.test.ts b/apps/web/src/lib/agent-harness/operation-maintenance.test.ts new file mode 100644 index 0000000000..c7ee379abc --- /dev/null +++ b/apps/web/src/lib/agent-harness/operation-maintenance.test.ts @@ -0,0 +1,369 @@ +import { beforeEach, expect, it, jest } from '@jest/globals'; +import { createHash } from 'node:crypto'; +import jwt from 'jsonwebtoken'; +import { TRPCError } from '@trpc/server'; +import { QuickChatAuthorityError } from '@kilocode/db/quick-chat-runtime'; +import { + access, + authority, + call, + capability, + conversationId, + operationId, + originalTime, + primary, + runId, + runtime, + toolCallId, +} from './operation-test-fixture'; +import type * as Contract from './operation-contract'; +import type * as Maintenance from './operation-maintenance'; +import type { drainLegacyHistoryWithProgress } from './history'; + +const { HarnessOperationSchema, harnessOperationFailure } = + jest.requireActual('./operation-contract'); +const { executeHarnessMaintenance, retirement } = + jest.requireActual('./operation-maintenance'); +const effects = new Map(); +const identity = { conversationId, operationId }; +const projection = { + id: toolCallId, + key: `agent-harness:${conversationId}:${toolCallId}`, + role: 'assistant', + content: 'Actual text', + clientId: null, + createdAt: new Date(originalTime).toISOString(), +}; +const inputs = [ + { type: 'read', ...identity, purpose: 'read' }, + { type: 'history', ...identity, limit: 1 }, + { type: 'projection', ...identity, projection }, +]; +const purge = { type: 'retirement', ...identity, generation: 0 }; +const now = originalTime / 1000; +function purgeToken(changes: Record = {}) { + const request = { type: 'purge', protocolVersion: 1, threadId: conversationId, generation: 0 }; + return jwt.sign( + JSON.stringify({ + operation: 'purge', + threadId: conversationId, + generation: 0, + dispatchId: operationId, + inputDigest: createHash('sha256').update(JSON.stringify(request)).digest('hex'), + iss: 'agent-harness', + aud: 'agent-harness:maintenance', + iat: now, + exp: now + 60, + ...changes, + }), + 'test-signing-key' + ); +} +let fence: { thread_id: string; generation: number } | undefined; +const fences = { findFirst: async () => fence }; +beforeEach(() => { + fence = { thread_id: conversationId, generation: 0 }; + Object.assign(authority, { userId: 'oauth/owner', organizationId: runId }); + effects.clear(); + Object.assign(primary.query, { agent_harness_retirements: fences }); + Object.assign(runtime, { + claimPending: async () => { + effects.set('history', true); + return []; + }, + hasPending: async () => false, + projectText: async (owner: unknown, value: { id: string; content: string }) => { + effects.set(value.id, { owner, content: value.content }); + return value.id; + }, + }); +}); +const invoke = async (raw: unknown, token = capability(raw)) => { + try { + const input = HarnessOperationSchema.parse(JSON.parse(JSON.stringify(raw))); + if (input.type === 'retirement') + return { status: 200, body: { result: await retirement(input, token) } as any }; + if (input.type !== 'read' && input.type !== 'history' && input.type !== 'projection') + throw new TRPCError({ code: 'BAD_REQUEST' }); + return { status: 200, body: (await executeHarnessMaintenance(input, token)) as any }; + } catch (error) { + return { status: 400, body: harnessOperationFailure(error) as any }; + } +}; + +it.each([null, runId])( + 'authorizes history and bound projection in context %s', + async organizationId => { + Object.assign(authority, { organizationId }); + expect((await invoke(inputs[0])).body).toEqual({ result: authority }); + expect((await invoke(inputs[1])).body).toEqual({ + result: { deliveries: [], backlog: 'drained' }, + }); + expect((await invoke(inputs[2])).body).toEqual({ result: toolCallId }); + expect(effects.get(toolCallId)).toEqual({ owner: authority, content: 'Actual text' }); + } +); +it.each<[string, () => unknown]>([ + ['expired token', () => jest.spyOn(Date, 'now').mockReturnValue(originalTime + 60000)], + ['expired grant', () => Object.assign(access, { expires: new Date(originalTime).toISOString() })], + [ + 'revoked grant', + async () => { + const grants = primary.query.agent_harness_conversation_grants; + const grant = await grants.findFirst(); + jest.spyOn(grants, 'findFirst').mockResolvedValue({ ...grant, revoked_at: 'revoked' } as any); + }, + ], + ['removed membership', () => Object.assign(access, { role: false })], + ['missing primary authority', () => Object.assign(access, { active: false })], +])('denies %s before protected access', async (_name, invalidate) => { + const requests = inputs.map(input => ({ input, token: capability(input) })); + await invalidate(); + for (const { input, token } of requests) + expect((await invoke(input, token)).body.error).toMatchObject({ + code: 'access_revoked', + retryable: false, + }); + expect(effects.size).toBe(0); +}); +it.each(inputs)('rejects changed $type scope and caller owner substitution', async input => { + const token = capability(input); + const changed = + input.type === 'read' + ? { purpose: 'import' } + : input.type === 'history' + ? { limit: 2 } + : { projection: { ...projection, content: 'forged' } }; + for (const fields of [ + { conversationId: runId }, + { operationId: runId }, + { ownerUserId: 'forged' }, + changed, + ]) + expect((await invoke({ ...input, ...fields }, token)).status).not.toBe(200); + expect(effects.size).toBe(0); +}); +it.each([{ userId: 'forged' }, { organizationId: null }, { generation: 1 }, { threadId: runId }])( + 'rejects validly signed substituted authority %j', + async changed => { + for (const input of inputs) { + const claims = jwt.decode(capability(input)) as jwt.JwtPayload; + const token = jwt.sign( + { ...claims, authority: { ...authority, ...changed } }, + 'test-signing-key' + ); + expect((await invoke(input, token)).body.error.code).toBe('access_revoked'); + } + expect(effects.size).toBe(0); + } +); +it('keeps an empty leased history batch pending', async () => { + Object.assign(runtime, { hasPending: async () => true }); + expect((await invoke(inputs[1])).body).toEqual({ + result: { deliveries: [], backlog: 'pending' }, + }); +}); +it.each(['acknowledged', 'retry', 'rejected'])('preserves history delivery %s', async status => { + Object.assign(runtime, { + claimPending: async () => [{ ...authority, ...projection, leaseToken: operationId }], + withClaim: async () => { + if (status === 'rejected') throw new QuickChatAuthorityError(); + return status === 'acknowledged'; + }, + hasPending: async () => status !== 'acknowledged', + }); + expect((await invoke(inputs[1])).body).toEqual({ + result: { + deliveries: [{ id: toolCallId, status }], + backlog: status === 'acknowledged' ? 'drained' : 'pending', + }, + }); +}); +it.each<[string, Record | Error, 'acknowledged' | 'retry']>([ + ['matching receipt', {}, 'acknowledged'], + ['non-durable receipt', { durable: false }, 'retry'], + ['wrong message receipt', { messageId: runId }, 'retry'], + ['wrong thread receipt', { threadId: runId }, 'retry'], + ['wrong owner receipt', { userId: 'oauth/other' }, 'retry'], + ['wrong context receipt', { organizationId: null }, 'retry'], + ['wrong generation receipt', { generation: 1 }, 'retry'], + ['transport failure', new Error('secret-transport-details'), 'retry'], +])('delivers signed history with %s', async (_name, reply, status) => { + jest.replaceProperty(process, 'env', { + AGENT_HARNESS_API_URL: 'https://harness.example/ignored', + }); + const claim = { + ...authority, + ...projection, + createdAt: '2026-04-29 01:16:12.945+00', + leaseToken: operationId, + }; + const pending = new Map([[claim.id, claim]]); + Object.assign(runtime, { + claimPending: async () => [...pending.values()], + withClaim: async (claimed, work) => work(async () => pending.delete(claimed.id)), + hasPending: async () => pending.size > 0, + } satisfies Parameters[0]); + const sent = Promise.withResolvers(); + const response = Promise.withResolvers(); + jest.spyOn(globalThis, 'fetch').mockImplementation((url, init) => { + sent.resolve(new Request(url, init)); + return response.promise; + }); + const delivery = invoke(inputs[1]); + const request = await Promise.race([ + sent.promise, + delivery.then(() => { + throw new Error('History finished before signed delivery'); + }), + ]); + const pendingBeforeReceipt = [...pending.values()]; + if (reply instanceof Error) response.reject(reply); + else + response.resolve( + Response.json({ ...authority, messageId: toolCallId, durable: true, ...reply }) + ); + expect(await delivery).toEqual({ + status: 200, + body: { + result: { + deliveries: [{ id: toolCallId, status }], + backlog: status === 'acknowledged' ? 'drained' : 'pending', + }, + }, + }); + expect(pendingBeforeReceipt).toEqual([claim]); + expect([...pending.values()]).toEqual(status === 'acknowledged' ? [] : [claim]); + expect(request.url).toBe('https://harness.example/internal/maintenance'); + expect(request.method).toBe('POST'); + expect(request.redirect).toBe('error'); + expect(request.headers.get('accept')).toBe('application/json'); + expect(request.headers.get('content-type')).toBe('application/json'); + expect(request.headers.get('x-internal-api-key')).toBe('test-service-key'); + const body = await request.text(); + expect(JSON.parse(body)).toEqual({ + type: 'importLegacy', + protocolVersion: 1, + authority, + message: { + id: toolCallId, + role: 'assistant', + content: 'Actual text', + clientId: null, + createdAt: '2026-04-29T01:16:12.945Z', + provenance: 'legacy', + parts: [{ type: 'text', text: 'Actual text' }], + }, + }); + expect( + jwt.verify(request.headers.get('authorization')?.slice(7) ?? '', 'test-signing-key', { + algorithms: ['HS256'], + issuer: 'agent-harness', + audience: 'agent-harness:maintenance', + }) + ).toEqual({ + ...authority, + operation: 'importLegacy', + dispatchId: toolCallId, + inputDigest: createHash('sha256').update(body).digest('hex'), + iss: 'agent-harness', + aud: 'agent-harness:maintenance', + iat: now, + exp: now + 60, + }); +}); +it.each([{ ids: [toolCallId, runId] }, { ids: ['invalid-id'] }])( + 'rejects invalid history delivery output $ids', + async ({ ids }) => { + Object.assign(runtime, { + claimPending: async () => ids.map(id => ({ ...authority, ...projection, id })), + withClaim: async () => true, + }); + expect((await invoke(inputs[1])).body.error.code).toBe('invalid_output'); + } +); +it.each<[Partial, string]>([ + [{ key: `agent-harness:${runId}:${toolCallId}` }, 'invalid_input'], + [{ key: `agent-harness:${conversationId}:${runId}` }, 'invalid_input'], + [{ content: 'é'.repeat(32768) }, 'limit_exceeded'], +])('rejects invalid projection before writing (case %#)', async (changed, code) => { + expect( + (await invoke({ ...inputs[2], projection: { ...projection, ...changed } })).body.error.code + ).toBe(code); + expect(effects.size).toBe(0); +}); +it.each(['invalid-id', runId])('rejects a mismatched projection result %s', async id => { + Object.assign(runtime, { projectText: async () => id }); + expect((await invoke(inputs[2])).body.error.code).toBe('invalid_output'); +}); +it('bounds authority output in UTF-8 bytes', async () => { + authority.userId = 'é'.repeat(32768); + expect((await invoke(inputs[0])).body.error.code).toBe('limit_exceeded'); +}); +it.each(['claimPending', 'hasPending', 'projectText'])( + 'denies late primary revocation in %s', + async method => { + Object.assign(runtime, { + [method]: async () => { + throw new QuickChatAuthorityError(); + }, + }); + expect((await invoke(inputs[method === 'projectText' ? 2 : 1])).body.error).toMatchObject({ + code: 'access_revoked', + retryable: false, + }); + } +); +it.each(inputs)('sanitizes primary storage failures for $type', async input => { + jest + .spyOn(primary.query.agent_harness_conversation_grants, 'findFirst') + .mockRejectedValue(new Error('secret-database-url')); + expect((await invoke(input)).body).toEqual({ + error: { + code: 'storage_unavailable', + message: 'The operation service is unavailable. Retry synchronization.', + retryable: true, + }, + }); + expect(effects.size).toBe(0); +}); +it('accepts an exact fenced purge without a live grant and rejects altered requests', async () => { + const token = purgeToken(); + access.expires = new Date(originalTime).toISOString(); + access.active = access.role = false; + expect((await invoke(purge, token)).body).toEqual({ result: { retired: true } }); + for (const changed of [{ ...purge, generation: 1 }, { ...purge, operationId: runId }, call()]) + expect((await invoke(changed, token)).status).not.toBe(200); +}); +it.each([ + { operation: 'importLegacy' }, + { threadId: runId }, + { generation: 1 }, + { dispatchId: runId }, + { inputDigest: '0'.repeat(64) }, + { iss: 'other' }, + { aud: 'agent-harness:operations' }, + { aud: ['agent-harness:maintenance', 'other'] }, + { iat: now + 1 }, + { iat: now - 61 }, + { exp: now }, + { exp: now + 61 }, + { exp: undefined }, + { iat: undefined }, + { authority }, +])('rejects invalid signed purge claims %j before fence access', async changed => { + const lookup = jest.spyOn(fences, 'findFirst'); + expect((await invoke(purge, purgeToken(changed))).body.error).toMatchObject({ + code: 'access_revoked', + retryable: false, + }); + expect(lookup).not.toHaveBeenCalled(); +}); +it.each([ + undefined, + { thread_id: runId, generation: 0 }, + { thread_id: conversationId, generation: 1 }, +])('requires the exact primary purge fence %j', async value => { + fence = value; + expect((await invoke(purge, purgeToken())).body.error.code).toBe('access_revoked'); +}); diff --git a/apps/web/src/lib/agent-harness/operation-maintenance.ts b/apps/web/src/lib/agent-harness/operation-maintenance.ts new file mode 100644 index 0000000000..4ab859926c --- /dev/null +++ b/apps/web/src/lib/agent-harness/operation-maintenance.ts @@ -0,0 +1,138 @@ +import 'server-only'; +import { createHash } from 'node:crypto'; +import jwt from 'jsonwebtoken'; +import { z } from 'zod'; +import { TRPCError } from '@trpc/server'; +import { and, eq } from 'drizzle-orm'; +import { + createQuickChatRuntime, + QuickChatAuthorityError, + QuickChatAuthoritySchema, +} from '@kilocode/db/quick-chat-runtime'; +import { agent_harness_retirements } from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { INTERNAL_API_SECRET, NEXTAUTH_SECRET } from '@/lib/config.server'; +import { + Id, + Time, + bounded, + harnessOperationScope, + invalid, + type HarnessOperation, +} from './operation-contract'; +import { authorizeHarnessCapability } from './authorization'; +import { drainLegacyHistoryWithProgress } from './history'; +import { sendHarnessMaintenance } from './retirement'; + +export async function retirement( + input: Extract, + token: string +) { + // A deleted account has no grant. Only the exact signed purge request can check its primary fence. + const request = { + type: 'purge', + protocolVersion: 1, + threadId: input.conversationId, + generation: input.generation, + }; + try { + const claims = z + .strictObject({ + iss: z.literal('agent-harness'), + aud: z.literal('agent-harness:maintenance'), + operation: z.literal('purge'), + threadId: z.literal(input.conversationId), + generation: z.literal(input.generation), + dispatchId: z.literal(input.operationId), + inputDigest: z.literal(createHash('sha256').update(JSON.stringify(request)).digest('hex')), + iat: Time, + exp: Time, + }) + .parse( + jwt.verify(token, NEXTAUTH_SECRET, { + algorithms: ['HS256'], + issuer: 'agent-harness', + audience: 'agent-harness:maintenance', + maxAge: 60, + }) + ); + if ( + claims.iat > Math.floor(Date.now() / 1000) || + claims.exp <= claims.iat || + claims.exp - claims.iat > 60 + ) + throw new Error('Invalid maintenance lifetime'); + } catch { + throw new TRPCError({ code: 'FORBIDDEN' }); + } + const fence = await db.query.agent_harness_retirements.findFirst({ + columns: { thread_id: true, generation: true }, + where: and( + eq(agent_harness_retirements.thread_id, input.conversationId), + eq(agent_harness_retirements.generation, input.generation) + ), + }); + if (!fence || fence.thread_id !== input.conversationId || fence.generation !== input.generation) + throw new TRPCError({ code: 'FORBIDDEN' }); + return { retired: true }; +} + +export async function executeHarnessMaintenance( + input: Extract, + token: string +) { + try { + const { authority } = await authorizeHarnessCapability( + token, + harnessOperationScope(bounded(input)) + ); + const source = createQuickChatRuntime(db); + switch (input.type) { + case 'read': + return { result: bounded(QuickChatAuthoritySchema.strict().parse(authority)) }; + case 'history': { + const progress = await drainLegacyHistoryWithProgress( + source, + work => + sendHarnessMaintenance( + process.env.AGENT_HARNESS_API_URL, + NEXTAUTH_SECRET, + INTERNAL_API_SECRET, + { type: 'importLegacy', protocolVersion: 1, ...work }, + work.message.id + ), + { authority, limit: input.limit } + ); + return { + result: bounded( + z + .strictObject({ + deliveries: z + .array( + z.strictObject({ + id: Id, + status: z.enum(['acknowledged', 'retry', 'rejected']), + }) + ) + .max(input.limit), + backlog: z.enum(['pending', 'drained']), + }) + .parse(progress) + ), + }; + } + case 'projection': { + if (input.projection.key !== `agent-harness:${input.conversationId}:${input.projection.id}`) + invalid(); + return { + result: Id.pipe(z.literal(input.projection.id)).parse( + await source.projectText(authority, input.projection) + ), + }; + } + } + } catch (error) { + if (error instanceof QuickChatAuthorityError) throw new TRPCError({ code: 'FORBIDDEN' }); + throw error; + } +}