From e4b4d13ac750a140bb7166b630bf8904ab290c07 Mon Sep 17 00:00:00 2001
From: Greg Harvell <greg.harvell@complexthings.com>
Date: Fri, 27 Mar 2026 18:50:56 -0400
Subject: [PATCH] fix: prevent orphaned tool_use blocks from compression and
 harden autocomplete

Bug 1 (400 API error): applyCompressionBlocks could remove toolResult
messages while leaving their paired assistant(tool_use) message intact.
This produced invalid API sequences that Claude rejected with:
  messages.N: tool_use ids found without tool_result blocks immediately after

The fix adds two boundary expansions before the splice:
- Expand lo backward: if messages[lo-1] is an assistant whose toolCall
  ids appear as toolResult.toolCallId values inside [lo..hi], pull the
  assistant into the range
- Expand hi forward: for assistants inside [lo..hi], extend hi to include
  any consecutive toolResult messages that immediately follow hi

This ensures tool_use/tool_result pairs are always removed together.

Bug 2 (autocomplete crash): harden getArgumentCompletions to return an
explicitly typed AutocompleteItem[] | null, filter on value (not label,
matching pi-tui's internal contract), and add a typeof guard ensuring no
item with a non-string value can reach getBestAutocompleteMatchIndex.
---
 commands.ts | 23 +++++++++++++----------
 pruner.ts   | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/commands.ts b/commands.ts
index 7034599..93c55a3 100644
--- a/commands.ts
+++ b/commands.ts
@@ -1,4 +1,5 @@
 import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent"
+import type { AutocompleteItem } from "@mariozechner/pi-tui"
 import type { DcpState } from "./state.js"
 import type { DcpConfig } from "./config.js"
 
@@ -289,17 +290,19 @@ export function registerCommands(
 ): void {
   pi.registerCommand("dcp", {
     description: "Dynamic Context Pruning — manage context window usage",
-    getArgumentCompletions(prefix: string) {
-      const subcommands = [
-        { label: "context", description: "Show context window usage breakdown" },
-        { label: "stats", description: "Show pruning statistics" },
-        { label: "sweep", description: "Prune tool outputs" },
-        { label: "manual", description: "Toggle manual mode" },
-        { label: "decompress", description: "List or restore compression blocks" },
-        { label: "compress", description: "Trigger LLM compression" },
-        { label: "help", description: "Show help" },
+    getArgumentCompletions(prefix: string): AutocompleteItem[] | null {
+      const subcommands: AutocompleteItem[] = [
+        { value: "context", label: "context", description: "Show context window usage breakdown" },
+        { value: "stats", label: "stats", description: "Show pruning statistics" },
+        { value: "sweep", label: "sweep", description: "Prune tool outputs" },
+        { value: "manual", label: "manual", description: "Toggle manual mode" },
+        { value: "decompress", label: "decompress", description: "List or restore compression blocks" },
+        { value: "compress", label: "compress", description: "Trigger LLM compression" },
+        { value: "help", label: "help", description: "Show help" },
       ]
-      const matched = subcommands.filter((s) => s.label.startsWith(prefix))
+      const matched = subcommands
+        .filter((s) => typeof s.value === "string")
+        .filter((s) => s.value.startsWith(prefix))
       return matched.length > 0 ? matched : null
     },
 
diff --git a/pruner.ts b/pruner.ts
index 31b2cad..285b4f3 100644
--- a/pruner.ts
+++ b/pruner.ts
@@ -54,8 +54,57 @@ function applyCompressionBlocks(messages: any[], state: DcpState): any[] {
 
     if (startIdx === -1 || endIdx === -1) continue;
 
-    const lo = Math.min(startIdx, endIdx);
-    const hi = Math.max(startIdx, endIdx);
+    let lo = Math.min(startIdx, endIdx);
+    let hi = Math.max(startIdx, endIdx);
+
+    // Expand lo backward: if the message immediately before lo is an assistant
+    // whose tool_use blocks have matching tool_results inside [lo..hi], pull
+    // it into the range so the pair is always removed together.
+    while (lo > 0) {
+      const prev = messages[lo - 1] as any;
+      if (prev.role !== "assistant") break;
+      const toolCallIdsInRange = new Set<string>();
+      for (let i = lo; i <= hi; i++) {
+        const m = messages[i] as any;
+        if (m.role === "toolResult" && typeof m.toolCallId === "string") {
+          toolCallIdsInRange.add(m.toolCallId);
+        }
+      }
+      const prevContent: any[] = Array.isArray(prev.content) ? prev.content : [];
+      const hasMatchingToolCalls = prevContent.some(
+        (block: any) => block.type === "toolCall" && toolCallIdsInRange.has(block.id)
+      );
+      if (!hasMatchingToolCalls) break;
+      lo--;
+    }
+
+    // Expand hi forward: for every assistant message in [lo..hi] that has
+    // tool_use blocks, include any immediately-following tool_result messages
+    // that correspond to those blocks. Loop to fixed point because expanding
+    // hi could expose more assistants in theory.
+    let prevHi: number;
+    do {
+      prevHi = hi;
+      const assistantToolCallIds = new Set<string>();
+      for (let i = lo; i <= hi; i++) {
+        const m = messages[i] as any;
+        if (m.role !== "assistant") continue;
+        const content: any[] = Array.isArray(m.content) ? m.content : [];
+        for (const block of content) {
+          if (block.type === "toolCall" && typeof block.id === "string") {
+            assistantToolCallIds.add(block.id);
+          }
+        }
+      }
+      while (hi + 1 < messages.length) {
+        const next = messages[hi + 1] as any;
+        if (next.role === "toolResult" && assistantToolCallIds.has(next.toolCallId)) {
+          hi++;
+        } else {
+          break;
+        }
+      }
+    } while (hi !== prevHi);
 
     // Estimate tokens removed
     let removedTokens = 0;