From ad1f168133840fb00258d31c15ec777635a5aa11 Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 19 Aug 2026 17:16:36 -0300 Subject: [PATCH] feat(omp): route broad tool calls and cap direct results --- src/adapters/omp/plugin.ts | 78 +++----- src/adapters/omp/routing-guard.ts | 290 ++++++++++++++++++++++++++++++ tests/adapters/omp-plugin.test.ts | 105 ++++++++++- 3 files changed, 416 insertions(+), 57 deletions(-) create mode 100644 src/adapters/omp/routing-guard.ts diff --git a/src/adapters/omp/plugin.ts b/src/adapters/omp/plugin.ts index 53061c044..145292151 100644 --- a/src/adapters/omp/plugin.ts +++ b/src/adapters/omp/plugin.ts @@ -2,13 +2,14 @@ * Oh My Pi (OMP) plugin entry point for context-mode. * * Mirrors the Pi extension shape (`src/adapters/pi/extension.ts`) for - * the four OMP hook events that materially protect the context window - * and persist session continuity: + * the OMP hook events that protect the context window and persist session + * continuity: * * - session_start — initialize the session row in our DB - * - tool_call — hard-block curl/wget/inline-HTTP in bash - * - tool_result — extract structured events into the session DB + * - tool_call — route broad analysis and block inline HTTP + * - tool_result — bound direct output and extract session events * - session_before_compact — persist a resume snapshot before compaction + * - turn_end — persist per-turn token and cost usage * * Loaded by OMP via the `omp` (or `pi`) field in package.json — see * upstream loader at refs/platforms/oh-my-pi/packages/coding-agent/src/ @@ -35,6 +36,8 @@ import type { HookInput } from "../../session/extract.js"; import { buildResumeSnapshot } from "../../session/snapshot.js"; import type { SessionEvent } from "../../types.js"; import { OMPAdapter } from "./index.js"; +import { classifyOmpToolCall, limitOmpToolResult } from "./routing-guard.js"; +import type { ToolResultReplacement } from "./routing-guard.js"; import { parseOmpUsage } from "./usage.js"; // ── Tool-name normalization ───────────────────────────── @@ -50,22 +53,6 @@ const OMP_TOOL_MAP: Record = { view: "Read", }; -// ── Routing patterns ───────────────────────────────────── -// Inline HTTP client patterns to hard-block in bash. Identical to the -// Pi extension list (src/adapters/pi/extension.ts:42). One unrouted -// curl can dump 56 KB into context. -const BLOCKED_BASH_PATTERNS: RegExp[] = [ - /\bcurl\s/, - /\bwget\s/, - /\bfetch\s*\(/, - /\brequests\.get\s*\(/, - /\brequests\.post\s*\(/, - /\bhttp\.get\s*\(/, - /\bhttp\.request\s*\(/, - /\burllib\.request/, - /\bInvoke-WebRequest\b/, -]; - // ── Module-level singletons ────────────────────────────── // Same shape as Pi: one DB per process, session ID rebound on each // session_start so multi-session reuse within a long-lived plugin @@ -235,7 +222,7 @@ export interface MinimalHookAPI { on(event: "session_start", handler: HookHandler<{ type: "session_start" }>): void; on(event: "session_before_compact", handler: HookHandler<{ type: "session_before_compact" }>): void; on(event: "tool_call", handler: HookHandler): void; - on(event: "tool_result", handler: HookHandler): void; + on(event: "tool_result", handler: HookHandler): void; // turn_end carries a single per-turn AssistantMessage with `.usage`/`.model` // (refs/.../extensibility/shared-events.ts:204-208). agent_end carries // `messages: AssistantMessage[]` (:191-194) — both flow through parseOmpUsage. @@ -247,9 +234,9 @@ export interface MinimalHookAPI { /** * OMP plugin default export. Called once by the OMP runtime per - * upstream `extensibility/plugins/loader.ts` after `omp plugin install - * context-mode`. Subsequent `pi.on(...)` registrations route the four - * lifecycle events to our SessionDB-backed handlers below. + * upstream `extensibility/plugins/loader.ts` after `omp plugin install context-mode`. + * Subsequent `pi.on(...)` registrations route lifecycle events to the + * SessionDB-backed handlers below. */ export default function ompPlugin(pi: MinimalHookAPI): void { // OMP upstream uses PI_-prefixed env vars only (verified against @@ -279,40 +266,21 @@ export default function ompPlugin(pi: MinimalHookAPI): void { return undefined; }); - // ── 2. tool_call — pre-tool-call hard-block ─────────── - // Returning `{block: true, reason}` per - // refs/.../hooks/types.ts:566 (ToolCallEventResult) terminates the - // tool call with the reason surfaced to the LLM. + // ── 2. tool_call — routing enforcement ───────────────── + // The guard covers broad analysis, inline HTTP, and oversized-output + // fallbacks while leaving scoped reads, edits, and unknown tools alone. pi.on("tool_call", (event) => { - try { - const toolName = String(event?.toolName ?? "").toLowerCase(); - if (toolName !== "bash") return undefined; - - const command = String((event?.input as { command?: unknown } | undefined)?.command ?? ""); - if (!command) return undefined; - - const isBlocked = BLOCKED_BASH_PATTERNS.some((p) => p.test(command)); - if (isBlocked) { - return { - block: true, - reason: - "Use context-mode MCP tools (ctx_execute, ctx_fetch_and_index) instead of inline HTTP. " + - "curl/wget/fetch dump raw HTTP into the context window.", - }; - } - } catch { - // routing failure → allow passthrough - } - return undefined; + return classifyOmpToolCall(event, { cwd: projectDir }); }); - // ── 3. tool_result — post-tool-call event capture ───── - // OMP `tool_result` payload (refs/.../hooks/types.ts:461 onward) is - // `{toolName, toolCallId, input, content[], isError}`. We adapt to - // the Claude Code-shaped HookInput consumed by extractEvents. + // ── 3. tool_result — output bound + event capture ─────── + // Apply the model-facing bound before the session guard so a malformed or + // sessionless call cannot bypass the fallback limiter. Persist the original + // result for session search; only the returned value is bounded. pi.on("tool_result", (event) => { + const replacement = limitOmpToolResult(event); try { - if (!_sessionId) return undefined; + if (!_sessionId) return replacement; const rawToolName = String(event?.toolName ?? ""); const mappedToolName = OMP_TOOL_MAP[rawToolName.toLowerCase()] ?? rawToolName; @@ -335,9 +303,9 @@ export default function ompPlugin(pi: MinimalHookAPI): void { db.insertEvent(_sessionId, ev as SessionEvent, "PostToolUse"); } } catch { - // best effort + // best effort — never break a tool result or its replacement } - return undefined; + return replacement; }); // ── 4. session_before_compact — resume snapshot ─────── diff --git a/src/adapters/omp/routing-guard.ts b/src/adapters/omp/routing-guard.ts new file mode 100644 index 000000000..e0e07ddfa --- /dev/null +++ b/src/adapters/omp/routing-guard.ts @@ -0,0 +1,290 @@ +import { statSync } from "node:fs"; +import { resolve } from "node:path"; + +const MAX_READ_FILE_BYTES = 32 * 1024; +export const MAX_RESULT_BYTES = 16 * 1024; +const MAX_RESULT_PREFIX_BYTES = 8 * 1024; +export const TRUNCATION_MARKER = + "[context-mode routing guard: output truncated; use ctx_execute or ctx_execute_file for analysis]"; + +const CONTEXT_MODE_TOOL_PATTERN = /(?:ctx_|context-mode|context_mode_ctx_)/i; +const WEB_COMMAND_PATTERN = + /\b(?:curl|wget|fetch|invoke-webrequest)(?:\.exe)?\b|\b(?:requests\.(?:get|post)|http\.(?:get|request)|urllib\.request)\b/i; +const ANALYSIS_COMMAND_PATTERNS = [ + /\b(?:cat|type|get-content|more|less|head|tail|grep|rg|find|findstr)\b/i, + /\bgit\s+(?:log|diff|show)\b/i, + /\b(?:npm|bun)\s+(?:test|run|build|lint)\b/i, + /\bpytest\b/i, + /\bcargo\s+test\b/i, + /\bgo\s+test\b/i, + /\b(?:gh|aws|docker|kubectl)\b/i, +]; +const CLEAR_MUTATION_PATTERN = + /^(?:(?:set\s+\w+=\S+|env\s+\w+=\S+)\s+)*(?:mkdir|md|mv|move-item|cp|copy-item|rm|del|remove-item|touch|ni|new-item|chmod|set-content|add-content|git\s+(?:add|commit|push|checkout|branch|merge)|npm\s+(?:install|publish)|pip\s+install|bun\s+(?:add|install)|echo|printf)\b/i; +const BROAD_PATHS = new Set([".", "./", "*", "**", "**/*"]); +const READ_SELECTOR_PATTERN = /(?::(?:raw|conflicts|\d+(?:-\d+|\+\d+)))+$/i; +const SPECIAL_SCHEME_PATTERN = /^(?:skill|local|artifact|memory|agent|history|issue|pr):\/\//i; + +type RecordLike = Record; + +export interface OmpRoutingToolCallEvent { + toolName?: unknown; + tool_name?: unknown; + name?: unknown; + input?: unknown; +} + +export interface OmpRoutingToolResultEvent { + toolName?: unknown; + tool_name?: unknown; + name?: unknown; + content?: unknown; + isError?: unknown; +} + +export interface OmpRoutingContext { + cwd?: string; +} + +export interface ToolCallBlock { + block: true; + reason: string; +} + +export interface ToolResultReplacement { + content: Array<{ type: "text"; text: string }>; + isError?: unknown; +} + +function asRecord(value: unknown): RecordLike | null { + if (value === null || typeof value !== "object" || Array.isArray(value)) return null; + return value as RecordLike; +} + +function toolNameOf(event: unknown): string { + const record = asRecord(event); + const value = record?.toolName ?? record?.tool_name ?? record?.name; + return typeof value === "string" ? value : ""; +} + +function inputOf(event: unknown): RecordLike { + const record = asRecord(event); + if (record && asRecord(record.input)) return record.input as RecordLike; + if (record && typeof record.input === "string") return { command: record.input }; + return {}; +} + +function isContextModeTool(toolName: string): boolean { + return CONTEXT_MODE_TOOL_PATTERN.test(toolName); +} + +function commandReason(command: string): string { + if (WEB_COMMAND_PATTERN.test(command)) { + return "Use context-mode MCP tools (ctx_execute, ctx_fetch_and_index) for command analysis and web/API access."; + } + return "Use context-mode MCP tools (ctx_execute) for command analysis instead of direct shell output."; +} + +function isClearlyPermittedMutation(command: string): boolean { + const trimmed = command.trim(); + if (trimmed.includes("|")) return false; + return CLEAR_MUTATION_PATTERN.test(trimmed); +} + +function classifyShellCall(input: RecordLike): ToolCallBlock | undefined { + const command = input.command; + if (typeof command !== "string" || command.trim() === "") return undefined; + + if (WEB_COMMAND_PATTERN.test(command)) { + return { block: true, reason: commandReason(command) }; + } + if (ANALYSIS_COMMAND_PATTERNS.some((pattern) => pattern.test(command))) { + return { block: true, reason: commandReason(command) }; + } + + const hasPipelineOrRedirect = /[|>]/.test(command) || /2>&1/.test(command); + if (hasPipelineOrRedirect && !isClearlyPermittedMutation(command)) { + return { + block: true, + reason: "Use context-mode MCP tools (ctx_execute) for command analysis; direct pipelines/redirections are reserved for clearly scoped mutations.", + }; + } + return undefined; +} + +function inputPath(input: RecordLike): string | undefined { + const pathValue = input.path ?? input.file_path; + return typeof pathValue === "string" ? pathValue : undefined; +} + +function hasReadSelector(pathValue: string): boolean { + return READ_SELECTOR_PATTERN.test(pathValue); +} + +function classifyReadCall(input: RecordLike, ctx: OmpRoutingContext): ToolCallBlock | undefined { + const pathValue = inputPath(input); + if (!pathValue || pathValue.trim() === "") return undefined; + + if (/^https?:\/\//i.test(pathValue)) { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_fetch_and_index) for web/API URLs instead of direct read.", + }; + } + if (SPECIAL_SCHEME_PATTERN.test(pathValue) || hasReadSelector(pathValue)) return undefined; + + try { + const cwd = typeof ctx.cwd === "string" && ctx.cwd ? ctx.cwd : process.cwd(); + const stats = statSync(resolve(cwd, pathValue)); + if (stats.isFile() && stats.size > MAX_READ_FILE_BYTES) { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_execute_file) for large local-file analysis or add a line selector to direct read.", + }; + } + } catch { + // Preserve the native OMP error when stat cannot classify the path. + } + return undefined; +} + +function isBroadPath(pathValue: string): boolean { + const normalized = pathValue.trim().toLowerCase().replaceAll("\\", "/"); + return BROAD_PATHS.has(normalized); +} + +function parseLimit(value: unknown): number | null | undefined { + if (typeof value === "number") return Number.isFinite(value) ? value : null; + if (typeof value === "string" && value.trim() !== "") { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : null; + } + return undefined; +} + +function classifyGrepCall(input: RecordLike): ToolCallBlock | undefined { + const pathValue = input.path; + if (typeof pathValue !== "string" || pathValue.trim() === "" || isBroadPath(pathValue)) { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_execute) for broad search/analysis; direct grep is reserved for a scoped file or directory.", + }; + } + return undefined; +} + +function classifyGlobCall(input: RecordLike): ToolCallBlock | undefined { + const pathValue = input.path; + if (typeof pathValue !== "string" || pathValue.trim() === "") { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_execute) for broad inventory/analysis; direct glob/list needs a scoped pattern or a limit of 50.", + }; + } + + const limit = parseLimit(input.limit); + if (limit !== undefined && (limit === null || limit > 50)) { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_execute) for broad inventory/analysis; direct glob/list is limited to 50 results.", + }; + } + if (isBroadPath(pathValue) && !(typeof limit === "number" && limit <= 50)) { + return { + block: true, + reason: "Use context-mode MCP tool (ctx_execute) for broad inventory/analysis; direct glob/list needs an explicit limit of 50 or less.", + }; + } + return undefined; +} + +/** Classify a tool call; return undefined for permitted or unknown inputs. */ +export function classifyOmpToolCall( + event: OmpRoutingToolCallEvent, + ctx: OmpRoutingContext = {}, +): ToolCallBlock | undefined { + try { + const toolName = toolNameOf(event); + if (!toolName || isContextModeTool(toolName)) return undefined; + + const normalizedName = toolName.toLowerCase(); + const input = inputOf(event); + if (["bash", "shell", "exec_command"].includes(normalizedName)) return classifyShellCall(input); + if (["read", "view"].includes(normalizedName)) return classifyReadCall(input, ctx); + if (normalizedName === "grep") return classifyGrepCall(input); + if (["glob", "list"].includes(normalizedName)) return classifyGlobCall(input); + } catch { + // Routing is best-effort; preserve native OMP behavior on malformed inputs. + } + return undefined; +} + +function utf8Prefix(value: string, maxBytes: number): string { + let byteLength = 0; + let end = 0; + for (let index = 0; index < value.length;) { + const codePoint = value.codePointAt(index)!; + const character = String.fromCodePoint(codePoint); + const characterBytes = Buffer.byteLength(character, "utf8"); + if (byteLength + characterBytes > maxBytes) break; + byteLength += characterBytes; + index += character.length; + end = index; + } + return value.slice(0, end); +} + +function utf8Suffix(value: string, maxBytes: number): string { + let byteLength = 0; + let start = value.length; + for (let index = value.length; index > 0;) { + let characterStart = index - 1; + const lastUnit = value.charCodeAt(index - 1); + if (lastUnit >= 0xdc00 && lastUnit <= 0xdfff && index >= 2) { + const precedingUnit = value.charCodeAt(index - 2); + if (precedingUnit >= 0xd800 && precedingUnit <= 0xdbff) characterStart = index - 2; + } + const character = value.slice(characterStart, index); + const characterBytes = Buffer.byteLength(character, "utf8"); + if (byteLength + characterBytes > maxBytes) break; + byteLength += characterBytes; + index = characterStart; + start = index; + } + return value.slice(start); +} + +/** Replace oversized non-context-mode tool text with a bounded head/tail view. */ +export function limitOmpToolResult(event: OmpRoutingToolResultEvent): ToolResultReplacement | undefined { + try { + if (isContextModeTool(toolNameOf(event))) return undefined; + + const record = asRecord(event); + const content = Array.isArray(record?.content) ? record.content : []; + const textParts = content + .map(asRecord) + .filter((part): part is RecordLike => part?.type === "text" && typeof part.text === "string") + .map((part) => part.text as string); + if (textParts.length === 0) return undefined; + + const text = textParts.join("\n"); + if (Buffer.byteLength(text, "utf8") <= MAX_RESULT_BYTES) return undefined; + + const markerBytes = Buffer.byteLength(TRUNCATION_MARKER, "utf8"); + const prefix = utf8Prefix(text, MAX_RESULT_PREFIX_BYTES); + const suffix = utf8Suffix( + text, + Math.max(0, MAX_RESULT_BYTES - markerBytes - Buffer.byteLength(prefix, "utf8")), + ); + const result: ToolResultReplacement = { + content: [{ type: "text", text: `${prefix}${TRUNCATION_MARKER}${suffix}` }], + }; + if (record && Object.prototype.hasOwnProperty.call(record, "isError")) { + result.isError = record.isError; + } + return result; + } catch { + // Result limiting is a fallback; never break a tool result on guard failure. + return undefined; + } +} \ No newline at end of file diff --git a/tests/adapters/omp-plugin.test.ts b/tests/adapters/omp-plugin.test.ts index 094490f62..b28a799c4 100644 --- a/tests/adapters/omp-plugin.test.ts +++ b/tests/adapters/omp-plugin.test.ts @@ -20,11 +20,11 @@ import "../setup-home"; */ import { describe, it, expect, beforeEach, afterEach } from "vitest"; -import { mkdtempSync, rmSync } from "node:fs"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; import { SessionDB } from "../../src/session/db.js"; - +import { MAX_RESULT_BYTES, TRUNCATION_MARKER } from "../../src/adapters/omp/routing-guard.js"; // ── Mock OMP HookAPI ──────────────────────────────────────── type HandlerFn = (...args: unknown[]) => unknown | Promise; @@ -171,6 +171,79 @@ describe("OMP plugin", () => { await expect(api._trigger("tool_call", {})).resolves.toBeUndefined(); await expect(api._trigger("tool_call", { toolName: "bash" })).resolves.toBeUndefined(); }); + it("blocks npm test and points to ctx_execute", async () => { + await registerOmpPlugin(api); + const result = (await api._trigger("tool_call", { + toolName: "bash", + input: { command: "npm test" }, + })) as { block?: boolean; reason?: string } | undefined; + + expect(result?.block).toBe(true); + expect(result?.reason).toMatch(/ctx_execute/); + }); + + it("blocks broad grep but allows a scoped grep", async () => { + await registerOmpPlugin(api); + const broad = await api._trigger("tool_call", { + toolName: "grep", + input: { pattern: "TODO", path: "." }, + }); + const scoped = await api._trigger("tool_call", { + toolName: "grep", + input: { pattern: "TODO", path: "src/file.ts" }, + }); + + expect((broad as { block?: boolean } | undefined)?.block).toBe(true); + expect(scoped).toBeUndefined(); + }); + + it("blocks a large local read but allows a line-selected read", async () => { + writeFileSync(join(tempDir, "large.txt"), "x".repeat(32 * 1024 + 1)); + await registerOmpPlugin(api); + + const large = await api._trigger("tool_call", { + toolName: "read", + input: { path: "large.txt" }, + }); + const selected = await api._trigger("tool_call", { + toolName: "read", + input: { path: "large.txt:10-30" }, + }); + + expect((large as { block?: boolean } | undefined)?.block).toBe(true); + expect(selected).toBeUndefined(); + }); + + it("allows scoped glob and context-mode tools", async () => { + await registerOmpPlugin(api); + const scopedGlob = await api._trigger("tool_call", { + toolName: "glob", + input: { path: "src/**/*.ts", limit: 50 }, + }); + const broadGlob = await api._trigger("tool_call", { + toolName: "glob", + input: { path: "**/*" }, + }); + const contextMode = await api._trigger("tool_call", { + toolName: "mcp__context_mode_ctx_execute", + input: { language: "javascript", code: "console.log(1)" }, + }); + + expect(scopedGlob).toBeUndefined(); + expect((broadGlob as { block?: boolean } | undefined)?.block).toBe(true); + expect(contextMode).toBeUndefined(); + }); + + it("blocks a direct read URL and points to ctx_fetch_and_index", async () => { + await registerOmpPlugin(api); + const result = (await api._trigger("tool_call", { + toolName: "read", + input: { path: "https://example.com/docs" }, + })) as { block?: boolean; reason?: string } | undefined; + + expect(result?.block).toBe(true); + expect(result?.reason).toMatch(/ctx_fetch_and_index/); + }); }); // ═══════════════════════════════════════════════════════════ @@ -225,6 +298,34 @@ describe("OMP plugin", () => { }), ).resolves.toBeUndefined(); }); + it("bounds oversized direct results and preserves isError", async () => { + await registerOmpPlugin(api); + const result = (await api._trigger("tool_result", { + toolName: "read", + content: [{ type: "text", text: "x".repeat(MAX_RESULT_BYTES + 100) }], + isError: true, + })) as { content?: Array<{ type?: string; text?: string }>; isError?: boolean } | undefined; + const text = result?.content?.[0]?.text ?? ""; + + expect(text).toContain(TRUNCATION_MARKER); + expect(Buffer.byteLength(text, "utf8")).toBeLessThanOrEqual(MAX_RESULT_BYTES); + expect(result?.isError).toBe(true); + }); + + it("does not alter small or context-mode results", async () => { + await registerOmpPlugin(api); + const small = await api._trigger("tool_result", { + toolName: "read", + content: [{ type: "text", text: "small" }], + }); + const contextMode = await api._trigger("tool_result", { + toolName: "mcp__context_mode_ctx_execute", + content: [{ type: "text", text: "x".repeat(MAX_RESULT_BYTES + 100) }], + }); + + expect(small).toBeUndefined(); + expect(contextMode).toBeUndefined(); + }); }); // ═══════════════════════════════════════════════════════════