From 6a4bfa31a925af67c026cd68700fe7779fd7c4cb Mon Sep 17 00:00:00 2001 From: amankumarsingh77 Date: Sat, 2 May 2026 16:00:58 +0530 Subject: feat: add defineBrowserTool helper, truncation, schemas/common, empty registry Co-Authored-By: Claude Sonnet 4.6 --- src/util/tool.ts | 158 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/util/truncate.ts | 35 ++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 src/util/tool.ts create mode 100644 src/util/truncate.ts (limited to 'src/util') diff --git a/src/util/tool.ts b/src/util/tool.ts new file mode 100644 index 0000000..fc1c5e2 --- /dev/null +++ b/src/util/tool.ts @@ -0,0 +1,158 @@ +import type { TSchema, Static } from "typebox"; +import type { + AgentToolResult, + AgentToolUpdateCallback, + ExtensionAPI, + ExtensionContext, + Theme, + ToolDefinition, + ToolRenderResultOptions, +} from "@mariozechner/pi-coding-agent"; +import type { Component } from "@mariozechner/pi-tui"; +import type { BrowserClient } from "../client"; +import type { Result } from "./result"; + +export type ToolOk = { + readonly text: string; + readonly details?: Readonly>; +}; + +export type ToolErrKind = + | "not_connected" + | "cdp_error" + | "timeout" + | "invalid_state" + | "io_error" + | "internal"; + +export type ToolErr = { + readonly kind: ToolErrKind; + readonly message: string; + readonly details?: Readonly>; +}; + +export type HandlerContext = { + readonly client: BrowserClient; + readonly signal: AbortSignal | undefined; + readonly onUpdate: (update: ToolOk) => void; + readonly extensionCtx: ExtensionContext; +}; + +export type ToolHandler = ( + args: Static, + ctx: HandlerContext, +) => Promise>; + +export type BrowserToolDefinition = { + readonly name: string; + readonly label: string; + readonly description: string; + readonly promptSnippet: string; + readonly promptGuidelines: ReadonlyArray; + readonly parameters: S; + readonly handler: ToolHandler; + readonly renderCall?: (args: Static, theme: Theme) => Component; + readonly renderResult?: (result: AgentToolResult, expanded: boolean, theme: Theme) => Component; + readonly ensureAlive?: boolean; +}; + +// Type-erased entry for arrays. AnyBrowserToolDefinition lets the registry +// hold a heterogeneous list without losing the per-tool Static binding. +export type AnyBrowserToolDefinition = BrowserToolDefinition; + +export const defineBrowserTool = ( + def: BrowserToolDefinition, +): BrowserToolDefinition => def; + +type OkDetails = { readonly ok: true } & Readonly>; +type ErrDetails = { readonly ok: false; readonly kind: ToolErrKind; readonly message: string } & Readonly>; + +const toToolResult = ( + r: Result, + toolName: string, +): AgentToolResult => { + if (r.success) { + const details: OkDetails = { ok: true, ...(r.data.details ?? {}) }; + return { + content: [{ type: "text", text: r.data.text }], + details, + }; + } + const details: ErrDetails = { + ok: false, + kind: r.error.kind, + message: r.error.message, + ...(r.error.details ?? {}), + }; + return { + // isError is an extension the pi-coding-agent runtime picks up from the result object; + // AgentToolResult itself doesn't declare it but the runtime reads it via duck-typing. + // We cast through unknown to avoid widening the return type. + ...(({ isError: true }) as unknown as object), + content: [{ type: "text", text: `${toolName} failed (${r.error.kind}): ${r.error.message}` }], + details, + }; +}; + +export const registerBrowserTool = ( + pi: ExtensionAPI, + client: BrowserClient, + def: BrowserToolDefinition, +): void => { + const td: ToolDefinition = { + name: def.name, + label: def.label, + description: def.description, + promptSnippet: def.promptSnippet, + promptGuidelines: [...def.promptGuidelines], + parameters: def.parameters, + ...(def.renderCall + ? { + renderCall: (args: Static, theme: Theme, _ctx: Parameters["renderCall"]>>[2]) => + // Non-null assertion safe: this branch only runs when renderCall was provided. + def.renderCall!(args, theme), + } + : {}), + ...(def.renderResult + ? { + renderResult: ( + result: AgentToolResult, + options: ToolRenderResultOptions, + theme: Theme, + _ctx: Parameters["renderResult"]>>[3], + ) => + // Non-null assertion safe: this branch only runs when renderResult was provided. + def.renderResult!(result, options.expanded ?? false, theme), + } + : {}), + async execute(_toolCallId, args, signal, onUpdate, extensionCtx) { + if (def.ensureAlive !== false) { + const alive = await client.ensureAlive(); + if (!alive.success) { + return toToolResult( + { success: false, error: { kind: "not_connected", message: alive.error.message } }, + def.name, + ); + } + } + const result = await def.handler(args, { + client, + signal, + onUpdate: (u) => { + if (onUpdate) { + // AgentToolUpdateCallback expects AgentToolResult; we use unknown for TDetails. + const update: AgentToolResult = { + content: [{ type: "text", text: u.text }], + details: { ok: true, ...(u.details ?? {}) }, + }; + // Cast needed: onUpdate is AgentToolUpdateCallback (generic TDetails). + (onUpdate as AgentToolUpdateCallback)(update); + } + }, + extensionCtx, + }); + return toToolResult(result, def.name); + }, + }; + pi.registerTool(td); +}; diff --git a/src/util/truncate.ts b/src/util/truncate.ts new file mode 100644 index 0000000..2bf0600 --- /dev/null +++ b/src/util/truncate.ts @@ -0,0 +1,35 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + DEFAULT_MAX_BYTES, + DEFAULT_MAX_LINES, + formatSize, + truncateHead, + withFileMutationQueue, +} from "@mariozechner/pi-coding-agent"; + +const tempDirs: string[] = []; + +export const cleanupTempDirs = async (): Promise => { + const dirs = tempDirs.splice(0); + await Promise.all(dirs.map((d) => rm(d, { recursive: true, force: true }).catch(() => {}))); +}; + +export type TruncatedOutput = { + readonly text: string; + readonly fullOutputPath?: string; + readonly wasTruncated: boolean; +}; + +export const applyTruncation = async (output: string, prefix: string): Promise => { + const t = truncateHead(output, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES }); + if (!t.truncated) return { text: t.content, wasTruncated: false }; + const dir = await mkdtemp(join(tmpdir(), `pi-bh-${prefix}-`)); + tempDirs.push(dir); + const file = join(dir, "output.txt"); + await withFileMutationQueue(file, async () => { await writeFile(file, output, "utf8"); }); + const omitted = t.totalBytes - t.outputBytes; + const text = `${t.content}\n\n[Output truncated: ${t.outputLines} of ${t.totalLines} lines (${formatSize(t.outputBytes)} of ${formatSize(t.totalBytes)}). ${formatSize(omitted)} omitted. Full output: ${file}]`; + return { text, fullOutputPath: file, wasTruncated: true }; +}; -- cgit v1.2.3