feat(agent-harness): normalize web sources and preserve citations

services/agent-harness/src/tools/web.test.ts

@@ -0,0 +1,236 @@
+import { expect, it } from 'vitest';
+import { z } from 'zod';
+import { RunLimitsSchema } from '../commands';
+import { reserve, ReservationSchema, type Reservation } from '../limits';
+import { executeWeb } from './web';
+
+const search = { name: 'web.search', arguments: { query: 'Kilo', limit: 5 } };
+const retrieve = { name: 'web.retrieve', arguments: { url: 'https://example.com/' } };
+const source = {
+  url: 'https://EXAMPLE.COM',
+  title: ' Source ',
+  text: 'abcdef',
+  publishedDate: '2026-08-01',
+};
+function fixture(overrides = {}) {
+  const limits = RunLimitsSchema.parse(overrides);
+  const admission = {
+    runId: crypto.randomUUID(),
+    messageId: crypto.randomUUID(),
+    context: { type: 'personal' as const },
+    limits,
+    model: { contextTokens: 32000, inputUsdPerMillion: 1, outputUsdPerMillion: 1 },
+  };
+  let reservations: Reservation[] = [],
+    dispatches = 0;
+  const nextBudget = () => {
+    const reservation = reserve(
+      admission,
+      reservations,
+      { kind: 'tool', step: 1, toolCallId: crypto.randomUUID(), webRequest: true },
+      Date.now()
+    );
+    // Reload the persisted representation before dispatch, including uncertain prior requests.
+    reservations = z
+      .array(ReservationSchema)
+      .parse(JSON.parse(JSON.stringify([...reservations, reservation])));
+    return {
+      limits,
+      reservation: ReservationSchema.parse(reservations.at(-1)),
+      signal: new AbortController().signal,
+    };
+  };
+  const run = async (
+    input: unknown = search,
+    body: unknown = { results: [source] },
+    lost = false
+  ) =>
+    executeWeb(input, nextBudget(), async () => {
+      dispatches++;
+      if (lost) throw new Error('Lost provider response');
+      return { status: 'succeeded', body, costMicrodollars: 2000 };
+    });
+  return { run, nextBudget, count: () => dispatches };
+}
+it('returns bounded normalized untrusted sources with actual citation parts', async () => {
+  const f = fixture({ searchResults: 2, snippetCharacters: 3 });
+  expect(await f.run(search, { results: [source, { ...source, title: '' }, source] })).toEqual({
+    outcome: {
+      status: 'succeeded',
+      output: [
+        {
+          url: 'https://example.com/',
+          title: 'Source',
+          text: 'abc',
+          publishedAt: '2026-08-01T00:00:00.000Z',
+          untrusted: true,
+        },
+        {
+          url: 'https://example.com/',
+          title: 'https://example.com/',
+          text: 'abc',
+          publishedAt: '2026-08-01T00:00:00.000Z',
+          untrusted: true,
+        },
+      ],
+    },
+    citations: [
+      { type: 'citation', url: 'https://example.com/', title: 'Source' },
+      { type: 'citation', url: 'https://example.com/', title: 'https://example.com/' },
+    ],
+    costMicrodollars: 2000,
+  });
+});
+it.each([
+  [4, '漢字', '漢'],
+  [3, '𠮷a', ''],
+  [4, '𠮷a', '𠮷'],
+  [5, 'é漢x', 'é漢'],
+] as const)(
+  'bounds page UTF-8 bytes at %i without corruption',
+  async (pageBytes, text, expected) => {
+    expect(
+      await fixture({ pageBytes }).run(retrieve, { results: [{ ...source, text }] })
+    ).toMatchObject({
+      outcome: { status: 'succeeded', output: { text: expected, untrusted: true } },
+    });
+  }
+);
+it('bounds snippets without splitting Unicode characters', async () => {
+  const result = await fixture({ snippetCharacters: 1 }).run(search, {
+    results: [{ ...source, text: '𠮷a' }],
+  });
+  expect(JSON.stringify(result)).toContain('"text":"𠮷"');
+});
+it.each([search, retrieve])('returns honest empty %j data without citations', async request => {
+  expect(await fixture().run(request, { results: [] })).toMatchObject({
+    outcome: { status: 'succeeded', output: request === search ? [] : { text: '' } },
+    citations: [],
+  });
+});
+it.each([undefined, '', '   '])('omits citations for empty page text %j', async text => {
+  expect(await fixture().run(retrieve, { results: [{ ...source, text }] })).toMatchObject({
+    outcome: { status: 'succeeded', output: { text: text ?? '', untrusted: true } },
+    citations: [],
+  });
+});
+it.each([
+  [{ results: [{ ...source, url: 'javascript:alert(1)' }] }, {}, 'invalid_output'],
+  [{ results: [{ ...source, publishedDate: 'invalid' }] }, {}, 'invalid_output'],
+  [{ results: [{ ...source, url: 'https://user:password@example.com/' }] }, {}, 'invalid_output'],
+  [{ results: [{ ...source, text: 42 }] }, {}, 'invalid_output'],
+  [{ results: [source] }, { toolOutputBytes: 1 }, 'limit_exceeded'],
+  [{ results: [source] }, { toolInputBytes: 1 }, 'limit_exceeded'],
+  [
+    { results: [], statuses: [{ status: 'error', error: { httpStatusCode: 404 } }] },
+    {},
+    'unavailable_tool',
+  ],
+] as const)('rejects malformed, oversized, and failed source data', async (body, limits, code) => {
+  expect(await fixture(limits).run(search, body)).toMatchObject({
+    outcome: { status: 'failed', error: { code, retryable: false } },
+    citations: [],
+  });
+});
+it.each([408, 429, 500, undefined])(
+  'allows budgeted recovery for provider status %j',
+  async status => {
+    const body = {
+      results: [],
+      statuses: [{ status: 'error', error: { httpStatusCode: status } }],
+    };
+    expect(await fixture().run(retrieve, body)).toMatchObject({
+      outcome: { status: 'failed', error: { code: 'unavailable_tool', retryable: true } },
+      citations: [],
+      costMicrodollars: 2000,
+    });
+  }
+);
+it.each([
+  null,
+  { id: '' },
+  { kind: 'model' },
+  { webRequest: false },
+  { status: 'released' },
+  { status: 'finished' },
+  { toolCallId: null },
+  { deadline: 0 },
+  { deadline: NaN },
+])('rejects an inadmissible reservation %j before dispatch', async patch => {
+  const budget = fixture().nextBudget();
+  const reservation = patch === null ? null : { ...budget.reservation, ...patch };
+  const result = await executeWeb(
+    search,
+    { ...budget, reservation: reservation as Reservation },
+    async () => ({
+      status: 'succeeded',
+      body: { results: [source] },
+      costMicrodollars: 2000,
+    })
+  );
+  expect(result).toMatchObject({
+    outcome: { status: 'failed', error: { code: 'limit_exceeded', retryable: false } },
+    citations: [],
+    costMicrodollars: null,
+  });
+});
+it.each([0, null, 2000])('preserves returned actual cost %j', async costMicrodollars => {
+  const result = await executeWeb(search, fixture().nextBudget(), async () => ({
+    status: 'succeeded',
+    body: { results: [] },
+    costMicrodollars,
+  }));
+  expect(result).toEqual({
+    outcome: { status: 'succeeded', output: [] },
+    citations: [],
+    costMicrodollars,
+  });
+});
+it.each([
+  [false, 'Access denied.', 'unavailable_tool'],
+  [true, 'Retry later.', 'unavailable_tool'],
+  [true, 'x'.repeat(1024), 'limit_exceeded'],
+] as const)('preserves bounded provider failures (%j)', async (retryable, message, code) => {
+  const result = await executeWeb(
+    search,
+    fixture({ toolOutputBytes: 512 }).nextBudget(),
+    async () => ({
+      status: 'failed',
+      error: { code: 'unavailable_tool', message, retryable },
+      costMicrodollars: null,
+    })
+  );
+  expect(result).toMatchObject({
+    outcome: {
+      status: 'failed',
+      error: { code, retryable: code === 'limit_exceeded' ? false : retryable },
+    },
+    citations: [],
+    costMicrodollars: null,
+  });
+});
+it('bounds the complete UTF-8 result including citations and its envelope', async () => {
+  const body = { results: [{ ...source, title: '漢'.repeat(30) }] };
+  const result = await fixture().run(search, body);
+  const size = new Blob([JSON.stringify(result)]).size;
+  expect(await fixture({ toolOutputBytes: size }).run(search, body)).toEqual(result);
+  expect(await fixture({ toolOutputBytes: size - 1 }).run(search, body)).toMatchObject({
+    outcome: { status: 'failed', error: { code: 'limit_exceeded', retryable: false } },
+    citations: [],
+    costMicrodollars: 2000,
+  });
+});
+it('allows five live requests, consumes uncertainty, and keeps Exa outside the model ceiling', async () => {
+  const f = fixture({ modelCostUsd: 0.000001 });
+  expect(await f.run(search, {}, true)).toMatchObject({
+    outcome: { status: 'failed', error: { retryable: true } },
+    costMicrodollars: null,
+  });
+  for (let count = 1; count < 5; count++)
+    expect(await f.run()).toMatchObject({
+      outcome: { status: 'succeeded' },
+      costMicrodollars: 2000,
+    });
+  await expect(f.run()).rejects.toMatchObject({ detail: { code: 'limit_exceeded' } });
+  expect(f.count()).toBe(5);
+});

services/agent-harness/src/tools/web.ts

@@ -0,0 +1,159 @@
+import { z } from 'zod';
+import { ErrorSchema, type Message, type ToolOutcome } from '@kilocode/agent-harness/contracts';
+import {
+  ToolRequestSchema,
+  toolDefinitions,
+  type ToolRequest,
+} from '@kilocode/agent-harness/tools';
+import { bytes, ReservationSchema, type Reservation, type RunLimits } from '../limits';
+
+const Reply = z.discriminatedUnion('status', [
+  z.object({
+    status: z.literal('succeeded'),
+    body: z.unknown(),
+    costMicrodollars: z.int().nonnegative().nullable(),
+  }),
+  z.object({
+    status: z.literal('failed'),
+    error: ErrorSchema,
+    costMicrodollars: z.int().nonnegative().nullable(),
+  }),
+]);
+const Sources = z.object({
+  results: z.array(
+    z.object({
+      url: z.url({ protocol: /^https?$/ }),
+      title: z.string().nullish(),
+      text: z.string().optional(),
+      publishedDate: z.string().optional(),
+    })
+  ),
+  statuses: z
+    .array(
+      z.object({
+        status: z.enum(['success', 'error']),
+        error: z.object({ httpStatusCode: z.int().nullish() }).nullish(),
+      })
+    )
+    .optional(),
+});
+export type WebResult = {
+  outcome: ToolOutcome;
+  citations: Message['parts'];
+  costMicrodollars: number | null;
+};
+
+/** The scheduler must supply its committed reservation. Never release it after an uncertain provider call. */
+export async function executeWeb(
+  input: unknown,
+  budget: { reservation: Reservation; limits: RunLimits; signal: AbortSignal },
+  invoke: (request: ToolRequest, httpResponseBytes: number, signal: AbortSignal) => Promise<unknown>
+): Promise<WebResult> {
+  let costMicrodollars: number | null = null;
+  const failed = (
+    code: z.infer<typeof ErrorSchema>['code'],
+    message: string,
+    retryable = false
+  ): WebResult => ({
+    outcome: { status: 'failed', error: { code, message, retryable } },
+    citations: [],
+    costMicrodollars,
+  });
+  const request = ToolRequestSchema.safeParse(input);
+  if (
+    !request.success ||
+    (request.data.name !== 'web.search' && request.data.name !== 'web.retrieve')
+  )
+    return failed('invalid_input', 'Invalid web request.');
+  const { reservation, limits, signal } = budget;
+  if (
+    !ReservationSchema.safeParse(reservation).success ||
+    reservation.kind !== 'tool' ||
+    !reservation.webRequest ||
+    reservation.status !== 'reserved' ||
+    !reservation.toolCallId ||
+    reservation.deadline <= Date.now()
+  )
+    return failed('limit_exceeded', 'A current web request reservation is required.');
+  if (bytes(request.data.arguments) > limits.toolInputBytes)
+    return failed('limit_exceeded', 'The web request is too large.');
+  let raw: unknown;
+  try {
+    signal.throwIfAborted();
+    raw = await invoke(request.data, limits.httpResponseBytes, signal);
+  } catch {
+    return failed(
+      signal.aborted ? 'cancelled' : 'unavailable_tool',
+      'The web response was not received; its cost is unknown.',
+      !signal.aborted
+    );
+  }
+  try {
+    const reply = Reply.parse(raw);
+    costMicrodollars = reply.costMicrodollars;
+    if (reply.status === 'failed') {
+      const result = failed(reply.error.code, reply.error.message, reply.error.retryable);
+      return bytes(result) > limits.toolOutputBytes
+        ? failed('limit_exceeded', 'The web provider error is too large.')
+        : result;
+    }
+    const response = Sources.parse(reply.body);
+    const error = response.statuses?.find(item => item.status === 'error');
+    if (error) {
+      const status = error.error?.httpStatusCode;
+      return failed(
+        'unavailable_tool',
+        'The provider could not retrieve the page.',
+        status == null || status >= 500 || status === 408 || status === 429
+      );
+    }
+    const searching = request.data.name === 'web.search';
+    const count =
+      request.data.name === 'web.search'
+        ? Math.min(request.data.arguments.limit, limits.searchResults)
+        : 1;
+    const pages = response.results.slice(0, count).map(source => {
+      const url = new URL(source.url);
+      if (url.username || url.password) throw new Error('Invalid source URL');
+      const text = source.text ?? '';
+      return {
+        url: url.href,
+        title: source.title?.trim() || url.href,
+        text: searching
+          ? Array.from(text).slice(0, limits.snippetCharacters).join('')
+          : new TextDecoder().decode(new TextEncoder().encode(text).subarray(0, limits.pageBytes), {
+              stream: true,
+            }),
+        ...(source.publishedDate
+          ? { publishedAt: new Date(source.publishedDate).toISOString() }
+          : {}),
+        untrusted: true as const,
+      };
+    });
+    const definition = toolDefinitions.find(tool => tool.name === request.data.name);
+    if (!definition) return failed('invalid_input', 'Invalid web request.');
+    const output = definition.outputSchema.parse(
+      request.data.name === 'web.search'
+        ? pages
+        : (pages[0] ?? {
+            url: request.data.arguments.url,
+            title: request.data.arguments.url,
+            text: '',
+            untrusted: true,
+          })
+    );
+    const citations: Message['parts'] = pages
+      .filter(page => page.text.trim())
+      .map(page => ({ type: 'citation', url: page.url, title: page.title }));
+    const result: WebResult = {
+      outcome: { status: 'succeeded', output },
+      citations,
+      costMicrodollars,
+    };
+    if (bytes(result) > limits.toolOutputBytes)
+      return failed('limit_exceeded', 'The normalized web output is too large.');
+    return result;
+  } catch {
+    return failed('invalid_output', 'The web provider returned invalid source data.');
+  }
+}