@jarcelao/pi-exa-api 0.1.2 → 0.2.1

package/LICENSE.md

@@ -19,4 +19,3 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
-

package/README.md

@@ -1,7 +1,6 @@
 # pi-exa-api
 
-Web search and content fetching for [pi](https://pi.dev) via the [Exa API](https://exa.ai/).
-
+Web search, content fetching, and code context for [pi](https://pi.dev) via the [Exa API](https://exa.ai/).
 
 ## Installation
 
@@ -67,6 +66,23 @@ Fetch the content from https://example.com/article
   - `summary` - AI-generated summary
 - `maxCharacters` (optional) - Maximum characters to return (1000-100000)
 
+### Code Context
+
+The agent can use `exa_code_context` to find code snippets and examples from open source libraries and repositories:
+
+```
+Find examples of React hooks for state management
+```
+
+It's ideal for understanding how libraries, frameworks, or programming concepts are implemented in practice.
+
+**Parameters:**
+
+- `query` (required) - Search query for code snippets and examples (1-2000 characters)
+- `tokensNum` (optional) - Token limit for the response:
+  - `"dynamic"` (default) - Automatically determine optimal response length
+  - `50-100000` - Specific number of tokens (5000 is a good default, use 10000 when more context is needed)
+
 ## Development
 
 ```bash

package/exa-search.test.ts

@@ -5,6 +5,8 @@ import exaSearchExtension, {
   mapFetchContentType,
   formatSearchResults,
   formatFetchResult,
+  formatCodeContextResult,
+  parseCostDollars,
   createMissingApiKeyError,
 } from "./extensions/exa-search.ts";
 import {
@@ -284,8 +286,19 @@ describe("Fetch Result Formatting", () => {
     expect(formatted).toContain("Full page content here");
   });
 
+  it("should include cost in fetch details", () => {
+    const details = {
+      url: "https://example.com",
+      title: "Test Page",
+      cost: { total: 0.000123 },
+    };
+    expect(details.cost).toBeDefined();
+    expect(details.cost?.total).toBe(0.000123);
+  });
+
   it("should format highlights", () => {
     const result = {
+      title: "Test Page",
       url: "https://example.com",
       highlights: ["Key point 1", "Key point 2"],
     };
@@ -299,6 +312,7 @@ describe("Fetch Result Formatting", () => {
 
   it("should format summary", () => {
     const result = {
+      title: "Test Page",
       url: "https://example.com",
       summary: "This page is about...",
     };
@@ -312,9 +326,9 @@ describe("Fetch Result Formatting", () => {
     const result = {
       url: "https://example.com",
       text: "Content only",
-    };
+    } as { title?: string; url: string; text: string };
 
-    const formatted = formatFetchResult(result, "text");
+    const formatted = formatFetchResult(result as Parameters<typeof formatFetchResult>[0], "text");
     expect(formatted).toContain("https://example.com");
     expect(formatted).not.toContain("Title:");
   });
@@ -329,6 +343,204 @@ describe("Error Handling", () => {
   });
 });
 
+describe("parseCostDollars", () => {
+  it("should parse JSON string costDollars", () => {
+    const costString = '{"total":0.007,"search":{"neural":0.007}}';
+    const parsed = parseCostDollars(costString);
+    expect(parsed).toEqual({ total: 0.007, search: { neural: 0.007 } });
+    expect(parsed.total).toBe(0.007);
+  });
+
+  it("should pass through object costDollars unchanged", () => {
+    const costObject = { total: 1.5 };
+    const parsed = parseCostDollars(costObject);
+    expect(parsed).toEqual({ total: 1.5 });
+  });
+
+  it("should handle various cost string formats", () => {
+    expect(parseCostDollars('{"total":0}').total).toBe(0);
+    expect(parseCostDollars('{"total":123.456}').total).toBe(123.456);
+    expect(parseCostDollars({ total: 99.9 }).total).toBe(99.9);
+  });
+});
+
+describe("Code Context Result Formatting", () => {
+  it("should format code context response with string costDollars", () => {
+    const response = {
+      requestId: "req_12345",
+      query: "how to use React hooks for state management",
+      response: "## useState Example\n\n```javascript\nconst [count, setCount] = useState(0);\n```",
+      resultsCount: 502,
+      costDollars: '{"total":0.007,"search":{"neural":0.007}}',
+      searchTime: 1.234,
+      outputTokens: 4805,
+    };
+
+    const formatted = formatCodeContextResult(response);
+    expect(formatted).toContain("Query: how to use React hooks for state management");
+    expect(formatted).toContain("Results: 502 sources");
+    expect(formatted).toContain("Output tokens: 4805");
+    expect(formatted).toContain("--- Code Context ---");
+    expect(formatted).toContain("## useState Example");
+    expect(formatted).toContain("const [count, setCount] = useState(0);");
+    expect(formatted).toContain("Cost: $0.007000");
+  });
+
+  it("should format code context response with object costDollars", () => {
+    const response = {
+      requestId: "req_67890",
+      query: "test query",
+      response: "Some code examples...",
+      resultsCount: 10,
+      costDollars: { total: 1.5 },
+      searchTime: 0.5,
+      outputTokens: 1000,
+    };
+
+    const formatted = formatCodeContextResult(response);
+    expect(formatted).toContain("Cost: $1.500000");
+  });
+
+  it("should include query in formatted output", () => {
+    const response = {
+      requestId: "req_abc",
+      query: "Express.js middleware authentication",
+      response: "Some code examples...",
+      resultsCount: 100,
+      costDollars: '{"total":0.5}',
+      searchTime: 0.5,
+      outputTokens: 2000,
+    };
+
+    const formatted = formatCodeContextResult(response);
+    expect(formatted).toContain("Query: Express.js middleware authentication");
+  });
+
+  it("should format cost correctly with decimals", () => {
+    const response = {
+      requestId: "req_xyz",
+      query: "test query",
+      response: "response content",
+      resultsCount: 10,
+      costDollars: '{"total":0.123456}',
+      searchTime: 0.1,
+      outputTokens: 500,
+    };
+
+    const formatted = formatCodeContextResult(response);
+    expect(formatted).toContain("Cost: $0.123456");
+  });
+
+  it("should display results count and output tokens", () => {
+    const response = {
+      requestId: "req_test",
+      query: "pandas dataframe operations",
+      response: "Code examples here",
+      resultsCount: 150,
+      costDollars: '{"total":0.75}',
+      searchTime: 0.8,
+      outputTokens: 3500,
+    };
+
+    const formatted = formatCodeContextResult(response);
+    expect(formatted).toContain("Results: 150 sources");
+    expect(formatted).toContain("Output tokens: 3500");
+  });
+});
+
+describe("Code Context Parameter Validation", () => {
+  it("should accept valid params with query only", () => {
+    const params = { query: "React hooks examples" };
+    expect(params.query).toBeTruthy();
+    expect(params.query.trim()).not.toBe("");
+  });
+
+  it("should accept valid params with dynamic tokens", () => {
+    const params = {
+      query: "Express middleware",
+      tokensNum: "dynamic" as const,
+    };
+    expect(typeof params.query).toBe("string");
+    expect(params.tokensNum).toBe("dynamic");
+  });
+
+  it("should accept valid params with numeric tokens", () => {
+    const params = {
+      query: "Next.js configuration",
+      tokensNum: 5000,
+    };
+    expect(params.tokensNum).toBeGreaterThanOrEqual(50);
+    expect(params.tokensNum).toBeLessThanOrEqual(100000);
+  });
+
+  it("should accept token counts in valid range", () => {
+    const validTokenCounts = [50, 1000, 5000, 10000, 50000, 100000];
+    for (const count of validTokenCounts) {
+      expect(count).toBeGreaterThanOrEqual(50);
+      expect(count).toBeLessThanOrEqual(100000);
+    }
+  });
+
+  it("should accept query strings up to 2000 characters", () => {
+    const longQuery = "a".repeat(2000);
+    expect(longQuery.length).toBe(2000);
+  });
+});
+
+describe("tokensNum Type Coercion", () => {
+  // Helper function that mirrors the type coercion logic in exa_code_context execute
+  function coerceTokensNum(tokensNum: string | number | undefined): string | number {
+    const value = tokensNum ?? "dynamic";
+    if (typeof value === "string" && value !== "dynamic") {
+      return Number(value);
+    }
+    return value;
+  }
+
+  it("should return 'dynamic' when tokensNum is undefined", () => {
+    const result = coerceTokensNum(undefined);
+    expect(result).toBe("dynamic");
+    expect(typeof result).toBe("string");
+  });
+
+  it("should preserve 'dynamic' string", () => {
+    const result = coerceTokensNum("dynamic");
+    expect(result).toBe("dynamic");
+    expect(typeof result).toBe("string");
+  });
+
+  it("should preserve numeric tokens", () => {
+    const result = coerceTokensNum(5000);
+    expect(result).toBe(5000);
+    expect(typeof result).toBe("number");
+  });
+
+  it("should coerce numeric string to number", () => {
+    const result = coerceTokensNum("5000");
+    expect(result).toBe(5000);
+    expect(typeof result).toBe("number");
+  });
+
+  it("should coerce various numeric strings", () => {
+    expect(coerceTokensNum("1000")).toBe(1000);
+    expect(coerceTokensNum("10000")).toBe(10000);
+    expect(coerceTokensNum("3000")).toBe(3000);
+    expect(typeof coerceTokensNum("5000")).toBe("number");
+  });
+
+  it("should handle boundary token values", () => {
+    expect(coerceTokensNum(50)).toBe(50);
+    expect(coerceTokensNum(100000)).toBe(100000);
+  });
+
+  it("should not modify actual numbers", () => {
+    const originalNumber = 7500;
+    const result = coerceTokensNum(originalNumber);
+    expect(result).toBe(originalNumber);
+    expect(result).toBe(7500);
+  });
+});
+
 describe("Extension Registration", () => {
   function createMockExtensionAPI() {
     const tools: unknown[] = [];
@@ -356,7 +568,7 @@ describe("Extension Registration", () => {
     exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
 
     const tools = api.getTools();
-    expect(tools.length).toBeGreaterThanOrEqual(2);
+    expect(tools.length).toBe(3);
 
     const exaSearchTool = tools.find((t: unknown) => (t as { name: string }).name === "exa_search");
     expect(exaSearchTool).toBeDefined();
@@ -375,6 +587,30 @@ describe("Extension Registration", () => {
     expect((exaFetchTool as { label: string }).label).toBe("Exa Fetch");
   });
 
+  it("should register exa_code_context tool", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const codeContextTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_code_context",
+    );
+    expect(codeContextTool).toBeDefined();
+    expect((codeContextTool as { name: string }).name).toBe("exa_code_context");
+    expect((codeContextTool as { label: string }).label).toBe("Exa Code Context");
+  });
+
+  it("should have execute function on exa_code_context tool", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const codeContextTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_code_context",
+    ) as { execute: unknown };
+    expect(typeof codeContextTool?.execute).toBe("function");
+  });
+
   it("should register /exa-status command", () => {
     const api = createMockExtensionAPI();
     exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
@@ -415,6 +651,162 @@ describe("Extension Registration", () => {
     ) as { execute: unknown };
     expect(typeof exaFetchTool?.execute).toBe("function");
   });
+
+  it("should display cost in exa_fetch renderResult", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const exaFetchTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_fetch",
+    ) as { renderResult: Function };
+
+    const mockTheme = {
+      fg: (name: string, text: string) => text,
+    };
+
+    const mockResult = {
+      content: [{ type: "text", text: "test" }],
+      details: {
+        url: "https://example.com",
+        title: "Test Page",
+        cost: { total: 0.000123 },
+      },
+    };
+
+    const rendered = exaFetchTool.renderResult(
+      mockResult,
+      { expanded: false, isPartial: false },
+      mockTheme,
+    );
+
+    expect(rendered.text).toContain("Test Page");
+    expect(rendered.text).toContain("$0.000123");
+  });
+
+  it("should display cost without title in exa_fetch renderResult", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const exaFetchTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_fetch",
+    ) as { renderResult: Function };
+
+    const mockTheme = {
+      fg: (name: string, text: string) => text,
+    };
+
+    const mockResult = {
+      content: [{ type: "text", text: "test" }],
+      details: {
+        url: "https://example.com",
+        cost: { total: 0.000456 },
+      },
+    };
+
+    const rendered = exaFetchTool.renderResult(
+      mockResult,
+      { expanded: false, isPartial: false },
+      mockTheme,
+    );
+
+    expect(rendered.text).toContain("Fetched");
+    expect(rendered.text).toContain("$0.000456");
+  });
+
+  it("should display stats and cost in exa_code_context renderResult", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const codeContextTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_code_context",
+    ) as { renderResult: Function };
+
+    const mockTheme = {
+      fg: (name: string, text: string) => text,
+    };
+
+    const mockResult = {
+      content: [{ type: "text", text: "test" }],
+      details: {
+        query: "React hooks",
+        resultsCount: 502,
+        outputTokens: 4805,
+        cost: { total: 1.0 },
+      },
+    };
+
+    const rendered = codeContextTool.renderResult(
+      mockResult,
+      { expanded: false, isPartial: false },
+      mockTheme,
+    );
+
+    expect(rendered.text).toContain("502 sources");
+    expect(rendered.text).toContain("4805 tokens");
+    expect(rendered.text).toContain("$1.000000");
+  });
+
+  it("should display exa_code_context renderResult without cost", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const codeContextTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_code_context",
+    ) as { renderResult: Function };
+
+    const mockTheme = {
+      fg: (name: string, text: string) => text,
+    };
+
+    const mockResult = {
+      content: [{ type: "text", text: "test" }],
+      details: {
+        query: "Express middleware",
+        resultsCount: 100,
+        outputTokens: 2000,
+      },
+    };
+
+    const rendered = codeContextTool.renderResult(
+      mockResult,
+      { expanded: false, isPartial: false },
+      mockTheme,
+    );
+
+    expect(rendered.text).toContain("100 sources");
+    expect(rendered.text).toContain("2000 tokens");
+    expect(rendered.text).not.toContain("$");
+  });
+
+  it("should handle exa_code_context renderResult without details", () => {
+    const api = createMockExtensionAPI();
+    exaSearchExtension(api as unknown as import("@mariozechner/pi-coding-agent").ExtensionAPI);
+
+    const tools = api.getTools();
+    const codeContextTool = tools.find(
+      (t: unknown) => (t as { name: string }).name === "exa_code_context",
+    ) as { renderResult: Function };
+
+    const mockTheme = {
+      fg: (name: string, text: string) => text,
+    };
+
+    const mockResult = {
+      content: [{ type: "text", text: "Some code context output here" }],
+    };
+
+    const rendered = codeContextTool.renderResult(
+      mockResult,
+      { expanded: false, isPartial: false },
+      mockTheme,
+    );
+
+    expect(rendered.text).toContain("Some code context");
+  });
 });
 
 describe("Tool Execute Validation", () => {

package/extensions/exa-search.ts

@@ -1,22 +1,26 @@
 /**
  * exa-search Extension
  *
- * Registers two tools for web search and content fetching using the Exa API:
+ * Registers three tools for web search, content fetching, and code context using the Exa API:
  * - exa_search: Natural language web search
  * - exa_fetch: Fetch and extract content from URLs
+ * - exa_code_context: Search for code snippets and examples from open source repos
  *
  * Also registers the /exa-status command to check API key configuration.
  */
 
+import { mkdtemp, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
-import { Text } from "@mariozechner/pi-tui";
-import { Type, type Static } from "@sinclair/typebox";
 import {
-  truncateHead,
   DEFAULT_MAX_BYTES,
   DEFAULT_MAX_LINES,
   formatSize,
+  truncateHead,
 } from "@mariozechner/pi-coding-agent";
+import { Text } from "@mariozechner/pi-tui";
+import { Type, type Static } from "@mariozechner/pi-ai";
 import Exa from "exa-js";
 
 // API Key Management
@@ -26,6 +30,15 @@ function getApiKey(): string | undefined {
   return key && key.length > 0 ? key : undefined;
 }
 
+// Temp File Helper
+
+async function writeTempFile(content: string): Promise<string> {
+  const tempDir = await mkdtemp(join(tmpdir(), "pi-exa-"));
+  const tempFile = join(tempDir, "output.txt");
+  await writeFile(tempFile, content, "utf8");
+  return tempFile;
+}
+
 // Type Definitions
 
 type SearchContentType = "text" | "highlights" | "summary" | "none";
@@ -40,6 +53,24 @@ interface SearchDetails {
 interface FetchDetails {
   url: string;
   title?: string;
+  cost?: { total: number };
+}
+
+interface CodeContextDetails {
+  query: string;
+  resultsCount: number;
+  outputTokens: number;
+  cost?: { total: number };
+}
+
+interface CodeContextResponse {
+  requestId: string;
+  query: string;
+  response: string;
+  resultsCount: number;
+  costDollars: string | { total: number };
+  searchTime: number;
+  outputTokens: number;
 }
 
 // Content Type Mapping
@@ -170,6 +201,31 @@ function formatFetchResult(result: ExaSearchResult, contentType: FetchContentTyp
   return lines.join("\n");
 }
 
+function parseCostDollars(costDollars: string | { total: number }): { total: number } {
+  if (typeof costDollars === "string") {
+    return JSON.parse(costDollars);
+  }
+  return costDollars;
+}
+
+function formatCodeContextResult(response: CodeContextResponse): string {
+  const lines: string[] = [];
+
+  lines.push(`Query: ${response.query}`);
+  lines.push(`Results: ${response.resultsCount} sources`);
+  lines.push(`Output tokens: ${response.outputTokens}`);
+  lines.push("");
+  lines.push("--- Code Context ---");
+  lines.push("");
+  lines.push(response.response);
+  lines.push("");
+
+  const cost = parseCostDollars(response.costDollars);
+  lines.push(`Cost: $${cost.total.toFixed(6)}`);
+
+  return lines.join("\n");
+}
+
 // Error Creation
 
 function createMissingApiKeyError(): Error {
@@ -180,7 +236,16 @@ function createMissingApiKeyError(): Error {
 
 // Exports
 
-export { getApiKey, mapSearchContentType, mapFetchContentType, formatSearchResults, formatFetchResult, createMissingApiKeyError };
+export {
+  getApiKey,
+  mapSearchContentType,
+  mapFetchContentType,
+  formatSearchResults,
+  formatFetchResult,
+  formatCodeContextResult,
+  parseCostDollars,
+  createMissingApiKeyError,
+};
 
 export default function exaSearchExtension(pi: ExtensionAPI): void {
   pi.on("session_start", async (_event: unknown, ctx: ExtensionContext) => {
@@ -248,10 +313,7 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
         response = await exa.search(params.query, searchOptions);
       } catch (error) {
         const message = error instanceof Error ? error.message : String(error);
-        return {
-          content: [{ type: "text", text: `Exa API error: ${message}` }],
-          details: { query: params.query, numResults: 0 } as SearchDetails,
-        };
+        throw new Error(`Exa API error: ${message}`);
       }
 
       let output = formatSearchResults({
@@ -267,8 +329,10 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
       let result = truncation.content;
 
       if (truncation.truncated) {
+        const tempFile = await writeTempFile(output);
         result += `\n\n[Output truncated: ${truncation.outputLines} of ${truncation.totalLines} lines`;
-        result += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).]`;
+        result += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).`;
+        result += ` Full output saved to: ${tempFile}]`;
       }
 
       return {
@@ -291,7 +355,7 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
       return new Text(text, 0, 0);
     },
 
-    renderResult(result, { expanded: _expanded }, theme) {
+    renderResult(result, _options, theme) {
       const details = result.details as SearchDetails | undefined;
 
       if (!details) {
@@ -361,16 +425,13 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
         response = await exa.getContents(params.url, contentsOptions);
       } catch (error) {
         const message = error instanceof Error ? error.message : String(error);
-        return {
-          content: [{ type: "text", text: `Exa API error: ${message}` }],
-          details: { url: params.url } as FetchDetails,
-        };
+        throw new Error(`Exa API error: ${message}`);
       }
 
       if (!response.results || response.results.length === 0) {
         return {
           content: [{ type: "text", text: "No content found at this URL." }],
-          details: { url: params.url } as FetchDetails,
+          details: { url: params.url, cost: response.costDollars } as FetchDetails,
         };
       }
 
@@ -385,8 +446,10 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
       let content = truncation.content;
 
       if (truncation.truncated) {
+        const tempFile = await writeTempFile(output);
         content += `\n\n[Output truncated: ${truncation.outputLines} of ${truncation.totalLines} lines`;
-        content += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).]`;
+        content += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).`;
+        content += ` Full output saved to: ${tempFile}]`;
       }
 
       return {
@@ -394,6 +457,7 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
         details: {
           url: params.url,
           title: result.title,
+          cost: response.costDollars,
         } as FetchDetails,
       };
     },
@@ -408,7 +472,7 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
       return new Text(text, 0, 0);
     },
 
-    renderResult(result, { expanded: _expanded }, theme) {
+    renderResult(result, _options, theme) {
       const details = result.details as FetchDetails | undefined;
 
       if (!details) {
@@ -416,11 +480,143 @@ export default function exaSearchExtension(pi: ExtensionAPI): void {
         return new Text(text?.type === "text" ? text.text.slice(0, 60) : "", 0, 0);
       }
 
+      const cost = details.cost ? ` • $${details.cost.total.toFixed(6)}` : "";
+
       if (details.title) {
-        return new Text(theme.fg("success", "✓ ") + theme.fg("accent", details.title), 0, 0);
+        return new Text(theme.fg("success", `✓ ${details.title}${cost}`), 0, 0);
+      }
+
+      return new Text(theme.fg("success", `✓ Fetched${cost}`), 0, 0);
+    },
+  });
+
+  // Register exa_code_context tool
+
+  const ExaCodeContextParams = Type.Object({
+    query: Type.String({
+      description: "Search query to find relevant code snippets and examples",
+    }),
+    tokensNum: Type.Optional(
+      Type.Union([
+        Type.String({
+          description: 'Token limit: "dynamic" for automatic sizing',
+        }),
+        Type.Number({
+          description: "Token limit: 50-100000 (5000 is a good default)",
+        }),
+      ]),
+    ),
+  });
+
+  pi.registerTool({
+    name: "exa_code_context",
+    label: "Exa Code Context",
+    description:
+      "Search for code snippets and examples from open source libraries and repositories. Use this to find working code examples that help understand how libraries, frameworks, or concepts are implemented.",
+    parameters: ExaCodeContextParams,
+
+    async execute(
+      _toolCallId: string,
+      params: Static<typeof ExaCodeContextParams>,
+      _signal: AbortSignal | undefined,
+      _onUpdate: unknown,
+      _ctx: ExtensionContext,
+    ) {
+      const apiKey = getApiKey();
+      if (!apiKey) {
+        throw createMissingApiKeyError();
       }
 
-      return new Text(theme.fg("muted", "Done"), 0, 0);
+      // Ensure tokensNum is the correct type: number or "dynamic"
+      // The schema accepts both string and number, but the Exa API requires:
+      // - A number (e.g., 5000)
+      // - The literal string "dynamic"
+      let tokensNum: string | number = params.tokensNum ?? "dynamic";
+      if (typeof tokensNum === "string" && tokensNum !== "dynamic") {
+        tokensNum = Number(tokensNum);
+      }
+
+      let response;
+      try {
+        const httpResponse = await fetch("https://api.exa.ai/context", {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            "x-api-key": apiKey,
+          },
+          body: JSON.stringify({
+            query: params.query,
+            tokensNum,
+          }),
+        });
+
+        if (!httpResponse.ok) {
+          const errorText = await httpResponse.text();
+          throw new Error(`HTTP ${httpResponse.status}: ${errorText}`);
+        }
+
+        response = (await httpResponse.json()) as CodeContextResponse;
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        throw new Error(`Exa Context API error: ${message}`);
+      }
+
+      let output = formatCodeContextResult(response);
+
+      const truncation = truncateHead(output, {
+        maxLines: DEFAULT_MAX_LINES,
+        maxBytes: DEFAULT_MAX_BYTES,
+      });
+
+      let result = truncation.content;
+
+      if (truncation.truncated) {
+        const tempFile = await writeTempFile(output);
+        result += `\n\n[Output truncated: ${truncation.outputLines} of ${truncation.totalLines} lines`;
+        result += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).`;
+        result += ` Full output saved to: ${tempFile}]`;
+      }
+
+      const cost = parseCostDollars(response.costDollars);
+
+      return {
+        content: [{ type: "text", text: result }],
+        details: {
+          query: params.query,
+          resultsCount: response.resultsCount,
+          outputTokens: response.outputTokens,
+          cost,
+        } as CodeContextDetails,
+      };
+    },
+
+    renderCall(args, theme) {
+      const preview = args.query.length > 50 ? args.query.slice(0, 50) + "..." : args.query;
+      const desc = `${args.tokensNum ?? "dynamic"} tokens`;
+      const text =
+        theme.fg("toolTitle", theme.bold("exa_code_context ")) +
+        theme.fg("muted", preview) +
+        theme.fg("dim", ` ${desc}`);
+      return new Text(text, 0, 0);
+    },
+
+    renderResult(result, _options, theme) {
+      const details = result.details as CodeContextDetails | undefined;
+
+      if (!details) {
+        const text = result.content[0];
+        return new Text(text?.type === "text" ? text.text.slice(0, 60) : "", 0, 0);
+      }
+
+      const cost = details.cost ? ` • $${details.cost.total.toFixed(6)}` : "";
+      return new Text(
+        theme.fg(
+          "success",
+          `✓ ${details.resultsCount} sources • ${details.outputTokens} tokens${cost}`,
+        ),
+        0,
+        0,
+      );
     },
   });
 

package/package.json

@@ -1,17 +1,17 @@
 {
   "name": "@jarcelao/pi-exa-api",
-  "version": "0.1.2",
+  "version": "0.2.1",
   "description": "Web search and content fetching for pi via the Exa API",
-  "author": "Jericho Renniel Arcelao <jerichoarcelao@gmail.com>",
+  "keywords": [
+    "pi-package"
+  ],
+  "homepage": "https://github.com/jarcelao/pi-exa-api",
   "license": "MIT",
+  "author": "Jericho Renniel Arcelao <jerichoarcelao@gmail.com>",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/jarcelao/pi-exa-api.git"
   },
-  "homepage": "https://github.com/jarcelao/pi-exa-api",
-  "keywords": [
-    "pi-package"
-  ],
   "main": "./extensions/exa-search.ts",
   "scripts": {
     "test": "vitest",
@@ -22,7 +22,7 @@
     "format:check": "oxfmt --check"
   },
   "dependencies": {
-    "exa-js": "^1.1.1"
+    "exa-js": "^2.10.1"
   },
   "devDependencies": {
     "@types/node": "^25.5.0",

package/tsconfig.json

@@ -8,8 +8,9 @@
     "skipLibCheck": true,
     "noEmit": true,
     "resolveJsonModule": true,
+    "allowImportingTsExtensions": true,
     "types": ["node"]
   },
-  "include": ["*.ts"],
+  "include": ["*.ts", "extensions/**/*.ts"],
   "exclude": ["node_modules"]
 }