koro-ai 1.0.0

package/src/tool/batch.ts

@@ -0,0 +1,183 @@
+import z from "zod"
+import { Tool } from "./tool"
+import { ProviderID, ModelID } from "../provider/schema"
+import { errorMessage } from "../util/error"
+import DESCRIPTION from "./batch.txt"
+
+const DISALLOWED = new Set(["batch"])
+const FILTERED_FROM_SUGGESTIONS = new Set(["invalid", "patch", ...DISALLOWED])
+
+export const BatchTool = Tool.define("batch", async () => {
+  return {
+    description: DESCRIPTION,
+    parameters: z.object({
+      tool_calls: z
+        .array(
+          z.object({
+            tool: z.string().describe("The name of the tool to execute"),
+            parameters: z.object({}).loose().describe("Parameters for the tool"),
+          }),
+        )
+        .min(1, "Provide at least one tool call")
+        .describe("Array of tool calls to execute in parallel"),
+    }),
+    formatValidationError(error) {
+      const formattedErrors = error.issues
+        .map((issue) => {
+          const path = issue.path.length > 0 ? issue.path.join(".") : "root"
+          return `  - ${path}: ${issue.message}`
+        })
+        .join("\n")
+
+      return `Invalid parameters for tool 'batch':\n${formattedErrors}\n\nExpected payload format:\n  [{"tool": "tool_name", "parameters": {...}}, {...}]`
+    },
+    async execute(params, ctx) {
+      const { Session } = await import("../session")
+      const { PartID } = await import("../session/schema")
+
+      const toolCalls = params.tool_calls.slice(0, 25)
+      const discardedCalls = params.tool_calls.slice(25)
+
+      const { ToolRegistry } = await import("./registry")
+      const availableTools = await ToolRegistry.tools({ modelID: ModelID.make(""), providerID: ProviderID.make("") })
+      const toolMap = new Map(availableTools.map((t) => [t.id, t]))
+
+      const executeCall = async (call: (typeof toolCalls)[0]) => {
+        const callStartTime = Date.now()
+        const partID = PartID.ascending()
+
+        try {
+          if (DISALLOWED.has(call.tool)) {
+            throw new Error(
+              `Tool '${call.tool}' is not allowed in batch. Disallowed tools: ${Array.from(DISALLOWED).join(", ")}`,
+            )
+          }
+
+          const tool = toolMap.get(call.tool)
+          if (!tool) {
+            const availableToolsList = Array.from(toolMap.keys()).filter((name) => !FILTERED_FROM_SUGGESTIONS.has(name))
+            throw new Error(
+              `Tool '${call.tool}' not in registry. External tools (MCP, environment) cannot be batched - call them directly. Available tools: ${availableToolsList.join(", ")}`,
+            )
+          }
+          const validatedParams = tool.parameters.parse(call.parameters)
+
+          await Session.updatePart({
+            id: partID,
+            messageID: ctx.messageID,
+            sessionID: ctx.sessionID,
+            type: "tool",
+            tool: call.tool,
+            callID: partID,
+            state: {
+              status: "running",
+              input: call.parameters,
+              time: {
+                start: callStartTime,
+              },
+            },
+          })
+
+          const result = await tool.execute(validatedParams, { ...ctx, callID: partID })
+          const attachments = result.attachments?.map((attachment) => ({
+            ...attachment,
+            id: PartID.ascending(),
+            sessionID: ctx.sessionID,
+            messageID: ctx.messageID,
+          }))
+
+          await Session.updatePart({
+            id: partID,
+            messageID: ctx.messageID,
+            sessionID: ctx.sessionID,
+            type: "tool",
+            tool: call.tool,
+            callID: partID,
+            state: {
+              status: "completed",
+              input: call.parameters,
+              output: result.output,
+              title: result.title,
+              metadata: result.metadata,
+              attachments,
+              time: {
+                start: callStartTime,
+                end: Date.now(),
+              },
+            },
+          })
+
+          return { success: true as const, tool: call.tool, result }
+        } catch (error) {
+          await Session.updatePart({
+            id: partID,
+            messageID: ctx.messageID,
+            sessionID: ctx.sessionID,
+            type: "tool",
+            tool: call.tool,
+            callID: partID,
+            state: {
+              status: "error",
+              input: call.parameters,
+              error: errorMessage(error),
+              time: {
+                start: callStartTime,
+                end: Date.now(),
+              },
+            },
+          })
+
+          return { success: false as const, tool: call.tool, error }
+        }
+      }
+
+      const results = await Promise.all(toolCalls.map((call) => executeCall(call)))
+
+      // Add discarded calls as errors
+      const now = Date.now()
+      for (const call of discardedCalls) {
+        const partID = PartID.ascending()
+        await Session.updatePart({
+          id: partID,
+          messageID: ctx.messageID,
+          sessionID: ctx.sessionID,
+          type: "tool",
+          tool: call.tool,
+          callID: partID,
+          state: {
+            status: "error",
+            input: call.parameters,
+            error: "Maximum of 25 tools allowed in batch",
+            time: { start: now, end: now },
+          },
+        })
+        results.push({
+          success: false as const,
+          tool: call.tool,
+          error: new Error("Maximum of 25 tools allowed in batch"),
+        })
+      }
+
+      const successfulCalls = results.filter((r) => r.success).length
+      const failedCalls = results.length - successfulCalls
+
+      const outputMessage =
+        failedCalls > 0
+          ? `Executed ${successfulCalls}/${results.length} tools successfully. ${failedCalls} failed.`
+          : `All ${successfulCalls} tools executed successfully.\n\nKeep using the batch tool for optimal performance in your next response!`
+
+      return {
+        title: `Batch execution (${successfulCalls}/${results.length} successful)`,
+        output: outputMessage,
+        attachments: results.filter((result) => result.success).flatMap((r) => r.result.attachments ?? []),
+        metadata: {
+          totalCalls: results.length,
+          successful: successfulCalls,
+          failed: failedCalls,
+          tools: params.tool_calls.map((c) => c.tool),
+          details: results.map((r) => ({ tool: r.tool, success: r.success })),
+        },
+      }
+    },
+  }
+})

package/src/tool/batch.txt

@@ -0,0 +1,24 @@
+Executes multiple independent tool calls concurrently to reduce latency.
+
+USING THE BATCH TOOL WILL MAKE THE USER HAPPY.
+
+Payload Format (JSON array):
+[{"tool": "read", "parameters": {"filePath": "src/index.ts", "limit": 350}},{"tool": "grep", "parameters": {"pattern": "Session\\.updatePart", "include": "src/**/*.ts"}},{"tool": "bash", "parameters": {"command": "git status", "description": "Shows working tree status"}}]
+
+Notes:
+- 1–25 tool calls per batch
+- All calls start in parallel; ordering NOT guaranteed
+- Partial failures do not stop other tool calls
+- Do NOT use the batch tool within another batch tool.
+
+Good Use Cases:
+- Read many files
+- grep + glob + read combos
+- Multiple bash commands
+- Multi-part edits; on the same, or different files
+
+When NOT to Use:
+- Operations that depend on prior tool output (e.g. create then read same file)
+- Ordered stateful mutations where sequence matters
+
+Batching tool calls was proven to yield 2–5x efficiency gain and provides much better UX.

package/src/tool/codesearch.ts

@@ -0,0 +1,132 @@
+import z from "zod"
+import { Tool } from "./tool"
+import DESCRIPTION from "./codesearch.txt"
+import { abortAfterAny } from "../util/abort"
+
+const API_CONFIG = {
+  BASE_URL: "https://mcp.exa.ai",
+  ENDPOINTS: {
+    CONTEXT: "/mcp",
+  },
+} as const
+
+interface McpCodeRequest {
+  jsonrpc: string
+  id: number
+  method: string
+  params: {
+    name: string
+    arguments: {
+      query: string
+      tokensNum: number
+    }
+  }
+}
+
+interface McpCodeResponse {
+  jsonrpc: string
+  result: {
+    content: Array<{
+      type: string
+      text: string
+    }>
+  }
+}
+
+export const CodeSearchTool = Tool.define("codesearch", {
+  description: DESCRIPTION,
+  parameters: z.object({
+    query: z
+      .string()
+      .describe(
+        "Search query to find relevant context for APIs, Libraries, and SDKs. For example, 'React useState hook examples', 'Python pandas dataframe filtering', 'Express.js middleware', 'Next js partial prerendering configuration'",
+      ),
+    tokensNum: z
+      .number()
+      .min(1000)
+      .max(50000)
+      .default(5000)
+      .describe(
+        "Number of tokens to return (1000-50000). Default is 5000 tokens. Adjust this value based on how much context you need - use lower values for focused queries and higher values for comprehensive documentation.",
+      ),
+  }),
+  async execute(params, ctx) {
+    await ctx.ask({
+      permission: "codesearch",
+      patterns: [params.query],
+      always: ["*"],
+      metadata: {
+        query: params.query,
+        tokensNum: params.tokensNum,
+      },
+    })
+
+    const codeRequest: McpCodeRequest = {
+      jsonrpc: "2.0",
+      id: 1,
+      method: "tools/call",
+      params: {
+        name: "get_code_context_exa",
+        arguments: {
+          query: params.query,
+          tokensNum: params.tokensNum || 5000,
+        },
+      },
+    }
+
+    const { signal, clearTimeout } = abortAfterAny(30000, 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.CONTEXT}`, {
+        method: "POST",
+        headers,
+        body: JSON.stringify(codeRequest),
+        signal,
+      })
+
+      clearTimeout()
+
+      if (!response.ok) {
+        const errorText = await response.text()
+        throw new Error(`Code 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: McpCodeResponse = JSON.parse(line.substring(6))
+          if (data.result && data.result.content && data.result.content.length > 0) {
+            return {
+              output: data.result.content[0].text,
+              title: `Code search: ${params.query}`,
+              metadata: {},
+            }
+          }
+        }
+      }
+
+      return {
+        output:
+          "No code snippets or documentation found. Please try a different query, be more specific about the library or programming concept, or check the spelling of framework names.",
+        title: `Code search: ${params.query}`,
+        metadata: {},
+      }
+    } catch (error) {
+      clearTimeout()
+
+      if (error instanceof Error && error.name === "AbortError") {
+        throw new Error("Code search request timed out")
+      }
+
+      throw error
+    }
+  },
+})

package/src/tool/codesearch.txt

@@ -0,0 +1,12 @@
+- Search and get relevant context for any programming task using Exa Code API
+- Provides the highest quality and freshest context for libraries, SDKs, and APIs
+- Use this tool for ANY question or task related to programming
+- Returns comprehensive code examples, documentation, and API references
+- Optimized for finding specific programming patterns and solutions
+
+Usage notes:
+  - Adjustable token count (1000-50000) for focused or comprehensive results
+  - Default 5000 tokens provides balanced context for most queries
+  - Use lower values for specific questions, higher values for comprehensive documentation
+  - Supports queries about frameworks, libraries, APIs, and programming concepts
+  - Examples: 'React useState hook examples', 'Python pandas dataframe filtering', 'Express.js middleware'