@oh-my-pi/pi-coding-agent 6.2.0 → 6.7.0

package/src/core/tools/complete.ts

@@ -4,8 +4,9 @@
  * Subagents must call this tool to finish and return structured JSON output.
  */
 
-import type { AgentTool } from "@oh-my-pi/pi-agent-core";
+import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import { StringEnum } from "@oh-my-pi/pi-ai";
+import type { Static, TObject } from "@sinclair/typebox";
 import { Type } from "@sinclair/typebox";
 import Ajv, { type ErrorObject, type ValidateFunction } from "ajv";
 import type { ToolSession } from "./index";
@@ -52,77 +53,86 @@ function formatAjvErrors(errors: ErrorObject[] | null | undefined): string {
 		.join("; ");
 }
 
-export function createCompleteTool(session: ToolSession) {
-	const schemaResult = normalizeSchema(session.outputSchema);
-	// Convert JTD to JSON Schema if needed (auto-detected)
-	const normalizedSchema =
-		schemaResult.normalized !== undefined ? jtdToJsonSchema(schemaResult.normalized) : undefined;
-	let validate: ValidateFunction | undefined;
-	let schemaError = schemaResult.error;
-
-	if (normalizedSchema !== undefined && !schemaError) {
-		try {
-			validate = ajv.compile(normalizedSchema as any);
-		} catch (err) {
-			schemaError = err instanceof Error ? err.message : String(err);
+export class CompleteTool implements AgentTool<TObject, CompleteDetails> {
+	public readonly name = "complete";
+	public readonly label = "Complete";
+	public readonly description =
+		"Finish the task with structured JSON output. Call exactly once at the end of the task.\n\n" +
+		"If you cannot complete the task, call with status='aborted' and an error message.";
+	public readonly parameters: TObject;
+
+	private readonly validate?: ValidateFunction;
+	private readonly schemaError?: string;
+
+	constructor(session: ToolSession) {
+		const schemaResult = normalizeSchema(session.outputSchema);
+		// Convert JTD to JSON Schema if needed (auto-detected)
+		const normalizedSchema =
+			schemaResult.normalized !== undefined ? jtdToJsonSchema(schemaResult.normalized) : undefined;
+		let schemaError = schemaResult.error;
+
+		if (normalizedSchema !== undefined && !schemaError) {
+			try {
+				this.validate = ajv.compile(normalizedSchema as any);
+			} catch (err) {
+				schemaError = err instanceof Error ? err.message : String(err);
+			}
 		}
+
+		this.schemaError = schemaError;
+
+		const schemaHint = formatSchema(normalizedSchema ?? session.outputSchema);
+
+		// Use actual schema if provided, otherwise fall back to Type.Any
+		// Merge description into the JSON schema for better tool documentation
+		const dataSchema = normalizedSchema
+			? Type.Unsafe({
+					...(normalizedSchema as object),
+					description: `Structured output matching the schema:\n${schemaHint}`,
+				})
+			: Type.Any({ description: "Structured JSON output (no schema specified)" });
+
+		this.parameters = Type.Object({
+			data: Type.Optional(dataSchema),
+			status: Type.Optional(
+				StringEnum(["success", "aborted"], {
+					description: "Use 'aborted' if the task cannot be completed, defaults to 'success'",
+				}),
+			),
+			error: Type.Optional(Type.String({ description: "Error message when status is 'aborted'" })),
+		});
 	}
 
-	const schemaHint = formatSchema(normalizedSchema ?? session.outputSchema);
-
-	// Use actual schema if provided, otherwise fall back to Type.Any
-	// Merge description into the JSON schema for better tool documentation
-	const dataSchema = normalizedSchema
-		? Type.Unsafe({
-				...(normalizedSchema as object),
-				description: `Structured output matching the schema:\n${schemaHint}`,
-			})
-		: Type.Any({ description: "Structured JSON output (no schema specified)" });
-
-	const completeParams = Type.Object({
-		data: Type.Optional(dataSchema),
-		status: Type.Optional(
-			StringEnum(["success", "aborted"], {
-				description: "Use 'aborted' if the task cannot be completed, defaults to 'success'",
-			}),
-		),
-		error: Type.Optional(Type.String({ description: "Error message when status is 'aborted'" })),
-	});
-
-	const tool: AgentTool<typeof completeParams, CompleteDetails> = {
-		name: "complete",
-		label: "Complete",
-		description:
-			"Finish the task with structured JSON output. Call exactly once at the end of the task.\n\n" +
-			"If you cannot complete the task, call with status='aborted' and an error message.",
-		parameters: completeParams,
-		execute: async (_toolCallId, params) => {
-			const status = params.status ?? "success";
-
-			// Skip validation when aborting - data is optional for aborts
-			if (status === "success") {
-				if (params.data === undefined) {
-					throw new Error("data is required when status is 'success'");
-				}
-				if (schemaError) {
-					throw new Error(`Invalid output schema: ${schemaError}`);
-				}
-				if (validate && !validate(params.data)) {
-					throw new Error(`Output does not match schema: ${formatAjvErrors(validate.errors)}`);
-				}
+	public async execute(
+		_toolCallId: string,
+		params: Static<TObject>,
+		_signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<CompleteDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<CompleteDetails>> {
+		const status = (params.status ?? "success") as "success" | "aborted";
+
+		// Skip validation when aborting - data is optional for aborts
+		if (status === "success") {
+			if (params.data === undefined) {
+				throw new Error("data is required when status is 'success'");
 			}
+			if (this.schemaError) {
+				throw new Error(`Invalid output schema: ${this.schemaError}`);
+			}
+			if (this.validate && !this.validate(params.data)) {
+				throw new Error(`Output does not match schema: ${formatAjvErrors(this.validate.errors)}`);
+			}
+		}
 
-			const responseText =
-				status === "aborted" ? `Task aborted: ${params.error || "No reason provided"}` : "Completion recorded.";
-
-			return {
-				content: [{ type: "text", text: responseText }],
-				details: { data: params.data, status, error: params.error },
-			};
-		},
-	};
+		const responseText =
+			status === "aborted" ? `Task aborted: ${params.error || "No reason provided"}` : "Completion recorded.";
 
-	return tool;
+		return {
+			content: [{ type: "text", text: responseText }],
+			details: { data: params.data, status, error: params.error as string | undefined },
+		};
+	}
 }
 
 // Register subprocess tool handler for extraction + termination.

package/src/core/tools/context.ts

@@ -11,31 +11,29 @@ declare module "@oh-my-pi/pi-agent-core" {
 	}
 }
 
-export interface ToolContextStore {
-	getContext(toolCall?: ToolCallContext): AgentToolContext;
-	setUIContext(uiContext: ExtensionUIContext, hasUI: boolean): void;
-	setToolNames(names: string[]): void;
-}
+export class ToolContextStore {
+	private uiContext: ExtensionUIContext | undefined;
+	private hasUI = false;
+	private toolNames: string[] = [];
 
-export function createToolContextStore(getBaseContext: () => CustomToolContext): ToolContextStore {
-	let uiContext: ExtensionUIContext | undefined;
-	let hasUI = false;
-	let toolNames: string[] = [];
+	constructor(private readonly getBaseContext: () => CustomToolContext) {}
 
-	return {
-		getContext: (toolCall) => ({
-			...getBaseContext(),
-			ui: uiContext,
-			hasUI,
-			toolNames,
+	getContext(toolCall?: ToolCallContext): AgentToolContext {
+		return {
+			...this.getBaseContext(),
+			ui: this.uiContext,
+			hasUI: this.hasUI,
+			toolNames: this.toolNames,
 			toolCall,
-		}),
-		setUIContext: (context, uiAvailable) => {
-			uiContext = context;
-			hasUI = uiAvailable;
-		},
-		setToolNames: (names) => {
-			toolNames = names;
-		},
-	};
+		};
+	}
+
+	setUIContext(uiContext: ExtensionUIContext, hasUI: boolean): void {
+		this.uiContext = uiContext;
+		this.hasUI = hasUI;
+	}
+
+	setToolNames(names: string[]): void {
+		this.toolNames = names;
+	}
 }

package/src/core/tools/exa/index.ts

@@ -48,11 +48,11 @@ export {
 	callExaTool,
 	callWebsetsTool,
 	createMCPToolFromServer,
-	createMCPWrappedTool,
 	fetchMCPToolSchema,
 	findApiKey,
 	formatSearchResults,
 	isSearchResponse,
+	MCPWrappedTool,
 } from "./mcp-client";
 export { renderExaCall, renderExaResult } from "./render";
 export { researcherTools } from "./researcher";

package/src/core/tools/exa/mcp-client.ts

@@ -1,14 +1,15 @@
 /**
  * Exa MCP Client
  *
- * Client for interacting with Exa MCP servers via JSON-RPC 2.0 over HTTPS.
+ * Client for interacting with Exa MCP servers.
  */
 
 import { existsSync, readFileSync } from "node:fs";
 import { homedir } from "node:os";
 import type { TSchema } from "@sinclair/typebox";
-import type { CustomTool } from "../../custom-tools/types";
+import type { CustomTool, CustomToolResult } from "../../custom-tools/types";
 import { logger } from "../../logger";
+import { callMCP } from "../../mcp/json-rpc";
 import type {
 	ExaRenderDetails,
 	ExaSearchResponse,
@@ -48,63 +49,6 @@ export async function findApiKey(): Promise<string | null> {
 	return null;
 }
 
-/** Parse SSE response format (lines starting with "data: ") */
-function parseSSE(text: string): unknown {
-	const lines = text.split("\n");
-	for (const line of lines) {
-		if (line.startsWith("data: ")) {
-			const data = line.slice(6).trim();
-			if (data === "[DONE]") continue;
-			try {
-				return JSON.parse(data);
-			} catch {
-				// Try next line
-			}
-		}
-	}
-	// Fallback: try parsing entire response as JSON
-	try {
-		return JSON.parse(text);
-	} catch {
-		return null;
-	}
-}
-
-/** Call MCP server with JSON-RPC 2.0 */
-export async function callMCP(url: string, method: string, params?: Record<string, unknown>): Promise<unknown> {
-	const body = {
-		jsonrpc: "2.0",
-		id: Math.random().toString(36).slice(2),
-		method,
-		params: params ?? {},
-	};
-
-	const response = await fetch(url, {
-		method: "POST",
-		headers: {
-			"Content-Type": "application/json",
-			Accept: "application/json, text/event-stream",
-		},
-		body: JSON.stringify(body),
-	});
-
-	if (!response.ok) {
-		const errorMsg = `MCP request failed: ${response.status} ${response.statusText}`;
-		logger.error(errorMsg, { url, method, params });
-		throw new Error(errorMsg);
-	}
-
-	const text = await response.text();
-	const result = parseSSE(text);
-
-	if (!result) {
-		logger.error("Failed to parse MCP response", { url, method, responseText: text.slice(0, 500) });
-		throw new Error("Failed to parse MCP response");
-	}
-
-	return result;
-}
-
 /** Fetch available tools from Exa MCP */
 export async function fetchExaTools(apiKey: string, toolNames: string[]): Promise<MCPTool[]> {
 	const url = `https://mcp.exa.ai/mcp?exaApiKey=${encodeURIComponent(apiKey)}&toolNames=${encodeURIComponent(toolNames.join(","))}`;
@@ -299,56 +243,67 @@ export async function fetchMCPToolSchema(
 }
 
 /**
- * Create a CustomTool dynamically from MCP tool metadata.
+ * CustomTool dynamically created from MCP tool metadata.
  *
  * This allows tools to be generated from MCP server schemas without hardcoding,
  * reducing drift when MCP servers add new parameters.
  */
-export function createMCPWrappedTool(
-	config: MCPToolWrapperConfig,
-	schema: TSchema,
-	description: string,
-): CustomTool<TSchema, ExaRenderDetails> {
-	return {
-		name: config.name,
-		label: config.label,
-		description,
-		parameters: schema,
-		async execute(_toolCallId, params, _onUpdate, _ctx, _signal) {
-			try {
-				const apiKey = await findApiKey();
-				if (!apiKey) {
-					return {
-						content: [{ type: "text" as const, text: "Error: EXA_API_KEY not found" }],
-						details: { error: "EXA_API_KEY not found", toolName: config.name },
-					};
-				}
-
-				const response = config.isWebsetsTool
-					? await callWebsetsTool(apiKey, config.mcpToolName, params as Record<string, unknown>)
-					: await callExaTool(config.mcpToolName, params as Record<string, unknown>, apiKey);
-
-				if (isSearchResponse(response)) {
-					const formatted = formatSearchResults(response);
-					return {
-						content: [{ type: "text" as const, text: formatted }],
-						details: { response, toolName: config.name },
-					};
-				}
+export class MCPWrappedTool implements CustomTool<TSchema, ExaRenderDetails> {
+	public readonly name: string;
+	public readonly label: string;
+	public readonly description: string;
+	public readonly parameters: TSchema;
+
+	private readonly config: MCPToolWrapperConfig;
+
+	constructor(config: MCPToolWrapperConfig, schema: TSchema, description: string) {
+		this.config = config;
+		this.name = config.name;
+		this.label = config.label;
+		this.description = description;
+		this.parameters = schema;
+	}
 
+	async execute(
+		_toolCallId: string,
+		params: unknown,
+		_onUpdate?: unknown,
+		_ctx?: unknown,
+		_signal?: AbortSignal,
+	): Promise<CustomToolResult<ExaRenderDetails>> {
+		try {
+			const apiKey = await findApiKey();
+			if (!apiKey) {
 				return {
-					content: [{ type: "text" as const, text: JSON.stringify(response, null, 2) }],
-					details: { raw: response, toolName: config.name },
+					content: [{ type: "text" as const, text: "Error: EXA_API_KEY not found" }],
+					details: { error: "EXA_API_KEY not found", toolName: this.config.name },
 				};
-			} catch (error) {
-				const message = error instanceof Error ? error.message : String(error);
+			}
+
+			const response = this.config.isWebsetsTool
+				? await callWebsetsTool(apiKey, this.config.mcpToolName, params as Record<string, unknown>)
+				: await callExaTool(this.config.mcpToolName, params as Record<string, unknown>, apiKey);
+
+			if (isSearchResponse(response)) {
+				const formatted = formatSearchResults(response);
 				return {
-					content: [{ type: "text" as const, text: `Error: ${message}` }],
-					details: { error: message, toolName: config.name },
+					content: [{ type: "text" as const, text: formatted }],
+					details: { response, toolName: this.config.name },
 				};
 			}
-		},
-	};
+
+			return {
+				content: [{ type: "text" as const, text: JSON.stringify(response, null, 2) }],
+				details: { raw: response, toolName: this.config.name },
+			};
+		} catch (error) {
+			const message = error instanceof Error ? error.message : String(error);
+			return {
+				content: [{ type: "text" as const, text: `Error: ${message}` }],
+				details: { error: message, toolName: this.config.name },
+			};
+		}
+	}
 }
 
 /**
@@ -361,9 +316,9 @@ export async function createMCPToolFromServer(
 	config: MCPToolWrapperConfig,
 	fallbackSchema: TSchema,
 	fallbackDescription: string,
-): Promise<CustomTool<TSchema, ExaRenderDetails>> {
+): Promise<MCPWrappedTool> {
 	const mcpTool = await fetchMCPToolSchema(apiKey, config.mcpToolName, config.isWebsetsTool);
 	const schema = mcpTool?.inputSchema ?? fallbackSchema;
 	const description = mcpTool?.description ?? fallbackDescription;
-	return createMCPWrappedTool(config, schema, description);
+	return new MCPWrappedTool(config, schema, description);
 }