pi-tool-guard 0.1.0 → 0.2.0

package/README.md

@@ -33,20 +33,19 @@ When the LLM calls a tool with wrong field names, the extension normalizes them
 
 ### 2. Bash pipeline extractor stripping
 
-When the LLM appends truncation commands (`tail`, `head`, `grep`, etc.) to slow commands, the extension strips them so the full output is available. This prevents the LLM from repeatedly re-running slow commands with different truncation parameters.
+When the LLM appends truncation commands (`tail`, `head`, `grep`, etc.) to bash commands, the extension strips them and applies the extractor intelligently.
 
 **Three-case strategy:**
 
-| Scenario | Behavior |
-|---|---|
-| Output truncated (has `Full output: <file>`) | Run extractor on the full output file, return filtered result |
-| Fast command (< 10s), no truncation | Pipe result through extractor, return filtered result |
-| Slow command, no truncation | Return full result with notice |
+| Scenario | UI Notification | LLM Response |
+|---|---|---|
+| Fast command (< 10s), truncated | `Filtered via \`head\`` | Filtered result only |
+| Fast command (< 10s), not truncated | `Filtered via \`head\`` | Filtered result only |
+| Slow command (truncated or not) | `The full output is above...` | Result + `This is a slow command. Avoid re-running...` |
 
 **Example:** `vitest run | tail -n 10`
-- If vitest output is truncated → run `tail` on the saved full output file
-- If vitest finishes in < 10s → pipe result through `tail`
-- If vitest is slow but not truncated → return full output + notice
+- If vitest finishes in < 10s → run `tail` on result (or full output file if truncated), notify UI
+- If vitest is slow → return result as-is, notify UI, append slow-command hint to LLM
 
 **Detected extractors:** `head`, `tail`, `grep`, `egrep`, `fgrep`, `rg`, `sed`, `awk`, `cut`, `sort`, `uniq`, `wc`, `less`, `more`, `column`, `jq`, `yq`, `tr`
 
@@ -62,8 +61,20 @@ Both features use `prepareArguments` on overridden built-in tools — the cleane
 - **bash**: `createBashToolDefinition(cwd)` + custom `execute` override:
   1. `prepareArguments` parses the command with [unbash](https://github.com/nicolo-ribaudo/unbash), strips trailing extractors
   2. `execute` runs the stripped command via the original built-in execute
-  3. If truncated → runs extractor on the full output file via `pi.exec`
-  4. If fast (< 10s) → pipes result through extractor via `pi.exec`
-  5. If slow → returns full result with notice
+  3. If fast (< 10s) + truncated → runs extractor on the full output file via `pi.exec`
+  4. If fast (< 10s) + not truncated → pipes result through extractor via `pi.exec`
+  5. If slow (truncated or not) → returns result as-is with notice
 
 No error recovery, no session scanning.
+
+## Development
+
+```bash
+pnpm install
+pnpm test        # Run all tests
+pnpm typecheck   # Type check
+```
+
+## License
+
+MIT

package/bash.test.ts

@@ -0,0 +1,437 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+import { buildNotice, FAST_THRESHOLD_MS, runExtractorOnText, runExtractorOnFile } from "./bash.ts";
+import { stripTrailingExtractors } from "./pipeline.ts";
+import { executeBashGuarded, prepareBashArguments, type ExecuteContext, type ToolResult } from "./execute.ts";
+
+describe("buildNotice", () => {
+  it("returns fast-mode notice when mode is 'fast'", () => {
+    const notice = buildNotice(["head"], "fast");
+    expect(notice).toContain("Filtered via");
+    expect(notice).toContain("small enough to pipe");
+    expect(notice).not.toContain("do NOT re-run");
+  });
+
+  it("returns slow-mode notice when mode is 'slow'", () => {
+    const notice = buildNotice(["head"], "slow");
+    expect(notice).toContain("The full output is above");
+    expect(notice).toContain("do NOT re-run");
+  });
+
+  it("defaults to slow mode when no mode provided", () => {
+    const notice = buildNotice(["head"]);
+    expect(notice).toContain("The full output is above");
+    expect(notice).toContain("do NOT re-run");
+  });
+
+  it("lists multiple removed commands", () => {
+    const notice = buildNotice(["grep", "head"], "fast");
+    expect(notice).toContain("`grep`, `head`");
+  });
+});
+
+describe("FAST_THRESHOLD_MS", () => {
+  it("is 10 seconds", () => {
+    expect(FAST_THRESHOLD_MS).toBe(10_000);
+  });
+});
+
+describe("fast command behavior", () => {
+  let mockPi: ExtensionAPI;
+
+  beforeEach(() => {
+    mockPi = {
+      exec: vi.fn(),
+    } as unknown as ExtensionAPI;
+  });
+
+  it("runExtractorOnText pipes text through extractor pipeline", async () => {
+    const text = "line1\nline2\nline3";
+    const pipeline = "head -2";
+    
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: "line1\nline2",
+    });
+
+    const result = await runExtractorOnText(mockPi, text, pipeline);
+    expect(result).toBe("line1\nline2");
+    expect(mockPi.exec).toHaveBeenCalledWith("bash", [
+      "-c",
+      expect.stringContaining("head -2"),
+    ]);
+  });
+
+  it("runExtractorOnFile pipes file content through extractor pipeline", async () => {
+    const filePath = "/tmp/output.txt";
+    const pipeline = "tail -1";
+    
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: "last line",
+    });
+
+    const result = await runExtractorOnFile(mockPi, filePath, pipeline);
+    expect(result).toBe("last line");
+    expect(mockPi.exec).toHaveBeenCalledWith("bash", [
+      "-c",
+      expect.stringContaining("tail -1"),
+    ]);
+  });
+
+  it("runExtractorOnText returns undefined on failure", async () => {
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 1,
+      stdout: "",
+    });
+
+    const result = await runExtractorOnText(mockPi, "text", "head -1");
+    expect(result).toBeUndefined();
+  });
+
+  it("runExtractorOnFile returns undefined on failure", async () => {
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockRejectedValue(new Error("fail"));
+
+    const result = await runExtractorOnFile(mockPi, "/tmp/file.txt", "head -1");
+    expect(result).toBeUndefined();
+  });
+});
+
+describe("stripTrailingExtractors integration", () => {
+  it("strips head from ls | head", () => {
+    const result = stripTrailingExtractors("ls -la | head -5");
+    expect(result).toBeDefined();
+    expect(result!.cleaned).toBe("ls -la");
+    expect(result!.removedNames).toEqual(["head"]);
+    expect(result!.removedPipeline).toBe("head -5");
+  });
+
+  it("strips tail from npm test | tail -20", () => {
+    const result = stripTrailingExtractors("npm test 2>&1 | tail -20");
+    expect(result).toBeDefined();
+    expect(result!.cleaned).toBe("npm test 2>&1");
+    expect(result!.removedNames).toEqual(["tail"]);
+  });
+
+  it("strips multiple extractors from grep FAIL | head -5", () => {
+    const result = stripTrailingExtractors("npm test | grep FAIL | head -5");
+    expect(result).toBeDefined();
+    expect(result!.removedNames).toEqual(["grep", "head"]);
+  });
+
+  it("returns undefined for commands without extractors", () => {
+    const result = stripTrailingExtractors("ls -la");
+    expect(result).toBeUndefined();
+  });
+});
+
+describe("executeBashGuarded", () => {
+  let mockPi: ExtensionAPI;
+  let mockCtx: { ui: { notify: ReturnType<typeof vi.fn> } };
+  let mockOriginalExecute: ReturnType<typeof vi.fn>;
+  let ctx: ExecuteContext;
+
+  beforeEach(() => {
+    mockPi = {
+      exec: vi.fn(),
+    } as unknown as ExtensionAPI;
+    mockCtx = {
+      ui: { notify: vi.fn() },
+    };
+    mockOriginalExecute = vi.fn().mockResolvedValue({
+      content: [{ type: "text", text: "" }],
+      details: undefined,
+    });
+    ctx = {
+      pi: mockPi,
+      ctx: mockCtx as unknown as ExecuteContext["ctx"],
+      originalExecute: mockOriginalExecute as unknown as ExecuteContext["originalExecute"],
+    };
+  });
+
+  it("runs original execute when no extractors removed", async () => {
+    const expected: ToolResult = {
+      content: [{ type: "text", text: "output" }],
+      details: undefined,
+    };
+    mockOriginalExecute.mockResolvedValue(expected);
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      { command: "ls -la" },
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    expect(result).toBe(expected);
+    expect(mockOriginalExecute).toHaveBeenCalledOnce();
+  });
+
+  it("fast command: pipes result through extractor and returns filtered output", async () => {
+    // Simulate fast command (< 10s)
+    const originalResult: ToolResult = {
+      content: [{ type: "text", text: "line1\nline2\nline3" }],
+      details: undefined,
+    };
+    mockOriginalExecute.mockImplementation(async () => {
+      // Simulate fast execution
+      return originalResult;
+    });
+
+    // Mock pi.exec to simulate head -2
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: "line1\nline2",
+    });
+
+    const params = { command: "ls -la" };
+    // prepareBashArguments would have set these
+    const prepared = prepareBashArguments({ command: "ls -la | head -2" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // Should return filtered output (head -2)
+    const textContent = result.content[0];
+    if (textContent.type === "text") {
+      expect(textContent.text).toBe("line1\nline2");
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("Filtered via"),
+      "info",
+    );
+    // Should NOT append notice to LLM response for fast commands
+    expect(result.content).toHaveLength(1);
+  });
+
+  it("fast command: produces same output as original bash with extractor", async () => {
+    // This is the key test: ensure guarded execution == original execution
+    const originalOutput = "file1.txt\nfile2.txt\nfile3.txt";
+    const headOutput = "file1.txt\nfile2.txt";
+
+    mockOriginalExecute.mockResolvedValue({
+      content: [{ type: "text", text: originalOutput }],
+      details: undefined,
+    });
+
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: headOutput,
+    });
+
+    const prepared = prepareBashArguments({ command: "ls | head -2" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // The filtered output should match what head -2 would produce
+    const textContent = result.content[0];
+    if (textContent.type === "text") {
+      expect(textContent.text).toBe(headOutput);
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalled();
+    // Should NOT append notice to LLM response for fast commands
+    expect(result.content).toHaveLength(1);
+  });
+
+  it("slow command: returns full output with LLM notice", async () => {
+    // Use fake timers to simulate slow command
+    vi.useFakeTimers();
+
+    mockOriginalExecute.mockImplementation(async () => {
+      // Advance time past the threshold
+      vi.advanceTimersByTime(FAST_THRESHOLD_MS + 100);
+      return {
+        content: [{ type: "text", text: "full output" }],
+        details: undefined,
+      };
+    });
+
+    const prepared = prepareBashArguments({ command: "npm test | tail -5" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // Should return full output with LLM notice
+    const textContent = result.content[0];
+    const llmNotice = result.content[1];
+    if (textContent.type === "text" && llmNotice.type === "text") {
+      expect(textContent.text).toBe("full output");
+      expect(llmNotice.text).toContain("slow command");
+      expect(llmNotice.text).toContain("Avoid re-running");
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("The full output is above"),
+      "info",
+    );
+
+    vi.useRealTimers();
+  });
+
+  it("slow command + truncated: returns truncated output with LLM notice, no extractor run", async () => {
+    // Use fake timers to simulate slow command
+    vi.useFakeTimers();
+
+    mockOriginalExecute.mockImplementation(async () => {
+      vi.advanceTimersByTime(FAST_THRESHOLD_MS + 100);
+      return {
+        content: [{ type: "text", text: "truncated output..." }],
+        details: { fullOutputPath: "/tmp/full-output.txt" },
+      };
+    });
+
+    const prepared = prepareBashArguments({ command: "npm test | tail -5" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // Should NOT run extractor on file for slow commands
+    expect(mockPi.exec).not.toHaveBeenCalled();
+    // Should return truncated result as-is with LLM notice
+    const textContent = result.content[0];
+    const llmNotice = result.content[1];
+    if (textContent.type === "text" && llmNotice.type === "text") {
+      expect(textContent.text).toBe("truncated output...");
+      expect(llmNotice.text).toContain("slow command");
+      expect(llmNotice.text).toContain("Avoid re-running");
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("The full output is above"),
+      "info",
+    );
+
+    vi.useRealTimers();
+  });
+
+  it("truncated output: fast command runs extractor on full output file", async () => {
+    mockOriginalExecute.mockResolvedValue({
+      content: [{ type: "text", text: "truncated..." }],
+      details: { fullOutputPath: "/tmp/full-output.txt" },
+    });
+
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: "filtered from file",
+    });
+
+    const prepared = prepareBashArguments({ command: "npm test | tail -5" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // Should run extractor on the full output file
+    expect(mockPi.exec).toHaveBeenCalledWith("bash", [
+      "-c",
+      expect.stringContaining("/tmp/full-output.txt"),
+    ]);
+    // Should return filtered output only (no LLM notice for fast commands)
+    const textContent = result.content[0];
+    if (textContent.type === "text") {
+      expect(textContent.text).toBe("filtered from file");
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("Filtered via"),
+      "info",
+    );
+    // Should NOT append notice to LLM response for fast commands
+    expect(result.content).toHaveLength(1);
+  });
+
+  it("strips multiple extractors: grep FAIL | head -5", async () => {
+    mockOriginalExecute.mockResolvedValue({
+      content: [{ type: "text", text: "FAIL test1\nFAIL test2\nPASS test3" }],
+      details: undefined,
+    });
+
+    (mockPi.exec as ReturnType<typeof vi.fn>).mockResolvedValue({
+      code: 0,
+      stdout: "FAIL test1\nFAIL test2",
+    });
+
+    const prepared = prepareBashArguments({ command: "npm test | grep FAIL | head -5" });
+
+    const result = await executeBashGuarded(
+      ctx,
+      "call-1",
+      prepared,
+      new AbortController().signal,
+      () => {},
+      {},
+    );
+
+    // Should pipe through both grep and head
+    expect(mockPi.exec).toHaveBeenCalledWith("bash", [
+      "-c",
+      expect.stringContaining("grep FAIL | head -5"),
+    ]);
+    const textContent = result.content[0];
+    if (textContent.type === "text") {
+      expect(textContent.text).toBe("FAIL test1\nFAIL test2");
+    }
+    // Should notify UI
+    expect(mockCtx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("Filtered via"),
+      "info",
+    );
+    // Should NOT append notice to LLM response for fast commands
+    expect(result.content).toHaveLength(1);
+  });
+});
+
+describe("prepareBashArguments", () => {
+  it("strips trailing extractor and stores metadata", () => {
+    const result = prepareBashArguments({ command: "ls | head -5" });
+    expect(result.command).toBe("ls");
+    expect(result._piToolGuardRemoved).toEqual(["head"]);
+    expect(result._piToolGuardPipeline).toBe("head -5");
+  });
+
+  it("does nothing for commands without extractors", () => {
+    const result = prepareBashArguments({ command: "ls -la" });
+    expect(result.command).toBe("ls -la");
+    expect(result._piToolGuardRemoved).toBeUndefined();
+  });
+
+  it("handles non-string command gracefully", () => {
+    const result = prepareBashArguments({ command: 123 } as Record<string, unknown>);
+    expect(result.command).toBe(123);
+  });
+});

package/bash.ts

@@ -6,9 +6,14 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
 export const FAST_THRESHOLD_MS = 10_000;
 
-export function buildNotice(removedNames: string[]): string {
+export type NoticeMode = "fast" | "slow";
+
+export function buildNotice(removedNames: string[], mode: NoticeMode = "slow"): string {
   const list = removedNames.map((r) => `\`${r}\``).join(", ");
-  return `[pi-tool-guard] Removed trailing pipeline commands: ${list}. The full output is above — do NOT re-run with different parameters.`;
+  const suffix = mode === "fast"
+    ? `Filtered via ${list} — full output was small enough to pipe.`
+    : `The full output is above — do NOT re-run with different parameters.`;
+  return `[pi-tool-guard] Removed trailing pipeline commands: ${list}. ${suffix}`;
 }
 
 /**

package/execute.ts

@@ -0,0 +1,136 @@
+import type { ExtensionAPI, ExtensionContext, AgentToolResult, BashToolDetails } from "@earendil-works/pi-coding-agent";
+import {
+  stripTrailingExtractors,
+  extractFullOutputPath,
+} from "./pipeline.ts";
+import {
+  FAST_THRESHOLD_MS,
+  buildNotice,
+  runExtractorOnFile,
+  runExtractorOnText,
+} from "./bash.ts";
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+export type ToolResult = AgentToolResult<BashToolDetails | undefined>;
+
+export interface ExecuteContext {
+  pi: ExtensionAPI;
+  ctx: ExtensionContext;
+  originalExecute: (
+    toolCallId: string,
+    params: unknown,
+    signal: AbortSignal | undefined,
+    onUpdate: ((update: unknown) => void) | undefined,
+    execCtx: unknown,
+  ) => Promise<ToolResult>;
+}
+
+// ---------------------------------------------------------------------------
+// Core execute logic (extracted for testability)
+// ---------------------------------------------------------------------------
+
+export async function executeBashGuarded(
+  ctx: ExecuteContext,
+  toolCallId: string,
+  params: unknown,
+  signal: AbortSignal | undefined,
+  onUpdate: ((update: unknown) => void) | undefined,
+  execCtx: unknown,
+): Promise<ToolResult> {
+  const input = params as Record<string, unknown>;
+  const removedNames = input._piToolGuardRemoved as string[] | undefined;
+  const removedPipeline = input._piToolGuardPipeline as string | undefined;
+
+  // No extractors — run normally
+  if (!removedNames || removedNames.length === 0 || !removedPipeline) {
+    return ctx.originalExecute(toolCallId, params, signal, onUpdate, execCtx);
+  }
+
+  // Clean transient fields before running
+  delete input._piToolGuardRemoved;
+  delete input._piToolGuardPipeline;
+
+  // Run the stripped command
+  const start = Date.now();
+  let result;
+  try {
+    result = await ctx.originalExecute(toolCallId, params, signal, onUpdate, execCtx);
+  } catch (err) {
+    // On error: check if output was saved to file, run extractor on it
+    const errText = err instanceof Error ? err.message : String(err);
+    const fullPath = extractFullOutputPath(errText);
+    if (fullPath) {
+      const extracted = await runExtractorOnFile(ctx.pi, fullPath, removedPipeline);
+      if (extracted !== undefined) {
+        const uiNotice = buildNotice(removedNames, "fast");
+        ctx.ctx.ui.notify(uiNotice, "info");
+        throw new Error(extracted);
+      }
+    }
+    throw err;
+  }
+
+  const elapsed = Date.now() - start;
+  const details = result.details as BashToolDetails | undefined;
+  const isTruncated = !!details?.fullOutputPath;
+
+  // Case 1: Fast command, truncated → run extractor on the full output file
+  if (elapsed < FAST_THRESHOLD_MS && isTruncated) {
+    const extracted = await runExtractorOnFile(ctx.pi, details!.fullOutputPath!, removedPipeline);
+    if (extracted !== undefined) {
+      const uiNotice = buildNotice(removedNames, "fast");
+      ctx.ctx.ui.notify(uiNotice, "info");
+      return {
+        content: [{ type: "text" as const, text: extracted }],
+        details: result.details,
+      };
+    }
+  }
+
+  // Case 2: Fast command, not truncated → pipe result through extractor
+  if (elapsed < FAST_THRESHOLD_MS && !isTruncated) {
+    const resultText = result.content
+      .filter((c): c is { type: "text"; text: string } => c.type === "text")
+      .map((c) => c.text)
+      .join("\n");
+
+    const extracted = await runExtractorOnText(ctx.pi, resultText, removedPipeline);
+    if (extracted !== undefined) {
+      const uiNotice = buildNotice(removedNames, "fast");
+      ctx.ctx.ui.notify(uiNotice, "info");
+      return {
+        content: [{ type: "text" as const, text: extracted }],
+        details: result.details,
+      };
+    }
+  }
+
+  // Case 3: Slow command (truncated or not) → return result with LLM notice
+  const uiNotice = buildNotice(removedNames, "slow");
+  ctx.ctx.ui.notify(uiNotice, "info");
+  const llmNotice = `This is a slow command. Avoid re-running; prefer reading from the full output path if available.`;
+  return {
+    content: [...result.content, { type: "text" as const, text: `\n\n${llmNotice}` }],
+    details: result.details,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// prepareArguments logic (extracted for testability)
+// ---------------------------------------------------------------------------
+
+export function prepareBashArguments(input: Record<string, unknown>): Record<string, unknown> {
+  if (typeof input.command !== "string") return input;
+
+  const stripped = stripTrailingExtractors(input.command);
+  if (stripped) {
+    input.command = stripped.cleaned;
+    input._piToolGuardRemoved = stripped.removedNames;
+    input._piToolGuardPipeline = stripped.removedPipeline;
+  }
+
+  return input;
+}

package/index.ts

@@ -1,7 +1,6 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import {
   type BashToolInput,
-  type BashToolDetails,
   type EditToolInput,
   type ReadToolInput,
   type WriteToolInput,
@@ -15,16 +14,7 @@ import {
   normalizeWriteArgs,
   normalizeReadArgs,
 } from "./aliases.ts";
-import {
-  stripTrailingExtractors,
-  extractFullOutputPath,
-} from "./pipeline.ts";
-import {
-  FAST_THRESHOLD_MS,
-  buildNotice,
-  runExtractorOnFile,
-  runExtractorOnText,
-} from "./bash.ts";
+import { executeBashGuarded, prepareBashArguments } from "./execute.ts";
 
 // ---------------------------------------------------------------------------
 // extension
@@ -73,91 +63,18 @@ export default function toolGuardExtension(pi: ExtensionAPI) {
       ...bashDef,
       prepareArguments: (args: unknown): BashToolInput => {
         if (!args || typeof args !== "object") return args as BashToolInput;
-
-        const input = args as Record<string, unknown>;
-        if (typeof input.command !== "string") return args as BashToolInput;
-
-        const stripped = stripTrailingExtractors(input.command);
-        if (stripped) {
-          input.command = stripped.cleaned;
-          input._piToolGuardRemoved = stripped.removedNames;
-          input._piToolGuardPipeline = stripped.removedPipeline;
-        }
-
-        return input as unknown as BashToolInput;
+        return prepareBashArguments(args as Record<string, unknown>) as unknown as BashToolInput;
       },
 
       async execute(toolCallId, params, signal, onUpdate, execCtx) {
-        const input = params as Record<string, unknown>;
-        const removedNames = input._piToolGuardRemoved as string[] | undefined;
-        const removedPipeline = input._piToolGuardPipeline as string | undefined;
-
-        // No extractors — run normally
-        if (!removedNames || removedNames.length === 0 || !removedPipeline) {
-          return originalExecute(toolCallId, params, signal, onUpdate, execCtx);
-        }
-
-        // Clean transient fields before running
-        delete input._piToolGuardRemoved;
-        delete input._piToolGuardPipeline;
-
-        // Run the stripped command
-        const start = Date.now();
-        let result;
-        try {
-          result = await originalExecute(toolCallId, params, signal, onUpdate, execCtx);
-        } catch (err) {
-          // On error: check if output was saved to file, run extractor on it
-          const errText = err instanceof Error ? err.message : String(err);
-          const fullPath = extractFullOutputPath(errText);
-          if (fullPath) {
-            const extracted = await runExtractorOnFile(pi, fullPath, removedPipeline);
-            if (extracted !== undefined) {
-              const notice = buildNotice(removedNames);
-              throw new Error(`${extracted}\n\n${notice}`);
-            }
-          }
-          throw err;
-        }
-
-        const elapsed = Date.now() - start;
-
-        // Case 1: Output was truncated → run extractor on the full output file
-        const details = result.details as BashToolDetails | undefined;
-        if (details?.fullOutputPath) {
-          const extracted = await runExtractorOnFile(pi, details.fullOutputPath, removedPipeline);
-          if (extracted !== undefined) {
-            const notice = buildNotice(removedNames);
-            return {
-              content: [{ type: "text" as const, text: `${extracted}\n\n${notice}` }],
-              details: result.details,
-            };
-          }
-        }
-
-        // Case 2: Fast command, no truncation → pipe result through extractor
-        if (elapsed < FAST_THRESHOLD_MS) {
-          const resultText = result.content
-            .filter((c): c is { type: "text"; text: string } => c.type === "text")
-            .map((c) => c.text)
-            .join("\n");
-
-          const extracted = await runExtractorOnText(pi, resultText, removedPipeline);
-          if (extracted !== undefined) {
-            const notice = buildNotice(removedNames);
-            return {
-              content: [{ type: "text" as const, text: `${extracted}\n\n${notice}` }],
-              details: result.details,
-            };
-          }
-        }
-
-        // Case 3: Slow, no truncation → return full result with notice
-        const notice = buildNotice(removedNames);
-        return {
-          content: [...result.content, { type: "text" as const, text: `\n\n${notice}` }],
-          details: result.details,
-        };
+        return executeBashGuarded(
+          { pi, ctx, originalExecute },
+          toolCallId,
+          params,
+          signal,
+          onUpdate,
+          execCtx,
+        );
       },
     });
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-tool-guard",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "type": "module",
   "description": "Pi extension that corrects LLM tool calls: normalizes argument aliases for edit/write/read and strips trailing pipeline extractors from bash commands",
   "main": "index.ts",