diff --git a/.changeset/kilo-exa-websearch.md b/.changeset/kilo-exa-websearch.md new file mode 100644 index 00000000000..58bd6ec1463 --- /dev/null +++ b/.changeset/kilo-exa-websearch.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Route the websearch tool's Exa requests through the Kilo proxy when signed into Kilo. The MCP-Exa transport is preserved as a fallback for users who set `EXA_API_KEY` or are not authenticated. A new `KILO_WEBSEARCH_PROVIDER=kilo-exa` env override forces the Kilo proxy path. Results are capped at 10. diff --git a/packages/opencode/src/kilocode/tool/websearch-kilo-exa.ts b/packages/opencode/src/kilocode/tool/websearch-kilo-exa.ts new file mode 100644 index 00000000000..1e590b6c18f --- /dev/null +++ b/packages/opencode/src/kilocode/tool/websearch-kilo-exa.ts @@ -0,0 +1,74 @@ +// kilocode_change - new file +import { Duration, Effect, Schema } from "effect" +import { HttpClient, HttpClientRequest } from "effect/unstable/http" +import { KILO_API_BASE } from "@kilocode/kilo-gateway" + +export const KILO_EXA_URL = `${KILO_API_BASE}/api/exa/search` +export const MAX_KILO_EXA_RESULTS = 10 + +const ExaResult = Schema.Struct({ + title: Schema.optional(Schema.String), + url: Schema.String, + publishedDate: Schema.optional(Schema.String), + author: Schema.optional(Schema.String), + highlights: Schema.optional(Schema.Array(Schema.String)), +}) + +const ExaResponse = Schema.Struct({ + results: Schema.Array(ExaResult), +}) + +const NO_RESULTS = "No search results found. Please try a different query." + +const formatResults = (data: Schema.Schema.Type): string => { + if (data.results.length === 0) return NO_RESULTS + return data.results + .map((r, i) => { + const head = `[${i + 1}] ${r.title ?? r.url}\n${r.url}${r.publishedDate ? ` (${r.publishedDate})` : ""}` + const hl = r.highlights?.length ? `\n${r.highlights.map((h) => `> ${h}`).join("\n")}` : "" + return `${head}${hl}` + }) + .join("\n\n") +} + +export type KiloExaParams = { + query: string + type?: string + numResults?: number +} + +export const callKiloExa = Effect.fn("WebSearchKiloExa.call")(function* ( + http: HttpClient.HttpClient, + params: KiloExaParams, + kiloToken: string, +) { + const numResults = Math.min(params.numResults ?? MAX_KILO_EXA_RESULTS, MAX_KILO_EXA_RESULTS) + const request = yield* HttpClientRequest.post(KILO_EXA_URL).pipe( + HttpClientRequest.bearerToken(kiloToken), + HttpClientRequest.acceptJson, + HttpClientRequest.bodyJson({ + query: params.query, + type: params.type ?? "auto", + numResults, + contents: { highlights: true }, + }), + ) + const response = yield* http.execute(request).pipe( + Effect.timeoutOrElse({ + duration: Duration.seconds(25), + orElse: () => Effect.die(new Error("kilo exa request timed out")), + }), + ) + const status = response.status + if (status === 401 || status === 403) { + return yield* Effect.die(new Error(`Kilo exa request unauthorized (${status}); sign in with \`kilo auth login\``)) + } + if (status < 200 || status >= 300) { + const body = yield* response.text + return yield* Effect.die(new Error(`Kilo exa request failed (${status}): ${body.slice(0, 200)}`)) + } + const data = yield* response.json + const decode = Schema.decodeUnknownEffect(ExaResponse) + const parsed = yield* decode(data).pipe(Effect.orDie) + return formatResults(parsed) +}) diff --git a/packages/opencode/src/tool/websearch.ts b/packages/opencode/src/tool/websearch.ts index 12da7fbb3aa..9a812cfc626 100644 --- a/packages/opencode/src/tool/websearch.ts +++ b/packages/opencode/src/tool/websearch.ts @@ -1,16 +1,20 @@ -import { Effect, Schema } from "effect" +import { Effect, Option, Schema } from "effect" // kilocode_change - Option added for kilo-exa transport dispatch import { HttpClient } from "effect/unstable/http" import * as Tool from "./tool" import * as McpWebSearch from "./mcp-websearch" +import * as KiloExa from "@/kilocode/tool/websearch-kilo-exa" // kilocode_change - Kilo-REST Exa transport import DESCRIPTION from "./websearch.txt" import { checksum } from "@opencode-ai/core/util/encode" import { InstallationVersion } from "@opencode-ai/core/installation/version" import { RuntimeFlags } from "@/effect/runtime-flags" +import { Auth } from "@/auth" // kilocode_change - source Kilo bearer for Kilo-REST transport + +const MAX_RESULTS = 10 // kilocode_change - cap numResults across all transports export const Parameters = Schema.Struct({ query: Schema.String.annotate({ description: "Websearch query" }), numResults: Schema.optional(Schema.Number).annotate({ - description: "Number of search results to return (default: 8)", + description: "Number of search results to return (default: 8, maximum: 10)", // kilocode_change - note MAX_RESULTS cap }), livecrawl: Schema.optional(Schema.Literals(["fallback", "preferred"])).annotate({ description: @@ -24,12 +28,12 @@ export const Parameters = Schema.Struct({ }), }) -const WebSearchProviderSchema = Schema.Literals(["exa", "parallel"]) +const WebSearchProviderSchema = Schema.Literals(["exa", "parallel", "kilo-exa"]) // kilocode_change - kilo-exa env override export type WebSearchProvider = Schema.Schema.Type export function selectWebSearchProvider(sessionID: string, flags = { exa: false, parallel: false }): WebSearchProvider { const override = process.env.KILO_WEBSEARCH_PROVIDER - if (override === "exa" || override === "parallel") return override + if (override === "exa" || override === "parallel" || override === "kilo-exa") return override // kilocode_change - kilo-exa env override if (flags.parallel) return "parallel" if (flags.exa) return "exa" @@ -38,7 +42,7 @@ export function selectWebSearchProvider(sessionID: string, flags = { exa: false, export function webSearchProviderLabel(provider: unknown) { if (provider === "parallel") return "Parallel Web Search" - if (provider === "exa") return "Exa Web Search" + if (provider === "exa" || provider === "kilo-exa") return "Exa Web Search" // kilocode_change - kilo-exa shares label return "Web Search" } @@ -88,7 +92,7 @@ function callProvider( { query: params.query, type: params.type || "auto", - numResults: params.numResults || 8, + numResults: Math.min(params.numResults || 8, MAX_RESULTS), // kilocode_change - cap at MAX_RESULTS livecrawl: params.livecrawl || "fallback", contextMaxCharacters: params.contextMaxCharacters, }, @@ -101,6 +105,7 @@ export const WebSearchTool = Tool.define( Effect.gen(function* () { const http = yield* HttpClient.HttpClient const flags = yield* RuntimeFlags.Service + const authSvc = yield* Auth.Service // kilocode_change - source Kilo bearer for Kilo-REST transport return { get description() { @@ -114,7 +119,36 @@ export const WebSearchTool = Tool.define( parallel: flags.enableParallel, }) const title = webSearchProviderLabel(provider) - yield* ctx.metadata({ title: `${title} "${params.query}"`, metadata: { provider } }) + // kilocode_change start - Kilo-REST Exa transport + // Precedence: + // provider="kilo-exa" -> kilo-rest (auth required) + // provider="exa" + EXA_API_KEY -> mcp-exa-byok (BYOK wins) + // provider="exa" + Kilo auth -> kilo-rest (new default for authed users) + // provider="exa" + no auth -> mcp-exa-unauth (preserves current fallback) + // provider="parallel" -> mcp-parallel (unchanged) + const kiloToken = yield* Effect.gen(function* () { + if (provider !== "exa" && provider !== "kilo-exa") return undefined as string | undefined + const info = yield* authSvc.get("kilo") + if (!info) return undefined + return info.type === "api" ? info.key : info.type === "oauth" ? info.access : undefined + }) + const transport = + provider === "kilo-exa" + ? "kilo-rest" + : provider === "parallel" + ? "mcp-parallel" + : provider === "exa" && process.env.EXA_API_KEY + ? "mcp-exa-byok" + : provider === "exa" && kiloToken + ? "kilo-rest" + : "mcp-exa-unauth" + // kilocode_change end + // kilocode_change start - add transport to metadata + yield* ctx.metadata({ + title: `${title} "${params.query}"`, + metadata: { provider, transport }, + }) + // kilocode_change end yield* ctx.ask({ permission: "websearch", @@ -130,12 +164,26 @@ export const WebSearchTool = Tool.define( }, }) - const result = yield* callProvider(http, provider, params, ctx) + // kilocode_change start - dispatch Kilo-REST transport + const result = yield* transport === "kilo-rest" + ? kiloToken + ? KiloExa.callKiloExa( + http, + { + query: params.query, + type: params.type, + numResults: params.numResults, + }, + kiloToken, + ) + : Effect.die(new Error("KILO_WEBSEARCH_PROVIDER=kilo-exa requires Kilo auth; run `kilo auth login`")) + : callProvider(http, provider, params, ctx) + // kilocode_change end return { output: result ?? "No search results found. Please try a different query.", title: `${title}: ${params.query}`, - metadata: { provider }, + metadata: { provider, transport }, // kilocode_change - add transport } }).pipe(Effect.orDie), } diff --git a/packages/opencode/test/kilocode/tool/websearch-kilo-exa.test.ts b/packages/opencode/test/kilocode/tool/websearch-kilo-exa.test.ts new file mode 100644 index 00000000000..5756e7c671f --- /dev/null +++ b/packages/opencode/test/kilocode/tool/websearch-kilo-exa.test.ts @@ -0,0 +1,180 @@ +// kilocode_change - new file +import { describe, expect, test } from "bun:test" +import { Effect, Exit, Layer } from "effect" +import { HttpBody, HttpClient, HttpClientResponse } from "effect/unstable/http" +import { + KILO_EXA_URL, + MAX_KILO_EXA_RESULTS, + type KiloExaParams, + callKiloExa, +} from "../../../src/kilocode/tool/websearch-kilo-exa" + +type Recorded = { + url?: string + method?: string + authorization?: string + body?: string +} + +const readBody = async (body: HttpBody.HttpBody): Promise => { + if (body._tag === "Uint8Array") return new TextDecoder().decode(body.body) + if (body._tag === "Raw") return JSON.stringify(body.body) + return "" +} + +const jsonResponse = (status: number, body: unknown): Response => + new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }) + +const okJson = (body: unknown) => jsonResponse(200, body) + +const fakeHttp = (respond: (status: number) => Response, recorded?: Recorded): HttpClient.HttpClient => + HttpClient.make((request) => + Effect.gen(function* () { + const url = request.url + const method = request.method + const authorization = request.headers["authorization"] + const body = yield* Effect.promise(() => readBody(request.body)) + if (recorded) { + recorded.url = url + recorded.method = method + recorded.authorization = authorization + recorded.body = body + } + return HttpClientResponse.fromWeb( + request as unknown as Parameters[0], + respond(200), + ) + }), + ) + +const runCall = async ( + params: KiloExaParams, + respond: (status: number) => Response, + recorded?: Recorded, + kiloToken = "kilo-test-token", +) => + Effect.runPromiseExit( + Effect.gen(function* () { + const http = fakeHttp(respond, recorded) + return yield* callKiloExa(http, params, kiloToken) + }), + ) + +describe("callKiloExa request shape", () => { + test("posts to KILO_EXA_URL with bearer token and highlights-only contents", async () => { + const recorded: Recorded = {} + const exit = await runCall({ query: "drone" }, () => okJson({ results: [] }), recorded) + expect(Exit.isSuccess(exit)).toBe(true) + expect(recorded.url).toContain("/api/exa/search") + expect(recorded.method).toBe("POST") + expect(recorded.authorization).toBe("Bearer kilo-test-token") + const parsed = JSON.parse(recorded.body!) + expect(parsed.query).toBe("drone") + expect(parsed.type).toBe("auto") + expect(parsed.numResults).toBe(MAX_KILO_EXA_RESULTS) + expect(parsed.contents).toEqual({ highlights: true }) + }) + + test("uses caller numResults when below cap", async () => { + const recorded: Recorded = {} + await runCall({ query: "x", numResults: 3 }, () => okJson({ results: [] }), recorded) + expect(JSON.parse(recorded.body!).numResults).toBe(3) + }) + + test("clamps numResults at MAX_KILO_EXA_RESULTS", async () => { + const recorded: Recorded = {} + await runCall({ query: "x", numResults: 25 }, () => okJson({ results: [] }), recorded) + expect(JSON.parse(recorded.body!).numResults).toBe(MAX_KILO_EXA_RESULTS) + }) + + test("passes through caller type", async () => { + const recorded: Recorded = {} + await runCall({ query: "x", type: "deep" }, () => okJson({ results: [] }), recorded) + expect(JSON.parse(recorded.body!).type).toBe("deep") + }) + + test("KILO_EXA_URL is built from KILO_API_BASE", () => { + expect(KILO_EXA_URL).toMatch(/\/api\/exa\/search$/) + }) +}) + +describe("callKiloExa response formatting", () => { + const okValue = (exit: Exit.Exit): string => { + if (Exit.isFailure(exit)) throw new Error("expected success") + return (exit as Extract).value as string + } + + test("formats results with title, url, date and highlights", async () => { + const exit = await runCall({ query: "x" }, () => + okJson({ + results: [ + { + title: "A drone", + url: "https://example.com/a", + publishedDate: "2025-01-02T00:00:00.000Z", + highlights: ["first", "second"], + }, + ], + }), + ) + const text = okValue(exit) + expect(text).toContain("[1] A drone") + expect(text).toContain("https://example.com/a") + expect(text).toContain("(2025-01-02T00:00:00.000Z)") + expect(text).toContain("> first") + expect(text).toContain("> second") + }) + + test("falls back to url when title is missing", async () => { + const exit = await runCall({ query: "x" }, () => okJson({ results: [{ url: "https://example.com/no-title" }] })) + expect(okValue(exit)).toContain("[1] https://example.com/no-title") + }) + + test("returns NO_RESULTS message on empty results", async () => { + const exit = await runCall({ query: "x" }, () => okJson({ results: [] })) + expect(okValue(exit)).toBe("No search results found. Please try a different query.") + }) + + test("ignores costDollars on the response (cost accounting out of scope)", async () => { + const exit = await runCall({ query: "x" }, () => + okJson({ + results: [{ url: "https://example.com" }], + costDollars: { total: 0.007, search: { neural: 0.007 } }, + requestId: "req-123", + }), + ) + expect(Exit.isSuccess(exit)).toBe(true) + }) +}) + +describe("callKiloExa error handling", () => { + test("dies with auth-required message on 401", async () => { + const exit = await runCall({ query: "x" }, () => jsonResponse(401, {})) + expect(Exit.isFailure(exit)).toBe(true) + if (Exit.isSuccess(exit)) return + expect(String((exit as Extract).cause)).toContain("unauthorized") + expect(String((exit as Extract).cause)).toContain("kilo auth login") + }) + + test("dies with auth-required message on 403", async () => { + const exit = await runCall({ query: "x" }, () => jsonResponse(403, {})) + expect(Exit.isFailure(exit)).toBe(true) + if (Exit.isSuccess(exit)) return + expect(String((exit as Extract).cause)).toContain("unauthorized") + }) + + test("dies with status code on other non-2xx", async () => { + const exit = await runCall({ query: "x" }, () => jsonResponse(500, { error: "boom" })) + expect(Exit.isFailure(exit)).toBe(true) + if (Exit.isSuccess(exit)) return + expect(String((exit as Extract).cause)).toContain("500") + }) + + test("dies when response body is not valid ExaResponse shape", async () => { + const exit = await runCall({ query: "x" }, () => okJson({ nope: true })) + expect(Exit.isFailure(exit)).toBe(true) + }) +}) diff --git a/packages/opencode/test/tool/__snapshots__/parameters.test.ts.snap b/packages/opencode/test/tool/__snapshots__/parameters.test.ts.snap index f4951dfa51b..368225175b6 100644 --- a/packages/opencode/test/tool/__snapshots__/parameters.test.ts.snap +++ b/packages/opencode/test/tool/__snapshots__/parameters.test.ts.snap @@ -462,7 +462,7 @@ exports[`tool parameters JSON Schema (wire shape) websearch 1`] = ` "type": "string", }, "numResults": { - "description": "Number of search results to return (default: 8)", + "description": "Number of search results to return (default: 8, maximum: 10)", "type": "number", }, "query": {