Răsfoiți Sursa

refactor(tool): convert websearch tool internals to Effect (#21810)

Kit Langton 2 săptămâni în urmă
părinte
comite
46b74e0873

+ 74 - 0
packages/opencode/src/tool/mcp-exa.ts

@@ -0,0 +1,74 @@
+import { Duration, Effect, Schema } from "effect"
+import { HttpClient, HttpClientRequest } from "effect/unstable/http"
+
+const URL = "https://mcp.exa.ai/mcp"
+
+const McpResult = Schema.Struct({
+  result: Schema.Struct({
+    content: Schema.Array(
+      Schema.Struct({
+        type: Schema.String,
+        text: Schema.String,
+      }),
+    ),
+  }),
+})
+
+const decode = Schema.decodeUnknownEffect(Schema.fromJsonString(McpResult))
+
+const parseSse = Effect.fn("McpExa.parseSse")(function* (body: string) {
+  for (const line of body.split("\n")) {
+    if (!line.startsWith("data: ")) continue
+    const data = yield* decode(line.substring(6))
+    if (data.result.content[0]?.text) return data.result.content[0].text
+  }
+  return undefined
+})
+
+export const SearchArgs = Schema.Struct({
+  query: Schema.String,
+  type: Schema.String,
+  numResults: Schema.Number,
+  livecrawl: Schema.String,
+  contextMaxCharacters: Schema.optional(Schema.Number),
+})
+
+export const CodeArgs = Schema.Struct({
+  query: Schema.String,
+  tokensNum: Schema.Number,
+})
+
+const McpRequest = <F extends Schema.Struct.Fields>(args: Schema.Struct<F>) =>
+  Schema.Struct({
+    jsonrpc: Schema.Literal("2.0"),
+    id: Schema.Literal(1),
+    method: Schema.Literal("tools/call"),
+    params: Schema.Struct({
+      name: Schema.String,
+      arguments: args,
+    }),
+  })
+
+export const call = <F extends Schema.Struct.Fields>(
+  http: HttpClient.HttpClient,
+  tool: string,
+  args: Schema.Struct<F>,
+  value: Schema.Struct.Type<F>,
+  timeout: Duration.Input,
+) =>
+  Effect.gen(function* () {
+    const request = yield* HttpClientRequest.post(URL).pipe(
+      HttpClientRequest.accept("application/json, text/event-stream"),
+      HttpClientRequest.schemaBodyJson(McpRequest(args))({
+        jsonrpc: "2.0" as const,
+        id: 1 as const,
+        method: "tools/call" as const,
+        params: { name: tool, arguments: value },
+      }),
+    )
+    const response = yield* HttpClient.filterStatusOk(http).execute(request).pipe(
+      Effect.timeoutOrElse({ duration: timeout, orElse: () => Effect.die(new Error(`${tool} request timed out`)) }),
+    )
+    const body = yield* response.text
+    return yield* parseSse(body)
+  })

+ 2 - 1
packages/opencode/src/tool/registry.ts

@@ -101,6 +101,7 @@ export namespace ToolRegistry {
       const lsptool = yield* LspTool
       const plan = yield* PlanExitTool
       const webfetch = yield* WebFetchTool
+      const websearch = yield* WebSearchTool
 
       const state = yield* InstanceState.make<State>(
         Effect.fn("ToolRegistry.state")(function* (ctx) {
@@ -168,7 +169,7 @@ export namespace ToolRegistry {
             task: Tool.init(task),
             fetch: Tool.init(webfetch),
             todo: Tool.init(todo),
-            search: Tool.init(WebSearchTool),
+            search: Tool.init(websearch),
             code: Tool.init(CodeSearchTool),
             skill: Tool.init(SkillTool),
             patch: Tool.init(ApplyPatchTool),

+ 52 - 126
packages/opencode/src/tool/websearch.ts

@@ -1,15 +1,9 @@
 import z from "zod"
+import { Effect } from "effect"
+import { HttpClient } from "effect/unstable/http"
 import { Tool } from "./tool"
+import * as McpExa from "./mcp-exa"
 import DESCRIPTION from "./websearch.txt"
-import { abortAfterAny } from "../util/abort"
-
-const API_CONFIG = {
-  BASE_URL: "https://mcp.exa.ai",
-  ENDPOINTS: {
-    SEARCH: "/mcp",
-  },
-  DEFAULT_NUM_RESULTS: 8,
-} as const
 
 const Parameters = z.object({
   query: z.string().describe("Websearch query"),
@@ -30,121 +24,53 @@ const Parameters = z.object({
     .describe("Maximum characters for context string optimized for LLMs (default: 10000)"),
 })
 
-interface McpSearchRequest {
-  jsonrpc: string
-  id: number
-  method: string
-  params: {
-    name: string
-    arguments: {
-      query: string
-      numResults?: number
-      livecrawl?: "fallback" | "preferred"
-      type?: "auto" | "fast" | "deep"
-      contextMaxCharacters?: number
-    }
-  }
-}
-
-interface McpSearchResponse {
-  jsonrpc: string
-  result: {
-    content: Array<{
-      type: string
-      text: string
-    }>
-  }
-}
-
-export const WebSearchTool = Tool.define("websearch", async () => {
-  return {
-    get description() {
-      return DESCRIPTION.replace("{{year}}", new Date().getFullYear().toString())
-    },
-    parameters: Parameters,
-    async execute(params, ctx) {
-      await ctx.ask({
-        permission: "websearch",
-        patterns: [params.query],
-        always: ["*"],
-        metadata: {
-          query: params.query,
-          numResults: params.numResults,
-          livecrawl: params.livecrawl,
-          type: params.type,
-          contextMaxCharacters: params.contextMaxCharacters,
-        },
-      })
-
-      const searchRequest: McpSearchRequest = {
-        jsonrpc: "2.0",
-        id: 1,
-        method: "tools/call",
-        params: {
-          name: "web_search_exa",
-          arguments: {
-            query: params.query,
-            type: params.type || "auto",
-            numResults: params.numResults || API_CONFIG.DEFAULT_NUM_RESULTS,
-            livecrawl: params.livecrawl || "fallback",
-            contextMaxCharacters: params.contextMaxCharacters,
-          },
-        },
-      }
-
-      const { signal, clearTimeout } = abortAfterAny(25000, ctx.abort)
-
-      try {
-        const headers: Record<string, string> = {
-          accept: "application/json, text/event-stream",
-          "content-type": "application/json",
-        }
-
-        const response = await fetch(`${API_CONFIG.BASE_URL}${API_CONFIG.ENDPOINTS.SEARCH}`, {
-          method: "POST",
-          headers,
-          body: JSON.stringify(searchRequest),
-          signal,
-        })
-
-        clearTimeout()
-
-        if (!response.ok) {
-          const errorText = await response.text()
-          throw new Error(`Search error (${response.status}): ${errorText}`)
-        }
-
-        const responseText = await response.text()
-
-        // Parse SSE response
-        const lines = responseText.split("\n")
-        for (const line of lines) {
-          if (line.startsWith("data: ")) {
-            const data: McpSearchResponse = JSON.parse(line.substring(6))
-            if (data.result && data.result.content && data.result.content.length > 0) {
-              return {
-                output: data.result.content[0].text,
-                title: `Web search: ${params.query}`,
-                metadata: {},
-              }
-            }
+export const WebSearchTool = Tool.defineEffect(
+  "websearch",
+  Effect.gen(function* () {
+    const http = yield* HttpClient.HttpClient
+
+    return {
+      get description() {
+        return DESCRIPTION.replace("{{year}}", new Date().getFullYear().toString())
+      },
+      parameters: Parameters,
+      execute: (params: z.infer<typeof Parameters>, ctx: Tool.Context) =>
+        Effect.gen(function* () {
+          yield* Effect.promise(() =>
+            ctx.ask({
+              permission: "websearch",
+              patterns: [params.query],
+              always: ["*"],
+              metadata: {
+                query: params.query,
+                numResults: params.numResults,
+                livecrawl: params.livecrawl,
+                type: params.type,
+                contextMaxCharacters: params.contextMaxCharacters,
+              },
+            }),
+          )
+
+          const result = yield* McpExa.call(
+            http,
+            "web_search_exa",
+            McpExa.SearchArgs,
+            {
+              query: params.query,
+              type: params.type || "auto",
+              numResults: params.numResults || 8,
+              livecrawl: params.livecrawl || "fallback",
+              contextMaxCharacters: params.contextMaxCharacters,
+            },
+            "25 seconds",
+          )
+
+          return {
+            output: result ?? "No search results found. Please try a different query.",
+            title: `Web search: ${params.query}`,
+            metadata: {},
           }
-        }
-
-        return {
-          output: "No search results found. Please try a different query.",
-          title: `Web search: ${params.query}`,
-          metadata: {},
-        }
-      } catch (error) {
-        clearTimeout()
-
-        if (error instanceof Error && error.name === "AbortError") {
-          throw new Error("Search request timed out")
-        }
-
-        throw error
-      }
-    },
-  }
-})
+        }).pipe(Effect.runPromise),
+    }
+  }),
+)

+ 1 - 1
packages/opencode/test/session/prompt-effect.test.ts

@@ -1,7 +1,7 @@
 import { NodeFileSystem } from "@effect/platform-node"
+import { FetchHttpClient } from "effect/unstable/http"
 import { expect } from "bun:test"
 import { Cause, Effect, Exit, Fiber, Layer } from "effect"
-import { FetchHttpClient } from "effect/unstable/http"
 import path from "path"
 import z from "zod"
 import { Agent as AgentSvc } from "../../src/agent/agent"