diff --git a/apps/web/src/lib/agent-harness/web-tools.test.ts b/apps/web/src/lib/agent-harness/web-tools.test.ts new file mode 100644 index 0000000000..7aa21885fb --- /dev/null +++ b/apps/web/src/lib/agent-harness/web-tools.test.ts @@ -0,0 +1,264 @@ +import { afterAll, beforeEach, expect, it, jest } from '@jest/globals'; +import { TRPCError } from '@trpc/server'; +import type { ToolRequest } from '@kilocode/agent-harness/tools'; +import type * as Web from './web-tools'; + +const conversationId = '11111111-1111-4111-8111-111111111111'; +const operationId = '22222222-2222-4222-8222-222222222222'; +let request: ToolRequest, response: Response, usage: number, balance: number, denied: boolean; +let ledger: Record[], sent: { url: string; body: unknown }[]; +jest.mock('@/lib/config.server', () => ({ EXA_API_KEY: 'provider-secret' })); +jest.mock('@/lib/drizzle', () => ({ readDb: {} })); +jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); +jest.mock('@/lib/exa-usage', () => ({ + getExaMonthlyUsage: async () => ({ usage, freeAllowance: 1000 }), + getExaFreeAllowanceMicrodollars: () => 1000, + recordExaUsage: async (entry: Record) => { + ledger.push(entry); + }, +})); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: async () => ({ balance }), +})); +jest.mock('node:dns/promises', () => ({ + resolve4: async (host: string) => + host === 'private.example' ? ['8.8.8.8', '10.0.0.1'] : ['8.8.8.8'], + resolve6: async () => [], +})); +jest.mock('./authorization', () => ({ + harnessInputDigest: JSON.stringify, + authorizeHarnessCapability: async (token: string, scope: Record) => { + if ( + denied || + token !== 'grant' || + JSON.stringify(scope) !== + JSON.stringify({ + audience: 'agent-harness:operations', + conversationId, + operation: request.name, + definitionVersion: '1', + inputDigest: JSON.stringify(request.arguments), + dispatchId: operationId, + target: { kind: 'backend' }, + }) + ) + throw new TRPCError({ code: 'FORBIDDEN' }); + return { authority: { organizationId: 'organization' }, ctx: { user: { id: 'oauth/owner' } } }; + }, +})); +const { executeHarnessWeb } = jest.requireActual('./web-tools'); +const originalFetch = globalThis.fetch; +const run = (limit = 1024 * 1024) => + executeHarnessWeb( + 'grant', + { conversationId, operationId, request }, + limit, + new AbortController().signal + ); +beforeEach(() => { + request = { name: 'web.search', arguments: { query: 'Kilo', limit: 5 } }; + response = Response.json({ results: [], costDollars: { total: 0.002 } }); + usage = 0; + balance = 1; + denied = false; + ledger = []; + sent = []; + globalThis.fetch = async (url, init) => { + sent.push({ url: String(url), body: JSON.parse(String(init?.body)) }); + return response; + }; +}); +afterAll(() => { + globalThis.fetch = originalFetch; +}); + +it.each([0, 1000])('records one authorized actual charge with monthly usage %s', async monthly => { + usage = monthly; + const body = { + results: [{ url: 'https://example.com', title: 'Kilo', text: 'Page text' }], + costDollars: { total: 0.002 }, + }; + response = Response.json(body); + expect(await run()).toEqual({ status: 'succeeded', body, costMicrodollars: 2000 }); + expect(ledger).toEqual([ + { + userId: 'oauth/owner', + organizationId: 'organization', + path: '/search', + costMicrodollars: 2000, + chargedToBalance: monthly === 1000, + freeAllowanceMicrodollars: 1000, + featureId: 'quick-chat', + type: undefined, + }, + ]); + expect(sent).toEqual([ + { + url: 'https://api.exa.ai/search', + body: { query: 'Kilo', numResults: 5, contents: { text: true } }, + }, + ]); +}); +it('blocks revoked authority and exhausted credit before provider dispatch', async () => { + denied = true; + expect(await run()).toMatchObject({ + status: 'failed', + error: { code: 'access_revoked', retryable: false }, + }); + denied = false; + usage = 1000; + balance = 0; + expect(await run()).toMatchObject({ + status: 'failed', + error: { code: 'limit_exceeded', retryable: false }, + costMicrodollars: 0, + }); + expect(sent).toEqual([]); + expect(ledger).toEqual([]); +}); +it.each([ + 'http://example.com', + 'file:///etc/passwd', + 'https://localhost', + 'https://127.0.0.1', + 'https://[::1]', + 'https://private.example', + 'https://user:password@example.com', + 'https://example.com/#fragment', +])('blocks unsafe retrieval: %s', async url => { + request = { name: 'web.retrieve', arguments: { url } }; + expect(await run()).toMatchObject({ + status: 'failed', + error: { code: 'invalid_input', retryable: false }, + }); + expect(sent).toEqual([]); +}); +it('uses only contents for public retrieval and cancels overflow before parsing or billing', async () => { + request = { name: 'web.retrieve', arguments: { url: 'https://EXAMPLE.COM/page' } }; + let cancelled = false; + response = new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(50)); + controller.enqueue(new Uint8Array(50)); + }, + cancel() { + cancelled = true; + }, + }), + { headers: { 'content-type': 'application/json' } } + ); + expect(await run(99)).toMatchObject({ + status: 'failed', + error: { code: 'limit_exceeded', retryable: false }, + costMicrodollars: null, + }); + expect(cancelled).toBe(true); + expect(ledger).toEqual([]); + expect(sent).toEqual([ + { url: 'https://api.exa.ai/contents', body: { ids: ['https://example.com/page'], text: true } }, + ]); +}); +it.each([ + [429, 'application/json', '{}', 'unavailable_tool', true], + [403, 'application/json', '{}', 'unavailable_tool', false], + [200, 'text/html', '', 'invalid_output', false], + [200, 'application/json', '{', 'invalid_output', false], +] as const)( + 'keeps provider failure %s %s explicit', + async (status, contentType, body, code, retryable) => { + response = new Response(body, { status, headers: { 'content-type': contentType } }); + expect(await run()).toMatchObject({ + status: 'failed', + error: { code, retryable }, + costMicrodollars: null, + }); + expect(ledger).toEqual([]); + } +); +it.each([ + { conversationId: operationId }, + { operationId: conversationId }, + { request: { name: 'web.search', arguments: { query: 'changed', limit: 5 } } }, + { request: { name: 'web.retrieve', arguments: { url: 'https://example.com' } } }, +])('rejects changed capability input %# before dispatch', async change => { + const result = await executeHarnessWeb( + 'grant', + { conversationId, operationId, request, ...change }, + 1024, + new AbortController().signal + ); + expect(result).toMatchObject({ + status: 'failed', + error: { code: 'access_revoked', retryable: false }, + costMicrodollars: 0, + }); + expect(sent).toEqual([]); + expect(ledger).toEqual([]); +}); +it.each(['web.search', 'web.retrieve'] as const)( + 'preserves empty provider data for %s', + async name => { + request = name === 'web.search' ? request : { name, arguments: { url: 'https://example.com' } }; + response = Response.json({ results: [] }); + expect(await run()).toEqual({ + status: 'succeeded', + body: { results: [] }, + costMicrodollars: null, + }); + } +); +it.each(['before', 'during', 'lost'])('reports cancellation or response loss: %s', async phase => { + const controller = new AbortController(); + let dispatched = false; + globalThis.fetch = async (_url, init) => { + dispatched = true; + if (phase === 'during') controller.abort(); + init?.signal?.throwIfAborted(); + throw new TypeError('Response lost'); + }; + if (phase === 'before') controller.abort(); + const result = await executeHarnessWeb( + 'grant', + { conversationId, operationId, request }, + 1024, + controller.signal + ); + expect(result).toMatchObject({ + status: 'failed', + error: { + code: phase === 'lost' ? 'unavailable_tool' : 'cancelled', + retryable: phase === 'lost', + }, + costMicrodollars: phase === 'before' ? 0 : null, + }); + expect(dispatched).toBe(phase !== 'before'); + expect(ledger).toEqual([]); +}); +it('keeps unsupported content terminal when cancellation fails', async () => { + response = new Response( + new ReadableStream({ + cancel() { + throw new Error('Cleanup failed'); + }, + }), + { headers: { 'content-type': 'text/html' } } + ); + expect(await run()).toMatchObject({ + status: 'failed', + error: { code: 'invalid_output', retryable: false }, + costMicrodollars: null, + }); + expect(ledger).toEqual([]); +}); +it.each([undefined, 0, -1, '0.002'])( + 'does not invent or bill unknown/invalid cost %s', + async total => { + response = Response.json({ results: [], costDollars: { total } }); + expect(await run()).toMatchObject({ + status: total === undefined || total === 0 ? 'succeeded' : 'failed', + costMicrodollars: total === 0 ? 0 : null, + }); + expect(ledger).toEqual([]); + } +); diff --git a/apps/web/src/lib/agent-harness/web-tools.ts b/apps/web/src/lib/agent-harness/web-tools.ts new file mode 100644 index 0000000000..ccfbc56be5 --- /dev/null +++ b/apps/web/src/lib/agent-harness/web-tools.ts @@ -0,0 +1,135 @@ +import 'server-only'; +import { z } from 'zod'; +import { TRPCError } from '@trpc/server'; +import type { ErrorSchema } from '@kilocode/agent-harness/contracts'; +import { ToolRequestSchema, toolDefinitions } from '@kilocode/agent-harness/tools'; +import { validatePublicHttpsDestination } from '@/lib/mcp-gateway/discovery-service'; +import { extractExaCostMicrodollars, prepareExaRequest } from '@/lib/exa-provider'; +import { authorizeHarnessCapability, harnessInputDigest } from './authorization'; + +const Invocation = z.strictObject({ + conversationId: z.uuid(), + operationId: z.uuid(), + request: ToolRequestSchema, +}); + +/** Internal adapter only. The Worker commits its web request reservation before calling this boundary. */ +export async function executeHarnessWeb( + token: string, + input: unknown, + httpResponseBytes: number, + signal: AbortSignal +) { + let costMicrodollars: number | null = 0; + const failed = ( + code: z.infer['code'], + message: string, + retryable = false + ) => ({ status: 'failed' as const, error: { code, message, retryable }, costMicrodollars }); + try { + const parsed = Invocation.safeParse(input); + if ( + !parsed.success || + !Number.isSafeInteger(httpResponseBytes) || + httpResponseBytes < 1 || + httpResponseBytes > 1024 * 1024 + ) + return failed('invalid_input', 'Invalid web request.'); + const { conversationId, operationId, request } = parsed.data; + if (request.name !== 'web.search' && request.name !== 'web.retrieve') + return failed('invalid_input', 'Invalid web request.'); + const definition = toolDefinitions.find(tool => tool.name === request.name); + if (!definition) return failed('invalid_input', 'Invalid web request.'); + const { authority, ctx } = await authorizeHarnessCapability(token, { + audience: 'agent-harness:operations', + conversationId, + operation: request.name, + definitionVersion: definition.version, + inputDigest: harnessInputDigest(request.arguments), + dispatchId: operationId, + target: { kind: 'backend' }, + }); + let target: string | undefined; + if (request.name === 'web.retrieve') { + try { + target = (await validatePublicHttpsDestination(request.arguments.url)).href; + } catch { + return failed('invalid_input', 'The page URL must resolve to a public HTTPS destination.'); + } + } + const provider = await prepareExaRequest(ctx.user, authority.organizationId ?? undefined); + if (provider instanceof Response) + return failed( + provider.status === 402 ? 'limit_exceeded' : 'unavailable_tool', + provider.status === 402 + ? 'Exa allowance and credit balance are exhausted.' + : 'Web search is unavailable.' + ); + signal.throwIfAborted(); + const path = request.name === 'web.search' ? '/search' : '/contents'; + const body = + request.name === 'web.search' + ? { + query: request.arguments.query, + numResults: request.arguments.limit, + contents: { text: true }, + } + : { ids: [target], text: true }; + // No trusted pre-call Exa price exists. A lost response leaves actual cost unknown. + costMicrodollars = null; + const response = await provider.send(path, body, signal, true); + if (!response.ok) { + await response.body?.cancel().catch(() => undefined); + return failed( + 'unavailable_tool', + 'The web provider rejected the request.', + response.status >= 500 || response.status === 408 || response.status === 429 + ); + } + if ( + response.headers.get('content-type')?.split(';')[0].trim().toLowerCase() !== + 'application/json' || + !response.body + ) { + await response.body?.cancel().catch(() => undefined); + return failed('invalid_output', 'The web provider returned unsupported content.'); + } + const reader = response.body.getReader(), + decoder = new TextDecoder(); + let text = '', + size = 0; + try { + while (true) { + const chunk = await reader.read(); + if (chunk.done) break; + size += chunk.value.byteLength; + if (size > httpResponseBytes) + return failed('limit_exceeded', 'The web response is too large.'); + text += decoder.decode(chunk.value, { stream: true }); + } + text += decoder.decode(); + } finally { + await reader.cancel().catch(() => undefined); + reader.releaseLock(); + } + const result: unknown = JSON.parse(text); + try { + costMicrodollars = extractExaCostMicrodollars(result) ?? null; + } catch { + return failed('invalid_output', 'The web provider returned an invalid cost.'); + } + await provider.record(path, costMicrodollars ?? undefined, 'quick-chat'); + return { status: 'succeeded' as const, body: result, costMicrodollars }; + } catch (error) { + if (signal.aborted) return failed('cancelled', 'The web request was cancelled.'); + if (error instanceof TRPCError && ['FORBIDDEN', 'UNAUTHORIZED'].includes(error.code)) + return failed('access_revoked', 'Access to this conversation is unavailable.'); + if (error instanceof SyntaxError) + return failed('invalid_output', 'The web provider returned invalid JSON.'); + return failed( + 'unavailable_tool', + 'The web request failed. Retry within the remaining request budget.', + true + ); + } +}