From 69a084eb587c8fce2c74e05b0c42f309dd172116 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Sun, 30 Aug 2026 12:08:34 +0200 Subject: [PATCH] feat(agent-harness): define closed internal operation contracts --- .../agent-harness/operation-contract.test.ts | 148 +++++++++++++++++ .../lib/agent-harness/operation-contract.ts | 149 ++++++++++++++++++ .../agent-harness/operation-test-fixture.ts | 116 ++++++++++++++ 3 files changed, 413 insertions(+) create mode 100644 apps/web/src/lib/agent-harness/operation-contract.test.ts create mode 100644 apps/web/src/lib/agent-harness/operation-contract.ts create mode 100644 apps/web/src/lib/agent-harness/operation-test-fixture.ts diff --git a/apps/web/src/lib/agent-harness/operation-contract.test.ts b/apps/web/src/lib/agent-harness/operation-contract.test.ts new file mode 100644 index 0000000000..ac3a4ab3d4 --- /dev/null +++ b/apps/web/src/lib/agent-harness/operation-contract.test.ts @@ -0,0 +1,148 @@ +import { expect, it, jest } from '@jest/globals'; +import { TRPCError } from '@trpc/server'; +import { ErrorSchema } from '@kilocode/agent-harness/contracts'; +import * as fixture from './operation-test-fixture'; +import type * as Contract from './operation-contract'; + +const { + HarnessOperationSchema, + harnessOperationScope, + harnessOperationFailure, + safeError, + bounded, +} = jest.requireActual('./operation-contract'); +const { call, capability, authorizedInput, conversationId, operationId, runId } = fixture; +const identity = { conversationId, operationId }; +const projection = { + id: runId, + key: 'projection', + role: 'assistant', + content: '', + createdAt: '2026-08-30T00:00:00.000Z', +}; +it.each([ + call(), + { ...call(), type: 'reconcile' }, + { ...identity, type: 'read', purpose: 'read' }, + { ...identity, type: 'history' }, + { ...identity, type: 'projection', projection }, + { ...identity, type: 'retirement', generation: 0 }, +])('authorizes strict $type wire input', async input => { + await expect(authorizedInput(input)).resolves.toMatchObject({ authority: fixture.authority }); + expect(HarnessOperationSchema.safeParse({ ...input, userId: 'forged' }).success).toBe(false); +}); +it.each([ + [{ ...identity, type: 'history', limit: undefined }, { limit: 50 }], + [ + call('web.search', { query: 'kilo', limit: undefined }), + { request: { arguments: { limit: 5 } } }, + ], + [{ ...identity, type: 'projection', projection }, { projection: { clientId: null } }], +])('preserves schema defaults through signed JSON %#', async (input, expected) => { + await expect(authorizedInput(input)).resolves.toMatchObject({ input: expected }); +}); +it('authorizes absent optional fields without inventing a dispatch timestamp', async () => { + const input = { + ...call('kilo.sessions.start', { prompt: 'Build', modelId: 'model', repository: undefined }), + dispatchStartedAt: undefined, + reservation: undefined, + }; + const result = await authorizedInput(input); + expect(harnessOperationScope(input)).toEqual(result.scope); + expect(result.input).not.toHaveProperty('dispatchStartedAt'); + expect(result.input).not.toHaveProperty('reservation'); + expect(result.input).toMatchObject({ + request: { arguments: { prompt: 'Build', modelId: 'model' } }, + }); +}); +it('authorizes reordered arguments with the original capability', async () => { + const token = capability( + call('kilo.invite', { recipient: 'member@example.com', role: 'member' }) + ); + const input = call('kilo.invite', { role: 'member', recipient: 'member@example.com' }); + await expect(authorizedInput(input, token)).resolves.toMatchObject({ + authority: fixture.authority, + }); +}); +it.each([ + { ...call(), type: 'arbitrary.trpc' }, + { ...call(), userId: 'forged' }, + { ...call(), userId: undefined }, + call('arbitrary.trpc'), + call('kilo.organizations', { userId: 'forged' }), + { ...call(), dispatchStartedAt: NaN }, + { ...call(), reservation: {} }, + { ...identity, type: 'history', limit: 51 }, + { ...identity, type: 'read', purpose: 'execute' }, + { ...identity, type: 'retirement', generation: -1 }, + { ...identity, type: 'projection', projection: { ...projection, token: 'secret' } }, +])('rejects malformed input before wire normalization %#', input => { + expect(HarnessOperationSchema.safeParse(input).success).toBe(false); + expect(() => harnessOperationScope(input)).toThrow(); +}); +it.each([ + { type: 'reconcile' }, + { conversationId: runId }, + { operationId: runId }, + { runId: operationId }, + { toolCallId: operationId }, + { dispatchStartedAt: fixture.originalTime + 1 }, + { request: { name: 'kilo.members', arguments: {} } }, + { request: { name: 'web.search', arguments: { query: 'changed' } } }, +])('rejects changed signed operation input %#', async patch => { + const input = call('web.search', { query: 'kilo' }); + await expect(authorizedInput({ ...input, ...patch }, capability(input))).rejects.toMatchObject({ + code: 'FORBIDDEN', + }); +}); +it.each([ + ['SERVICE_UNAVAILABLE', 'unavailable_tool', true, false], + ['PRECONDITION_FAILED', 'reauthorization_required', false, false], + ['BAD_REQUEST', 'invalid_input', false, false], + ['FORBIDDEN', 'access_revoked', false, false], + ['TIMEOUT', 'storage_unavailable', true, false], + ['TIMEOUT', 'outcome_unknown', false, true], + ['PAYLOAD_TOO_LARGE', 'limit_exceeded', false, false], + ['PAYLOAD_TOO_LARGE', 'outcome_unknown', false, true], + ['BAD_REQUEST', 'invalid_input', false, true], +] as const)( + 'sanitizes %s without merging recovery states %#', + (code, expected, retryable, uncertain) => { + const failure = new TRPCError({ + code, + message: code === 'PRECONDITION_FAILED' ? 'reauthorization_required' : 'provider-secret', + cause: new Error('secret'), + }); + const result = harnessOperationFailure(failure, uncertain); + expect(result).toMatchObject({ error: { code: expected, retryable } }); + expect(JSON.stringify(result)).not.toContain('secret'); + } +); +it.each(ErrorSchema.shape.code.options)('redacts extra fields from %s errors', code => { + const error = { code, message: 'secret', retryable: false, credentials: 'secret' }; + const result = safeError(error); + expect(result).toEqual({ code, message: expect.any(String), retryable: false }); + expect(JSON.stringify(result)).not.toContain('secret'); +}); +it('bounds UTF-8 output including JSON overhead', () => { + expect(bounded('é', 4)).toBe('é'); + expect(() => bounded('é', 3)).toThrow(expect.objectContaining({ code: 'PAYLOAD_TOO_LARGE' })); + expect(bounded('a'.repeat(65534))).toHaveLength(65534); + expect(() => bounded('a'.repeat(65535))).toThrow( + expect.objectContaining({ code: 'PAYLOAD_TOO_LARGE' }) + ); +}); +it.each([[], {}, null, ''])('preserves empty output %j', value => { + expect(bounded(value)).toEqual(value); +}); +it.each([undefined, 1n])('classifies non-JSON output %# as invalid, not retryable', value => { + let failure: unknown; + try { + bounded(value); + } catch (error) { + failure = error; + } + expect(harnessOperationFailure(failure)).toMatchObject({ + error: { code: 'invalid_output', retryable: false }, + }); +}); diff --git a/apps/web/src/lib/agent-harness/operation-contract.ts b/apps/web/src/lib/agent-harness/operation-contract.ts new file mode 100644 index 0000000000..b9989bba0d --- /dev/null +++ b/apps/web/src/lib/agent-harness/operation-contract.ts @@ -0,0 +1,149 @@ +import 'server-only'; +import { z } from 'zod'; +import { TRPCError } from '@trpc/server'; +import type { ErrorSchema } from '@kilocode/agent-harness/contracts'; +import { ToolRequestSchema } from '@kilocode/agent-harness/tools'; +import { harnessInputDigest } from './authorization'; + +export const Id = z.uuid().transform(value => value.toLowerCase()); +export const Time = z.int().nonnegative(); +const identity = { conversationId: Id, operationId: Id }; +// The Worker supplies these fields from the committed reservation and its owning run checkpoint. +const WebReservation = z.strictObject({ + id: Id, + runId: Id, + toolCallId: Id, + startedAt: Time, + deadline: Time, + kind: z.literal('tool'), + status: z.literal('reserved'), + webRequest: z.literal(true), +}); +export const HarnessOperationSchema = z.discriminatedUnion('type', [ + z.strictObject({ + type: z.literal(['execute', 'reconcile']), + ...identity, + runId: Id, + toolCallId: Id, + request: ToolRequestSchema, + // Legacy reconciliation omits this time; retain absence until those stored attempts are gone. + dispatchStartedAt: Time.optional(), + reservation: WebReservation.optional(), + }), + z.strictObject({ + type: z.literal('read'), + ...identity, + purpose: z.enum(['read', 'import', 'project', 'drain']), + }), + z.strictObject({ + type: z.literal('history'), + ...identity, + limit: z.int().min(1).max(50).default(50), + }), + z.strictObject({ + type: z.literal('projection'), + ...identity, + projection: z.strictObject({ + id: Id, + key: z.string().min(1), + role: z.enum(['user', 'assistant']), + content: z.string(), + clientId: z.string().nullable().default(null), + createdAt: z.iso.datetime(), + }), + }), + z.strictObject({ type: z.literal('retirement'), ...identity, generation: Time }), +]); +export type HarnessOperation = z.infer; +export function harnessOperationScope(raw: unknown) { + const input = HarnessOperationSchema.parse(raw); + return { + audience: 'agent-harness:operations', + conversationId: input.conversationId, + operation: input.type, + definitionVersion: '1', + // Validate before serialization so malformed fields cannot disappear from the signed input. + inputDigest: harnessInputDigest(JSON.parse(JSON.stringify(input))), + dispatchId: input.operationId, + target: { kind: 'backend' as const }, + }; +} +export const messages = { + stale_revision: 'Refresh the conversation before continuing.', + command_conflict: 'This operation has different recorded input.', + access_revoked: 'Access to this conversation is unavailable.', + retired: 'This conversation is retired.', + storage_unavailable: 'The operation service is unavailable. Retry synchronization.', + unsupported_protocol: 'Update the client before continuing.', + unavailable_tool: 'This tool is unavailable in the current context.', + reauthorization_required: 'Reconnect the integration in this context.', + invalid_input: 'The operation input is invalid.', + invalid_output: 'The operation returned invalid output.', + limit_exceeded: 'The operation exceeds its limit.', + cancelled: 'The operation was cancelled.', + outcome_unknown: 'Check the recorded outcome; do not repeat this mutation.', +} satisfies Record['code'], string>; +export const safeError = (error: z.infer) => ({ + code: error.code, + message: messages[error.code], + retryable: error.retryable, +}); +export function harnessOperationFailure(error: unknown, uncertain = false) { + const codes: Partial['code']>> = { + FORBIDDEN: 'access_revoked', + UNAUTHORIZED: 'access_revoked', + BAD_REQUEST: 'invalid_input', + PRECONDITION_FAILED: 'unavailable_tool', + NOT_FOUND: 'unavailable_tool', + CONFLICT: 'command_conflict', + PAYMENT_REQUIRED: 'limit_exceeded', + PAYLOAD_TOO_LARGE: 'limit_exceeded', + UNPROCESSABLE_CONTENT: 'invalid_output', + SERVICE_UNAVAILABLE: 'unavailable_tool', + }; + const rejection = + error instanceof TRPCError && + [ + 'FORBIDDEN', + 'UNAUTHORIZED', + 'BAD_REQUEST', + 'PRECONDITION_FAILED', + 'PAYMENT_REQUIRED', + 'CONFLICT', + ].includes(error.code); + const code = + uncertain && !rejection + ? 'outcome_unknown' + : error instanceof TRPCError + ? error.code === 'PRECONDITION_FAILED' && error.message === 'reauthorization_required' + ? 'reauthorization_required' + : (codes[error.code] ?? 'storage_unavailable') + : error instanceof z.ZodError + ? 'invalid_output' + : 'storage_unavailable'; + return { + error: safeError({ + code, + message: '', + retryable: + code === 'storage_unavailable' || + (code === 'unavailable_tool' && + error instanceof TRPCError && + error.code === 'SERVICE_UNAVAILABLE'), + }), + }; +} +export function bounded(value: T, limit = 64 * 1024): T { + let json: string | undefined; + try { + json = JSON.stringify(value); + } catch { + throw new TRPCError({ code: 'UNPROCESSABLE_CONTENT' }); + } + if (json === undefined) throw new TRPCError({ code: 'UNPROCESSABLE_CONTENT' }); + if (Buffer.byteLength(json, 'utf8') > limit) throw new TRPCError({ code: 'PAYLOAD_TOO_LARGE' }); + return value; +} +export const invalid = () => { + throw new TRPCError({ code: 'BAD_REQUEST' }); +}; diff --git a/apps/web/src/lib/agent-harness/operation-test-fixture.ts b/apps/web/src/lib/agent-harness/operation-test-fixture.ts new file mode 100644 index 0000000000..a0ece142b9 --- /dev/null +++ b/apps/web/src/lib/agent-harness/operation-test-fixture.ts @@ -0,0 +1,116 @@ +import { afterEach, beforeEach, jest } from '@jest/globals'; +import jwt from 'jsonwebtoken'; +import { TRPCError } from '@trpc/server'; +import type * as Authorization from './authorization'; +import type * as Contract from './operation-contract'; +import type * as Runtime from '@kilocode/db/quick-chat-runtime'; + +export const conversationId = '11111111-1111-4111-8111-111111111111'; +export const operationId = '22222222-2222-4222-8222-222222222222'; +export const runId = '33333333-3333-4333-8333-333333333333'; +export const toolCallId = '44444444-4444-4444-8444-444444444444'; +export const originalTime = 1788000000000; +export const authority = { + threadId: conversationId, + userId: 'oauth/owner', + organizationId: runId, + generation: 0, +}; +export const access = { + active: true, + role: true, + expires: new Date(originalTime + 3600000).toISOString(), +}; +export const primary = { + query: { + agent_harness_conversation_grants: { + findFirst: async () => ({ + id: runId, + thread_id: conversationId, + user_id: authority.userId, + generation: 0, + revoked_at: null, + expires_at: access.expires, + }), + }, + agent_harness_conversation_registry: { + findFirst: async () => ({ + thread_id: conversationId, + user_id: authority.userId, + organization_id: authority.organizationId, + generation: 0, + }), + }, + kilocode_users: { findFirst: async () => ({ id: authority.userId, blocked_reason: null }) }, + }, + update: () => ({ set: () => ({ where: async () => undefined }) }), +}; +export const runtime = { lookupThread: async () => (access.active ? authority : null) }; +jest.mock('@/lib/config.server', () => ({ + NEXTAUTH_SECRET: 'test-signing-key', + INTERNAL_API_SECRET: 'test-service-key', +})); +jest.mock('./clients', () => ({ + harnessAccessDenied: () => { + throw new TRPCError({ code: 'FORBIDDEN' }); + }, +})); +jest.mock('@/routers/organizations/utils', () => ({ + ensureOrganizationAccess: async () => { + if (!access.role) throw new TRPCError({ code: 'UNAUTHORIZED', message: 'secret-role-details' }); + return 'owner'; + }, +})); +jest.mock('@/lib/drizzle', () => ({ db: primary })); +jest.mock('@kilocode/db/quick-chat-runtime', () => ({ + ...jest.requireActual('@kilocode/db/quick-chat-runtime'), + createQuickChatRuntime: () => runtime, +})); +const { authorizeHarnessCapability, harnessInputDigest } = + jest.requireActual('./authorization'); +const { HarnessOperationSchema, harnessOperationScope } = + jest.requireActual('./operation-contract'); +export const call = (name = 'kilo.organizations', args: unknown = {}) => ({ + type: 'execute', + conversationId, + operationId, + runId, + toolCallId, + request: { name, arguments: args }, + dispatchStartedAt: originalTime, +}); +export function capability(raw: unknown) { + const input = HarnessOperationSchema.parse(raw); + return jwt.sign( + { + grantId: runId, + authority, + scope: { + audience: 'agent-harness:operations', + conversationId, + operation: input.type, + definitionVersion: '1', + inputDigest: harnessInputDigest(JSON.parse(JSON.stringify(input))), + dispatchId: operationId, + target: { kind: 'backend' }, + }, + }, + 'test-signing-key', + { + issuer: 'agent-harness', + audience: 'agent-harness:operations', + expiresIn: 60, + } + ); +} +export async function authorizedInput(raw: unknown, token = capability(raw)) { + const input = HarnessOperationSchema.parse(JSON.parse(JSON.stringify(raw))); + const scope = harnessOperationScope(input); + return { input, scope, ...(await authorizeHarnessCapability(token, scope)) }; +} +beforeEach(() => { + jest.spyOn(Date, 'now').mockReturnValue(originalTime + 1); + access.expires = new Date(originalTime + 3600000).toISOString(); + access.active = access.role = true; +}); +afterEach(() => jest.restoreAllMocks());