pi-ui-extend 0.1.1

package/external/pi-tools-suite/src/compress/pruner-candidates.ts

@@ -0,0 +1,176 @@
+import type { DcpConfig } from "./config.js";
+import type { DcpState } from "./state.js";
+import type { CompressionCandidate, MessageCompressionCandidate } from "./pruner-types.js";
+import {
+  estimateMessageTokens,
+  extractBlockId,
+  extractMessageId,
+  messageText,
+} from "./pruner-metadata.js";
+
+interface CandidateBoundary {
+  id: string;
+  role: string;
+  timestamp: number;
+  tokenEstimate: number;
+  blockId?: number;
+  isSystemReminder: boolean;
+}
+
+export function detectCompressionCandidate(
+  messages: any[],
+  _state: DcpState,
+  config: DcpConfig,
+  contextPercent: number,
+): CompressionCandidate | null {
+  const settings = config.compress.autoCandidates;
+  if (!settings.enabled) return null;
+  if (contextPercent < settings.minContextPercent) return null;
+
+  const boundaries: CandidateBoundary[] = [];
+  for (const msg of messages) {
+    const text = messageText(msg);
+    const blockId = extractBlockId(text);
+    const messageId = extractMessageId(text);
+    const id = blockId !== undefined ? `b${blockId}` : messageId;
+    if (!id) continue;
+    if (!Number.isFinite(msg.timestamp)) continue;
+    boundaries.push({
+      id,
+      role: msg.role ?? "",
+      timestamp: msg.timestamp,
+      tokenEstimate: estimateMessageTokens(msg),
+      blockId,
+      isSystemReminder: text.includes("<dcp-system-reminder>"),
+    });
+  }
+
+  if (boundaries.length < settings.minMessages) return null;
+
+  const keepRecentTurns = Math.max(1, settings.keepRecentTurns);
+  let recentUserTurns = 0;
+  let cutoffIndex = -1;
+
+  for (let i = boundaries.length - 1; i >= 0; i--) {
+    const boundary = boundaries[i]!;
+    const isRealUserMessage =
+      boundary.role === "user" && boundary.blockId === undefined && !boundary.isSystemReminder;
+    if (!isRealUserMessage) continue;
+    recentUserTurns++;
+    if (recentUserTurns >= keepRecentTurns) {
+      cutoffIndex = i - 1;
+      break;
+    }
+  }
+
+  if (cutoffIndex < 0) return null;
+
+  let candidate = boundaries.slice(0, cutoffIndex + 1);
+  while (candidate.length > 0 && candidate[0]!.isSystemReminder) candidate = candidate.slice(1);
+  while (candidate.length > 0 && candidate[candidate.length - 1]!.isSystemReminder) {
+    candidate = candidate.slice(0, -1);
+  }
+
+  if (candidate.length < settings.minMessages) return null;
+
+  const estimatedTokens = candidate.reduce((sum, item) => sum + item.tokenEstimate, 0);
+  if (estimatedTokens < settings.minTokens) return null;
+
+  const includedBlockIds = Array.from(
+    new Set(candidate.map((item) => item.blockId).filter((id): id is number => id !== undefined)),
+  );
+
+  return {
+    startId: candidate[0]!.id,
+    endId: candidate[candidate.length - 1]!.id,
+    messageCount: candidate.length,
+    estimatedTokens,
+    includedBlockIds,
+    reason: `older than the most recent ${keepRecentTurns} user turn(s)`,
+  };
+}
+
+export function formatCompressionCandidateHint(candidate: CompressionCandidate): string {
+  const blockHint = candidate.includedBlockIds.length > 0
+    ? `\nThis candidate includes compressed block(s): ${candidate.includedBlockIds
+        .map((id) => `b${id}`)
+        .join(", ")}. If you compress this range, include each required \`(bN)\` placeholder exactly once in the summary.`
+    : "";
+
+  return `\n\nSuggested compression candidate: ${candidate.startId}..${candidate.endId} (${candidate.messageCount} messages, ~${candidate.estimatedTokens} tokens, ${candidate.reason}).${blockHint}`;
+}
+
+export function detectMessageCompressionCandidates(
+  messages: any[],
+  config: DcpConfig,
+  contextPercent: number,
+): MessageCompressionCandidate[] {
+  const settings = config.compress.messageMode;
+  if (!settings?.enabled) return [];
+  if (contextPercent < settings.minContextPercent) return [];
+
+  const boundaries: CandidateBoundary[] = [];
+  for (const msg of messages) {
+    const text = messageText(msg);
+    const blockId = extractBlockId(text);
+    const messageId = extractMessageId(text);
+    if (!messageId || blockId !== undefined) continue;
+    if (!Number.isFinite(msg.timestamp)) continue;
+    boundaries.push({
+      id: messageId,
+      role: msg.role ?? "",
+      timestamp: msg.timestamp,
+      tokenEstimate: estimateMessageTokens(msg),
+      blockId,
+      isSystemReminder: text.includes("<dcp-system-reminder>"),
+    });
+  }
+
+  const keepRecentTurns = Math.max(1, settings.keepRecentTurns ?? 2);
+  let recentUserTurns = 0;
+  let cutoffIndex = boundaries.length - 1;
+
+  for (let i = boundaries.length - 1; i >= 0; i--) {
+    const boundary = boundaries[i]!;
+    const isRealUserMessage = boundary.role === "user" && !boundary.isSystemReminder;
+    if (!isRealUserMessage) continue;
+    recentUserTurns++;
+    if (recentUserTurns >= keepRecentTurns) {
+      cutoffIndex = i - 1;
+      break;
+    }
+  }
+
+  if (cutoffIndex < 0) return [];
+
+  const mediumTokens = Math.max(1, settings.mediumTokens ?? 500);
+  const highTokens = Math.max(mediumTokens, settings.highTokens ?? 5000);
+  const maxSuggestions = Math.max(1, settings.maxSuggestions ?? 5);
+
+  return boundaries
+    .slice(0, cutoffIndex + 1)
+    .filter((candidate) => !candidate.isSystemReminder)
+    .filter((candidate) => candidate.role !== "user" || !config.compress.protectUserMessages)
+    .filter((candidate) => candidate.tokenEstimate >= mediumTokens)
+    .map((candidate): MessageCompressionCandidate => ({
+      messageId: candidate.id,
+      role: candidate.role,
+      estimatedTokens: candidate.tokenEstimate,
+      priority: candidate.tokenEstimate >= highTokens ? "high" : "medium",
+      reason: `older than the most recent ${keepRecentTurns} user turn(s)`,
+    }))
+    .sort((a, b) => {
+      const priorityDiff = (b.priority === "high" ? 1 : 0) - (a.priority === "high" ? 1 : 0);
+      if (priorityDiff !== 0) return priorityDiff;
+      return b.estimatedTokens - a.estimatedTokens;
+    })
+    .slice(0, maxSuggestions);
+}
+
+export function formatMessageCompressionCandidateHint(candidates: MessageCompressionCandidate[]): string {
+  if (candidates.length === 0) return "";
+  const entries = candidates
+    .map((candidate) => `${candidate.messageId} (${candidate.priority}, ${candidate.role}, ~${candidate.estimatedTokens} tokens)`)
+    .join(", ");
+  return `\n\nSuggested individual message compression candidates: ${entries}. To compress individual messages, call \`compress\` with a \`messages\` array: { messageId, topic, summary }. Use this for large stale messages when a full range would be too broad.`;
+}

package/external/pi-tools-suite/src/compress/pruner-compression-blocks.ts

@@ -0,0 +1,260 @@
+import type { DcpState } from "./state.js";
+import { PASSTHROUGH_ROLES, estimateMessageTokens } from "./pruner-metadata.js";
+import { stableMessageId } from "./pruner-message-ids.js";
+
+function messageMatchesBoundary(msg: any, messageIndex: number, stableId: string | undefined, timestamp: number): boolean {
+  if (stableId && stableMessageId(msg, messageIndex) === stableId) return true;
+  return msg.timestamp === timestamp;
+}
+
+function findBoundaryIndex(messages: any[], stableId: string | undefined, timestamp: number): number {
+  return messages.findIndex((m, index) => messageMatchesBoundary(m, index, stableId, timestamp));
+}
+
+function collectToolCallIds(messages: any[]): Set<string> {
+  const ids = new Set<string>();
+  for (const msg of messages) {
+    if (typeof msg?.toolCallId === "string") ids.add(msg.toolCallId);
+    if (msg?.role !== "assistant" || !Array.isArray(msg.content)) continue;
+    for (const block of msg.content) {
+      if (block?.type === "toolCall" && typeof block.id === "string") ids.add(block.id);
+    }
+  }
+  return ids;
+}
+
+export function syncCompressionBlocks(messages: any[], state: DcpState): void {
+  if (state.compressionBlocks.length === 0) return;
+
+  const toolCallIds = collectToolCallIds(messages);
+
+  for (const block of state.compressionBlocks) {
+    if (!block.active || block.deactivatedByUser) continue;
+
+    if (
+      typeof block.createdByToolCallId === "string" &&
+      state.toolCalls.has(block.createdByToolCallId) &&
+      !toolCallIds.has(block.createdByToolCallId)
+    ) {
+      block.active = false;
+      block.deactivatedReason = "missing-origin-compress-call";
+      continue;
+    }
+
+    const hasStableBoundaries = !!block.startMessageId && !!block.endMessageId;
+    if (!hasStableBoundaries) continue;
+
+    const startIdx = findBoundaryIndex(messages, block.startMessageId, block.startTimestamp);
+    const endIdx = findBoundaryIndex(messages, block.endMessageId, block.endTimestamp);
+    if (startIdx === -1 || endIdx === -1) {
+      block.active = false;
+      block.deactivatedReason = "missing-origin-message";
+    }
+  }
+}
+
+export function applyCompressionBlocks(messages: any[], state: DcpState): any[] {
+  const activeBlocks = state.compressionBlocks
+    .filter((b) => b.active)
+    .sort((a, b) => a.startTimestamp - b.startTimestamp);
+  if (activeBlocks.length === 0) return messages;
+
+  for (const block of activeBlocks) {
+    // Skip blocks with corrupted timestamps (from pre-fix sessions)
+    if (!Number.isFinite(block.startTimestamp) || !Number.isFinite(block.endTimestamp)) continue;
+
+    // Find start and end indices by timestamp
+    const startIdx = findBoundaryIndex(messages, block.startMessageId, block.startTimestamp);
+    const endIdx = findBoundaryIndex(messages, block.endMessageId, block.endTimestamp);
+
+    if (startIdx === -1 || endIdx === -1) continue;
+
+    let lo = Math.min(startIdx, endIdx);
+    let hi = Math.max(startIdx, endIdx);
+
+    // Expand lo backward: if there is an assistant before lo whose tool_use
+    // blocks have matching tool_results inside [lo..hi], pull the entire
+    // assistant + any intermediate result messages into the range so the
+    // group is always removed atomically.
+    //
+    // Critically we must skip backward past any toolResult / bashExecution
+    // messages before lo, because an assistant with multiple tool_calls emits
+    // N consecutive result messages — the assistant itself sits further back.
+    while (lo > 0) {
+      // Walk backward past tool-result messages to find the preceding assistant
+      let scanIdx = lo - 1;
+      while (scanIdx >= 0) {
+        const r = (messages[scanIdx] as any).role as string;
+        if (r !== "toolResult" && r !== "bashExecution" && !PASSTHROUGH_ROLES.has(r)) break;
+        scanIdx--;
+      }
+      if (scanIdx < 0 || (messages[scanIdx] as any).role !== "assistant") break;
+
+      const prev = messages[scanIdx] as any;
+      const toolCallIdsInRange = new Set<string>();
+      for (let i = lo; i <= hi; i++) {
+        const m = messages[i] as any;
+        if (
+          (m.role === "toolResult" || m.role === "bashExecution") &&
+          typeof m.toolCallId === "string"
+        ) {
+          toolCallIdsInRange.add(m.toolCallId);
+        }
+      }
+      const prevContent: any[] = Array.isArray(prev.content) ? prev.content : [];
+      const hasMatchingToolCalls = prevContent.some(
+        (contentBlock: any) => contentBlock.type === "toolCall" && toolCallIdsInRange.has(contentBlock.id),
+      );
+      if (!hasMatchingToolCalls) break;
+      // Pull assistant + all intermediate result messages into the range
+      lo = scanIdx;
+    }
+
+    // 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 contentBlock of content) {
+          if (contentBlock.type === "toolCall" && typeof contentBlock.id === "string") {
+            assistantToolCallIds.add(contentBlock.id);
+          }
+        }
+      }
+      while (hi + 1 < messages.length) {
+        const next = messages[hi + 1] as any;
+        if (
+          (next.role === "toolResult" || next.role === "bashExecution") &&
+          assistantToolCallIds.has(next.toolCallId)
+        ) {
+          hi++;
+        } else if (PASSTHROUGH_ROLES.has(next.role)) {
+          hi++;
+        } else {
+          break;
+        }
+      }
+    } while (hi !== prevHi);
+
+    // Estimate tokens removed
+    let removedTokens = 0;
+    for (let i = lo; i <= hi; i++) {
+      removedTokens += estimateMessageTokens(messages[i]);
+    }
+
+    // Remove the range (inclusive)
+    messages.splice(lo, hi - lo + 1);
+
+    // Build synthetic user message for the compressed block
+    const syntheticMsg = {
+      role: "user",
+      content: [
+        {
+          type: "text",
+          text:
+            "[Compressed section: " +
+            block.topic +
+            "]\n\n" +
+            block.summary +
+            "\n\n[dcp-block-id]: # (b" +
+            block.id +
+            ")",
+        },
+      ],
+      // anchorTimestamp is always finite (resolveAnchorTimestamp returns
+      // endTimestamp + 1 instead of Infinity), but guard against corrupted
+      // state from older sessions where Infinity/null could leak in.
+      timestamp: Number.isFinite(block.anchorTimestamp) ? block.anchorTimestamp - 0.5 : block.endTimestamp + 0.5,
+    };
+
+    // Estimate tokens added by the summary
+    const addedTokens = estimateMessageTokens(syntheticMsg);
+
+    // Insert the synthetic message
+    messages.push(syntheticMsg);
+
+    // Re-sort by timestamp
+    messages.sort((a, b) => (a.timestamp ?? 0) - (b.timestamp ?? 0));
+
+    // Update tokens saved exactly once per compression block.
+    if (!state.accountedCompressionBlockIds.has(block.id)) {
+      state.accountedCompressionBlockIds.add(block.id);
+      state.totalPruneCount++;
+      const rawSaved = Math.max(0, removedTokens - addedTokens);
+      const coveredSavings = (block.coveredBlockIds ?? []).reduce(
+        (sum, id) => sum + (state.compressionTokenSavings.get(id) ?? 0),
+        0,
+      );
+      const adjustment = rawSaved - coveredSavings;
+      state.tokensSaved = Math.max(0, state.tokensSaved + adjustment);
+      state.compressionTokenSavings.set(block.id, rawSaved);
+    }
+  }
+
+  return messages;
+}
+
+/**
+ * Remove orphaned toolResult/bashExecution messages whose corresponding
+ * assistant toolCall was removed, and strip orphaned toolCall blocks from
+ * assistant messages whose toolResult was removed.
+ *
+ * This is a safety net that runs after all compression blocks are applied.
+ */
+export function repairOrphanedToolPairs(messages: any[]): void {
+  // 1. Build set of all toolCall IDs present in assistant messages
+  const assistantToolCallIds = new Set<string>();
+  for (const msg of messages) {
+    if (msg.role !== "assistant") continue;
+    const content: any[] = Array.isArray(msg.content) ? msg.content : [];
+    for (const contentBlock of content) {
+      if (contentBlock.type === "toolCall" && typeof contentBlock.id === "string") {
+        assistantToolCallIds.add(contentBlock.id);
+      }
+    }
+  }
+
+  // 2. Build set of all toolCallIds present in toolResult/bashExecution messages
+  const resultToolCallIds = new Set<string>();
+  for (const msg of messages) {
+    if (msg.role !== "toolResult" && msg.role !== "bashExecution") continue;
+    if (typeof msg.toolCallId === "string") {
+      resultToolCallIds.add(msg.toolCallId);
+    }
+  }
+
+  // 3. Remove orphaned toolResult/bashExecution messages (no matching assistant toolCall)
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i];
+    if (msg.role !== "toolResult" && msg.role !== "bashExecution") continue;
+    if (typeof msg.toolCallId === "string" && !assistantToolCallIds.has(msg.toolCallId)) {
+      messages.splice(i, 1);
+    }
+  }
+
+  // 4. Strip orphaned toolCall blocks from assistant messages (no matching toolResult)
+  for (const msg of messages) {
+    if (msg.role !== "assistant") continue;
+    const content: any[] = Array.isArray(msg.content) ? msg.content : [];
+    const hasToolCalls = content.some((b: any) => b.type === "toolCall");
+    if (!hasToolCalls) continue;
+
+    const filtered = content.filter((contentBlock: any) => {
+      if (contentBlock.type !== "toolCall") return true;
+      return typeof contentBlock.id === "string" && resultToolCallIds.has(contentBlock.id);
+    });
+
+    // Only update if we actually removed something
+    if (filtered.length !== content.length) {
+      // If the assistant has no content left at all, keep at least an empty array
+      msg.content = filtered.length > 0 ? filtered : [];
+    }
+  }
+}

package/external/pi-tools-suite/src/compress/pruner-message-ids.ts

@@ -0,0 +1,147 @@
+import type { DcpState } from "./state.js";
+import type { DcpConfig } from "./config.js";
+import {
+  ID_ELIGIBLE_ROLES,
+  PASSTHROUGH_ROLES,
+  estimateMessageTokens,
+  extractBlockId,
+  messageText,
+} from "./pruner-metadata.js";
+
+export interface InjectMessageIdsOptions {
+  /**
+   * When false, rebuild internal mNNN snapshots without appending visible
+   * DCP marker lines to provider-visible message content.
+   */
+  visible?: boolean;
+  /** Config enables priority markers for message-mode candidates. */
+  config?: DcpConfig;
+}
+
+export function stableMessageId(msg: any, fallbackIndex = 0): string {
+  const candidates = [
+    msg?.id,
+    msg?.entryId,
+    msg?.messageId,
+    msg?._dcpEntryId,
+    msg?.metadata?.id,
+    msg?.metadata?.entryId,
+  ];
+  for (const candidate of candidates) {
+    if (typeof candidate === "string" && candidate.trim().length > 0) {
+      return `id:${candidate.trim()}`;
+    }
+  }
+
+  if (typeof msg?.toolCallId === "string" && msg.toolCallId.trim().length > 0) {
+    return `tool:${msg.toolCallId.trim()}`;
+  }
+
+  if (Number.isFinite(msg?.timestamp)) return `ts:${msg.timestamp}`;
+  return `idx:${fallbackIndex}`;
+}
+
+function priorityForMessage(tokenEstimate: number, config: DcpConfig | undefined): "low" | "medium" | "high" {
+  const settings = config?.compress?.messageMode;
+  const mediumTokens = Math.max(1, settings?.mediumTokens ?? 500);
+  const highTokens = Math.max(mediumTokens, settings?.highTokens ?? 5000);
+
+  if (tokenEstimate >= highTokens) return "high";
+  if (tokenEstimate >= mediumTokens) return "medium";
+  return "low";
+}
+
+function formatIdTag(id: string, _priority: "low" | "medium" | "high"): string {
+  // Keep the provider-visible marker itself short and single-line. Priority is
+  // retained in state and emitted in message-compression candidate hints.
+  return `\n[dcp-id]: # (${id})`;
+}
+
+export function injectMessageIds(
+  messages: any[],
+  state: DcpState,
+  options: InjectMessageIdsOptions = {},
+): void {
+  const visible = options.visible ?? true;
+
+  // Clear the snapshots and rebuild
+  state.messageIdSnapshot.clear();
+  state.messageMetaSnapshot.clear();
+
+  let counter = 1;
+
+  for (let messageIndex = 0; messageIndex < messages.length; messageIndex++) {
+    const msg = messages[messageIndex];
+    const role: string = msg.role ?? "";
+
+    // Skip PI-internal passthrough messages
+    if (PASSTHROUGH_ROLES.has(role)) continue;
+    // Skip non-eligible roles
+    if (!ID_ELIGIBLE_ROLES.has(role)) continue;
+
+    const id = "m" + String(counter).padStart(3, "0");
+    counter++;
+
+    const originalText = messageText(msg);
+    const blockId = extractBlockId(originalText);
+    const tokenEstimate = estimateMessageTokens(msg);
+    const priority = priorityForMessage(tokenEstimate, options.config);
+    const idTag = formatIdTag(id, priority);
+    const rawStableId = stableMessageId(msg, messageIndex);
+
+    if (!visible) {
+      // Snapshot-only mode: keep mNNN mappings fresh for DCP internals, but
+      // do not expose synthetic IDs to the agent.
+    } else if (role === "user") {
+      if (typeof msg.content === "string") {
+        msg.content = msg.content + `\n${idTag}`;
+      } else if (Array.isArray(msg.content)) {
+        msg.content = [...msg.content, { type: "text", text: idTag }];
+      }
+    } else if (role === "toolResult" || role === "bashExecution") {
+      if (Array.isArray(msg.content)) {
+        msg.content = [...msg.content, { type: "text", text: idTag }];
+      } else if (typeof msg.content === "string") {
+        msg.content = msg.content + idTag;
+      }
+    } else if (role === "assistant") {
+      if (Array.isArray(msg.content)) {
+        // Insert the ID tag before any tool_use (toolCall) blocks.
+        // Anthropic requires: thinking → text → tool_use.
+        // Appending after tool_use blocks violates that constraint.
+        const firstToolCallIdx = msg.content.findIndex(
+          (b: any) => b.type === "toolCall",
+        );
+        const idBlock = { type: "text", text: idTag };
+        if (firstToolCallIdx === -1) {
+          // No tool_use blocks — append as usual
+          msg.content = [...msg.content, idBlock];
+        } else {
+          // Insert immediately before the first tool_use block
+          msg.content = [
+            ...msg.content.slice(0, firstToolCallIdx),
+            idBlock,
+            ...msg.content.slice(firstToolCallIdx),
+          ];
+        }
+      } else if (typeof msg.content === "string") {
+        msg.content = msg.content + idTag;
+      }
+    }
+
+    if (msg.timestamp !== undefined) {
+      state.messageIdSnapshot.set(id, msg.timestamp);
+      state.messageMetaSnapshot.set(id, {
+        timestamp: msg.timestamp,
+        stableId: rawStableId,
+        role,
+        blockId,
+        toolCallId: typeof msg.toolCallId === "string" ? msg.toolCallId : undefined,
+        toolName: typeof msg.toolName === "string" ? msg.toolName : undefined,
+        text: originalText,
+        tokenEstimate,
+        priority,
+      });
+    }
+  }
+}