diff --git a/apps/web/src/lib/agent-harness/cloud-agent.test.ts b/apps/web/src/lib/agent-harness/cloud-agent.test.ts new file mode 100644 index 0000000000..18aedce02e --- /dev/null +++ b/apps/web/src/lib/agent-harness/cloud-agent.test.ts @@ -0,0 +1,265 @@ +import { beforeEach, describe, expect, it, jest } from '@jest/globals'; +import { TRPCClientError } from '@trpc/client'; +import { getTRPCErrorFromUnknown, TRPCError } from '@trpc/server'; +import { getHTTPStatusCodeFromError } from '@trpc/server/http'; +import { TRPC_ERROR_CODES_BY_KEY } from '@trpc/server/rpc'; +import type { z } from 'zod'; +import { + basePrepareSessionNextSchema, + baseSendMessageNextSchema, +} from '@/routers/cloud-agent-next-schemas'; +import { + caller, + cloudId, + fixture, + guard, + invocation, + message, + operationId, + org, + reference, + sessionId, +} from './cloud-agent-test-fixture'; +import type * as Adapter from './cloud-agent'; + +let loss: 'prepare' | 'admission' | 'response' | undefined; +let admissionError: unknown; +let ledger: Record | undefined; +let responseChange: Record; +jest.mock('@/lib/drizzle', () => ({ + db: { + query: { + operation_ledgers: { findFirst: async () => (fixture.hideEvidence ? undefined : ledger) }, + }, + }, +})); +function cloud(scoped: boolean) { + const check = (input: { organizationId?: string }) => { + guard(scoped ? input.organizationId : null); + if (admissionError) throw admissionError; + }; + const settle = (effect: string, result: T) => { + fixture.effects.push(effect); + if (loss) + throw getTRPCErrorFromUnknown(TRPCClientError.from(new Error('secret-provider-error'))); + return result; + }; + return { + prepareSession: async ( + input: z.input & { organizationId?: string } + ) => { + check(input); + const prepared = basePrepareSessionNextSchema.parse(input); + if (fixture.unavailable) throw new TRPCError({ code: 'PAYMENT_REQUIRED' }); + if ( + !prepared.autoInitiate || + prepared.operationKey !== operationId || + !prepared.initialMessageId + ) + throw new Error('Missing durable admission'); + const result = { kiloSessionId: sessionId, cloudAgentSessionId: cloudId }; + ledger = { + organization_id: fixture.organizationId, + status: 'admitted', + canonical_result: { ...result, initialMessageId: prepared.initialMessageId }, + }; + if (loss !== 'prepare') { + fixture.messages = [message(prepared.initialMessageId, prepared.prompt ?? '')]; + ledger.status = loss === 'admission' ? 'reconcile_pending' : 'completed'; + } + return settle('start', result); + }, + sendMessage: async ( + input: z.input & { organizationId?: string } + ) => { + check(input); + const sent = baseSendMessageNextSchema.parse(input); + if (fixture.unavailable) throw new TRPCError({ code: 'PAYMENT_REQUIRED' }); + if ( + sent.cloudAgentSessionId !== cloudId || + !sent.messageId || + sent.payload.type !== 'prompt' || + sent.payload.model !== 'model' || + sent.payload.mode !== fixture.mode + ) + throw new Error('Wrong continuation'); + fixture.messages = [message(sent.messageId, sent.payload.prompt)]; + return settle('continue', { + cloudAgentSessionId: cloudId, + messageId: sent.messageId, + streamUrl: 'secret-stream-ticket', + ...responseChange, + }); + }, + interruptSession: async (input: { sessionId: string; organizationId?: string }) => { + check(input); + if (input.sessionId !== cloudId) throw new Error('Wrong stop target'); + return settle('stop', { success: true }); + }, + }; +} +Object.assign(caller.cloudAgentNext, cloud(false)); +Object.assign(caller.organizations.cloudAgentNext, cloud(true)); +// The repository transformer does not hoist mocks. +const { executeHarnessCloudAgent, reconcileHarnessCloudAgent } = + jest.requireActual('./cloud-agent'); +const cases = [ + ['start', { prompt: 'fix', modelId: 'model', repository: 'owner/repo' }], + ['continue', { ...reference, message: 'continue' }], + ['stop', reference], +] as const; +const invoke = (name: string, args: unknown, reconcile = false) => + (reconcile ? reconcileHarnessCloudAgent : executeHarnessCloudAgent)( + `kilo.sessions.${name}`, + invocation(name, args) + ); +beforeEach(() => { + loss = undefined; + admissionError = undefined; + ledger = undefined; + responseChange = {}; +}); +describe.each([null, org])('authorized Cloud Agent context %s', scope => { + beforeEach(() => { + fixture.organizationId = fixture.sessionScope = scope; + }); + it.each(cases)('keeps remote %s rejection terminal without an effect', async (name, args) => { + for (const code of [ + 'UNAUTHORIZED', + 'FORBIDDEN', + 'BAD_REQUEST', + 'PRECONDITION_FAILED', + 'PAYMENT_REQUIRED', + ] as const) { + admissionError = getTRPCErrorFromUnknown( + TRPCClientError.from({ + error: { + message: 'secret-provider-error', + code: TRPC_ERROR_CODES_BY_KEY[code], + data: { code, httpStatus: getHTTPStatusCodeFromError(new TRPCError({ code })) }, + }, + }) + ); + await expect(invoke(name, args)).rejects.toMatchObject({ + code, + message: 'Cloud Agent rejected this operation.', + cause: undefined, + }); + expect(ledger).toBeUndefined(); + expect(fixture.effects).toEqual([]); + } + }); + it.each(cases)('executes %s with private output and real session linkage', async (name, args) => { + const outcome = await invoke(name, args); + expect(outcome).toEqual({ status: 'succeeded', output: reference }); + expect(fixture.effects).toEqual([name]); + if (name === 'start' && outcome.status === 'succeeded') + expect(await invoke('progress', outcome.output)).toEqual({ + status: 'succeeded', + output: { ...reference, status: 'running' }, + }); + }); + it.each(cases)('denies removed membership for %s', async (name, args) => { + fixture.revoked = true; + await expect(invoke(name, args)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(fixture.effects).toEqual([]); + }); + it.each([cases[1], cases[2]])('rejects a context mismatch for %s', async (name, args) => { + fixture.sessionScope = scope === null ? org : null; + await expect(invoke(name, args)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(fixture.effects).toEqual([]); + }); + it.each([cases[0], cases[1]])('preserves model eligibility for %s', async (name, args) => { + fixture.unavailable = true; + await expect(invoke(name, args)).rejects.toMatchObject({ code: 'PAYMENT_REQUIRED' }); + expect(fixture.effects).toEqual([]); + }); + it.each([ + ['response', ...cases[0]], + ['response', ...cases[1]], + ['response', ...cases[2]], + ['prepare', ...cases[0]], + ['admission', ...cases[0]], + ] as const)('reconciles %s loss for %s without redispatch', async (phase, name, args) => { + loss = phase; + expect(await invoke(name, args)).toMatchObject({ status: 'outcome_unknown' }); + fixture.hideEvidence = true; + expect(await invoke(name, args, true)).toMatchObject({ status: 'outcome_unknown' }); + fixture.hideEvidence = false; + expect(await invoke(name, args, true)).toMatchObject({ + status: name === 'stop' || phase === 'prepare' ? 'outcome_unknown' : 'succeeded', + }); + if (name === 'stop') { + jest + .spyOn(caller.cliSessionsV2, 'get') + .mockRejectedValueOnce(new TRPCError({ code: 'NOT_FOUND' })); + expect(await invoke(name, args, true)).toMatchObject({ status: 'outcome_unknown' }); + } + expect(fixture.effects).toEqual([name]); + }); + it.each(['failed', 'no_op', 'interrupted', 'superseded', 'unknown'])( + 'never confirms a %s create through matching history', + async status => { + loss = 'admission'; + await invoke(...cases[0]); + ledger!.status = status; + expect(await invoke(...cases[0], true)).toMatchObject({ status: 'outcome_unknown' }); + expect(fixture.effects).toEqual(['start']); + } + ); + it.each([cases[0], cases[1]])('requires the exact user turn for %s', async (name, args) => { + loss = 'admission'; + await invoke(name, args); + const turn = fixture.messages[0]; + const content = name === 'start' ? cases[0][1].prompt : cases[1][1].message; + for (const messages of [ + [], + [message('another-id', content)], + [message(turn.info.id, 'wrong content')], + [message(turn.info.id, content, 'assistant')], + ]) { + fixture.messages = messages; + expect(await invoke(name, args, true)).toMatchObject({ status: 'outcome_unknown' }); + } + fixture.messages = [turn]; + expect(await invoke(name, args, true)).toEqual({ status: 'succeeded', output: reference }); + expect(fixture.effects).toEqual([name]); + }); + it.each(['cloudAgentSessionId', 'messageId'])('rejects a mismatched response %s', async key => { + responseChange = { [key]: 'another-identity' }; + expect(await invoke(...cases[1])).toMatchObject({ status: 'outcome_unknown' }); + fixture.messages = []; + expect(await invoke(...cases[1], true)).toMatchObject({ status: 'outcome_unknown' }); + expect(fixture.effects).toEqual(['continue']); + }); + it.each([cases[1], cases[2]])('rechecks the grant before %s', async (name, args) => { + const api = scope === null ? caller.cloudAgentNext : caller.organizations.cloudAgentNext; + const getState = api.getSession, + getOwned = caller.cliSessionsV2.get; + const revoke = (result: T) => { + fixture.grantRevoked = true; + return result; + }; + if (name === 'continue') + jest.spyOn(api, 'getSession').mockImplementationOnce(input => getState(input).then(revoke)); + else + jest + .spyOn(caller.cliSessionsV2, 'get') + .mockImplementationOnce(input => getOwned(input).then(revoke)); + await expect(invoke(name, args)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(fixture.effects).toEqual([]); + }); + it.each([ + ['search', { query: 'absent' }, []], + ['attach', reference, { ...reference, untrusted: true, messages: [] }], + ['progress', reference, { ...reference, status: 'idle' }], + ] as const)( + 'delegates empty %s reads during execution and reconciliation', + async (name, args, output) => { + fixture.hideEvidence = true; + for (const reconcile of [false, true]) + expect(await invoke(name, args, reconcile)).toEqual({ status: 'succeeded', output }); + expect(fixture.effects).toEqual([]); + } + ); +}); diff --git a/apps/web/src/lib/agent-harness/cloud-agent.ts b/apps/web/src/lib/agent-harness/cloud-agent.ts new file mode 100644 index 0000000000..f2d80c4ef9 --- /dev/null +++ b/apps/web/src/lib/agent-harness/cloud-agent.ts @@ -0,0 +1,171 @@ +import 'server-only'; +import { TRPCError } from '@trpc/server'; +import { and, eq } from 'drizzle-orm'; +import { z } from 'zod'; +import { type ToolOutcome } from '@kilocode/agent-harness/contracts'; +import { operation_ledgers } from '@kilocode/db/schema'; +import { OPERATION_NON_TERMINAL_STATUSES } from '@kilocode/db/operation-ledger'; +import { db } from '@/lib/drizzle'; +import { + basePrepareSessionNextSchema, + baseSendMessageNextSchema, +} from '@/routers/cloud-agent-next-schemas'; +import { + createHarnessCloudAgentContext, + normalizeCloudAgentAdmissionError, +} from './cloud-agent-context'; + +const Reference = z.object({ + kiloSessionId: z.string().startsWith('ses_').length(30), + cloudAgentSessionId: z.string().min(1), + initialMessageId: z.string(), +}); + +async function sessionOperation( + token: string, + input: unknown, + reconcile: boolean +): Promise { + const context = createHarnessCloudAgentContext(token, input); + const { invocation, request, messageId, fresh, owned, history } = context; + const unknown: ToolOutcome = { + status: 'outcome_unknown', + reason: 'Check Cloud Agent status; do not repeat this operation.', + providerReference: invocation.operationId, + }; + const hasUserTurn = async (sessionId: string, content: string) => + (await history(sessionId)).some( + message => + message.info.id === messageId && + message.info.role === 'user' && + message.parts + .filter(part => part.type === 'text') + .map(part => part.text) + .join('\n') === content + ); + if (request.name === 'kilo.sessions.search') return context.search(); + if (request.name === 'kilo.sessions.attach') return context.attachContext(); + if (request.name === 'kilo.sessions.progress') return context.progress(); + let dispatched = false; + try { + const { caller, authority } = await fresh(); + const organizationId = authority.organizationId; + switch (request.name) { + case 'kilo.sessions.start': { + if (reconcile) { + const row = await db.query.operation_ledgers.findFirst({ + where: and( + eq(operation_ledgers.kilo_user_id, authority.userId), + eq(operation_ledgers.operation_key, invocation.operationId), + eq(operation_ledgers.domain, 'session'), + eq(operation_ledgers.intent, 'create_cloud') + ), + }); + if (!row) return unknown; + if (row.organization_id !== organizationId) throw new TRPCError({ code: 'FORBIDDEN' }); + const reference = Reference.safeParse(row.canonical_result); + if (!reference.success) return unknown; + if (reference.data.initialMessageId !== messageId) + throw new TRPCError({ code: 'CONFLICT' }); + const { session } = await owned(reference.data.kiloSessionId); + if (session.cloud_agent_session_id !== reference.data.cloudAgentSessionId) + throw new TRPCError({ code: 'FORBIDDEN' }); + // Allocation proves no admission. Failed settles cannot be overridden by history. + if ( + row.status !== 'completed' && + (!OPERATION_NON_TERMINAL_STATUSES.some(status => status === row.status) || + !(await hasUserTurn(session.session_id, request.arguments.prompt))) + ) + return unknown; + return context.succeeded({ sessionId: session.session_id }); + } + const prepared = basePrepareSessionNextSchema.parse({ + prompt: request.arguments.prompt, + model: request.arguments.modelId, + githubRepo: request.arguments.repository, + mode: 'code', + autoInitiate: true, + operationKey: invocation.operationId, + initialMessageId: messageId, + }); + dispatched = true; + // The existing grouped path persists both session IDs before initial admission. + // Never call its retry ladder during reconciliation: its ledger has finite retention. + const result = + organizationId === null + ? await caller.cloudAgentNext.prepareSession(prepared) + : await caller.organizations.cloudAgentNext.prepareSession({ + ...prepared, + organizationId, + }); + return context.succeeded({ sessionId: result.kiloSessionId }); + } + case 'kilo.sessions.continue': + case 'kilo.sessions.stop': { + const sessionId = request.arguments.sessionId; + const current = await context.cloudSession(sessionId); + const cloudAgentSessionId = current.cloudAgentSessionId; + if (reconcile) { + // Stop has no operation receipt. An idle/absent query cannot establish this Stop's outcome. + if ( + request.name !== 'kilo.sessions.continue' || + !(await hasUserTurn(sessionId, request.arguments.message)) + ) + return unknown; + } else if (request.name === 'kilo.sessions.stop') { + const { caller } = await fresh(); + dispatched = true; + const result = + organizationId === null + ? await caller.cloudAgentNext.interruptSession({ sessionId: cloudAgentSessionId }) + : await caller.organizations.cloudAgentNext.interruptSession({ + sessionId: cloudAgentSessionId, + organizationId, + }); + if (!result.success) return unknown; + } else if (request.name === 'kilo.sessions.continue') { + const state = await context.sessionState(current, cloudAgentSessionId); + const message = baseSendMessageNextSchema.parse({ + cloudAgentSessionId, + messageId, + payload: { + type: 'prompt', + prompt: request.arguments.message, + mode: state.mode, + model: state.model, + variant: state.variant, + }, + }); + const { caller } = await fresh(); + dispatched = true; + const result = + organizationId === null + ? await caller.cloudAgentNext.sendMessage(message) + : await caller.organizations.cloudAgentNext.sendMessage({ + ...message, + organizationId, + }); + if (result.cloudAgentSessionId !== cloudAgentSessionId || result.messageId !== messageId) + return unknown; + } + // The Kilo session ID is the shared app.openScreen session destination, not a stream ticket. + return context.succeeded({ sessionId }); + } + } + throw new TRPCError({ code: 'BAD_REQUEST' }); + } catch (error) { + // Preserve authoritative admission rejection without exposing provider text or transport causes. + const rejection = normalizeCloudAgentAdmissionError(error); + if (rejection) throw rejection; + if (!dispatched && error instanceof TRPCError && error.code === 'CONFLICT') throw error; + if (dispatched || reconcile) return unknown; + throw error; + } +} + +// The durable scheduler dispatches an admitted call once and retains its canonical result. +// On response loss it must use reconciliation, never execute again, including after ledger expiry. +export const executeHarnessCloudAgent = (token: string, input: unknown) => + sessionOperation(token, input, false); +export const reconcileHarnessCloudAgent = (token: string, input: unknown) => + sessionOperation(token, input, true);