@poolzin/pool-bot 2026.1.39 → 2026.2.1

package/dist/agents/pi-embedded-runner/system-prompt.js

@@ -25,9 +25,17 @@ export function buildEmbeddedSystemPrompt(params) {
         userTime: params.userTime,
         userTimeFormat: params.userTimeFormat,
         contextFiles: params.contextFiles,
+        memoryCitationsMode: params.memoryCitationsMode,
     });
 }
 export function createSystemPromptOverride(systemPrompt) {
-    const trimmed = systemPrompt.trim();
-    return () => trimmed;
+    const override = systemPrompt.trim();
+    return (_defaultPrompt) => override;
+}
+export function applySystemPromptOverrideToSession(session, override) {
+    const prompt = typeof override === "function" ? override() : override.trim();
+    session.agent.setSystemPrompt(prompt);
+    const mutableSession = session;
+    mutableSession._baseSystemPrompt = prompt;
+    mutableSession._rebuildSystemPrompt = () => prompt;
 }

package/dist/agents/pi-embedded-runner/tool-result-truncation.js

@@ -0,0 +1,275 @@
+import { SessionManager } from "@mariozechner/pi-coding-agent";
+import { log } from "./logger.js";
+/**
+ * Maximum share of the context window a single tool result should occupy.
+ * This is intentionally conservative – a single tool result should not
+ * consume more than 30% of the context window even without other messages.
+ */
+const MAX_TOOL_RESULT_CONTEXT_SHARE = 0.3;
+/**
+ * Hard character limit for a single tool result text block.
+ * Even for the largest context windows (~2M tokens), a single tool result
+ * should not exceed ~400K characters (~100K tokens).
+ * This acts as a safety net when we don't know the context window size.
+ */
+export const HARD_MAX_TOOL_RESULT_CHARS = 400_000;
+/**
+ * Minimum characters to keep when truncating.
+ * We always keep at least the first portion so the model understands
+ * what was in the content.
+ */
+const MIN_KEEP_CHARS = 2_000;
+/**
+ * Suffix appended to truncated tool results.
+ */
+const TRUNCATION_SUFFIX = "\n\n⚠️ [Content truncated — original was too large for the model's context window. " +
+    "The content above is a partial view. If you need more, request specific sections or use " +
+    "offset/limit parameters to read smaller chunks.]";
+/**
+ * Truncate a single text string to fit within maxChars, preserving the beginning.
+ */
+export function truncateToolResultText(text, maxChars) {
+    if (text.length <= maxChars) {
+        return text;
+    }
+    const keepChars = Math.max(MIN_KEEP_CHARS, maxChars - TRUNCATION_SUFFIX.length);
+    // Try to break at a newline boundary to avoid cutting mid-line
+    let cutPoint = keepChars;
+    const lastNewline = text.lastIndexOf("\n", keepChars);
+    if (lastNewline > keepChars * 0.8) {
+        cutPoint = lastNewline;
+    }
+    return text.slice(0, cutPoint) + TRUNCATION_SUFFIX;
+}
+/**
+ * Calculate the maximum allowed characters for a single tool result
+ * based on the model's context window tokens.
+ *
+ * Uses a rough 4 chars ≈ 1 token heuristic (conservative for English text;
+ * actual ratio varies by tokenizer).
+ */
+export function calculateMaxToolResultChars(contextWindowTokens) {
+    const maxTokens = Math.floor(contextWindowTokens * MAX_TOOL_RESULT_CONTEXT_SHARE);
+    // Rough conversion: ~4 chars per token on average
+    const maxChars = maxTokens * 4;
+    return Math.min(maxChars, HARD_MAX_TOOL_RESULT_CHARS);
+}
+/**
+ * Get the total character count of text content blocks in a tool result message.
+ */
+function getToolResultTextLength(msg) {
+    if (!msg || msg.role !== "toolResult") {
+        return 0;
+    }
+    const content = msg.content;
+    if (!Array.isArray(content)) {
+        return 0;
+    }
+    let totalLength = 0;
+    for (const block of content) {
+        if (block && typeof block === "object" && block.type === "text") {
+            const text = block.text;
+            if (typeof text === "string") {
+                totalLength += text.length;
+            }
+        }
+    }
+    return totalLength;
+}
+/**
+ * Truncate a tool result message's text content blocks to fit within maxChars.
+ * Returns a new message (does not mutate the original).
+ */
+function truncateToolResultMessage(msg, maxChars) {
+    const content = msg.content;
+    if (!Array.isArray(content)) {
+        return msg;
+    }
+    // Calculate total text size
+    const totalTextChars = getToolResultTextLength(msg);
+    if (totalTextChars <= maxChars) {
+        return msg;
+    }
+    // Distribute the budget proportionally among text blocks
+    const newContent = content.map((block) => {
+        if (!block || typeof block !== "object" || block.type !== "text") {
+            return block; // Keep non-text blocks (images) as-is
+        }
+        const textBlock = block;
+        if (typeof textBlock.text !== "string") {
+            return block;
+        }
+        // Proportional budget for this block
+        const blockShare = textBlock.text.length / totalTextChars;
+        const blockBudget = Math.max(MIN_KEEP_CHARS, Math.floor(maxChars * blockShare));
+        return {
+            ...textBlock,
+            text: truncateToolResultText(textBlock.text, blockBudget),
+        };
+    });
+    return { ...msg, content: newContent };
+}
+/**
+ * Find oversized tool result entries in a session and truncate them.
+ *
+ * This operates on the session file by:
+ * 1. Opening the session manager
+ * 2. Walking the current branch to find oversized tool results
+ * 3. Branching from before the first oversized tool result
+ * 4. Re-appending all entries from that point with truncated tool results
+ *
+ * @returns Object indicating whether any truncation was performed
+ */
+export async function truncateOversizedToolResultsInSession(params) {
+    const { sessionFile, contextWindowTokens } = params;
+    const maxChars = calculateMaxToolResultChars(contextWindowTokens);
+    try {
+        const sessionManager = SessionManager.open(sessionFile);
+        const branch = sessionManager.getBranch();
+        if (branch.length === 0) {
+            return { truncated: false, truncatedCount: 0, reason: "empty session" };
+        }
+        // Find oversized tool result entries and their indices in the branch
+        const oversizedIndices = [];
+        for (let i = 0; i < branch.length; i++) {
+            const entry = branch[i];
+            if (entry.type !== "message") {
+                continue;
+            }
+            const msg = entry.message;
+            if (msg.role !== "toolResult") {
+                continue;
+            }
+            const textLength = getToolResultTextLength(msg);
+            if (textLength > maxChars) {
+                oversizedIndices.push(i);
+                log.info(`[tool-result-truncation] Found oversized tool result: ` +
+                    `entry=${entry.id} chars=${textLength} maxChars=${maxChars} ` +
+                    `sessionKey=${params.sessionKey ?? params.sessionId ?? "unknown"}`);
+            }
+        }
+        if (oversizedIndices.length === 0) {
+            return { truncated: false, truncatedCount: 0, reason: "no oversized tool results" };
+        }
+        // Branch from the parent of the first oversized entry
+        const firstOversizedIdx = oversizedIndices[0];
+        const firstOversizedEntry = branch[firstOversizedIdx];
+        const branchFromId = firstOversizedEntry.parentId;
+        if (!branchFromId) {
+            // The oversized entry is the root - very unusual but handle it
+            sessionManager.resetLeaf();
+        }
+        else {
+            sessionManager.branch(branchFromId);
+        }
+        // Re-append all entries from the first oversized one onwards,
+        // with truncated tool results
+        const oversizedSet = new Set(oversizedIndices);
+        let truncatedCount = 0;
+        for (let i = firstOversizedIdx; i < branch.length; i++) {
+            const entry = branch[i];
+            if (entry.type === "message") {
+                let message = entry.message;
+                if (oversizedSet.has(i)) {
+                    message = truncateToolResultMessage(message, maxChars);
+                    truncatedCount++;
+                    const newLength = getToolResultTextLength(message);
+                    log.info(`[tool-result-truncation] Truncated tool result: ` +
+                        `originalEntry=${entry.id} newChars=${newLength} ` +
+                        `sessionKey=${params.sessionKey ?? params.sessionId ?? "unknown"}`);
+                }
+                // appendMessage expects Message | CustomMessage | BashExecutionMessage
+                sessionManager.appendMessage(message);
+            }
+            else if (entry.type === "compaction") {
+                sessionManager.appendCompaction(entry.summary, entry.firstKeptEntryId, entry.tokensBefore, entry.details, entry.fromHook);
+            }
+            else if (entry.type === "thinking_level_change") {
+                sessionManager.appendThinkingLevelChange(entry.thinkingLevel);
+            }
+            else if (entry.type === "model_change") {
+                sessionManager.appendModelChange(entry.provider, entry.modelId);
+            }
+            else if (entry.type === "custom") {
+                sessionManager.appendCustomEntry(entry.customType, entry.data);
+            }
+            else if (entry.type === "custom_message") {
+                sessionManager.appendCustomMessageEntry(entry.customType, entry.content, entry.display, entry.details);
+            }
+            else if (entry.type === "branch_summary") {
+                // Branch summaries reference specific entry IDs - skip to avoid inconsistency
+                continue;
+            }
+            else if (entry.type === "label") {
+                // Labels reference specific entry IDs - skip to avoid inconsistency
+                continue;
+            }
+            else if (entry.type === "session_info") {
+                if (entry.name) {
+                    sessionManager.appendSessionInfo(entry.name);
+                }
+            }
+        }
+        log.info(`[tool-result-truncation] Truncated ${truncatedCount} tool result(s) in session ` +
+            `(contextWindow=${contextWindowTokens} maxChars=${maxChars}) ` +
+            `sessionKey=${params.sessionKey ?? params.sessionId ?? "unknown"}`);
+        return { truncated: true, truncatedCount };
+    }
+    catch (err) {
+        const errMsg = err instanceof Error ? err.message : String(err);
+        log.warn(`[tool-result-truncation] Failed to truncate: ${errMsg}`);
+        return { truncated: false, truncatedCount: 0, reason: errMsg };
+    }
+}
+/**
+ * Truncate oversized tool results in an array of messages (in-memory).
+ * Returns a new array with truncated messages.
+ *
+ * This is used as a pre-emptive guard before sending messages to the LLM,
+ * without modifying the session file.
+ */
+export function truncateOversizedToolResultsInMessages(messages, contextWindowTokens) {
+    const maxChars = calculateMaxToolResultChars(contextWindowTokens);
+    let truncatedCount = 0;
+    const result = messages.map((msg) => {
+        if (msg.role !== "toolResult") {
+            return msg;
+        }
+        const textLength = getToolResultTextLength(msg);
+        if (textLength <= maxChars) {
+            return msg;
+        }
+        truncatedCount++;
+        return truncateToolResultMessage(msg, maxChars);
+    });
+    return { messages: result, truncatedCount };
+}
+/**
+ * Check if a tool result message exceeds the size limit for a given context window.
+ */
+export function isOversizedToolResult(msg, contextWindowTokens) {
+    if (msg.role !== "toolResult") {
+        return false;
+    }
+    const maxChars = calculateMaxToolResultChars(contextWindowTokens);
+    return getToolResultTextLength(msg) > maxChars;
+}
+/**
+ * Estimate whether the session likely has oversized tool results that caused
+ * a context overflow. Used as a heuristic to decide whether to attempt
+ * tool result truncation before giving up.
+ */
+export function sessionLikelyHasOversizedToolResults(params) {
+    const { messages, contextWindowTokens } = params;
+    const maxChars = calculateMaxToolResultChars(contextWindowTokens);
+    for (const msg of messages) {
+        if (msg.role !== "toolResult") {
+            continue;
+        }
+        const textLength = getToolResultTextLength(msg);
+        if (textLength > maxChars) {
+            return true;
+        }
+    }
+    return false;
+}

package/dist/agents/pi-embedded-runner/utils.js

@@ -1,5 +1,5 @@
 export function mapThinkingLevel(level) {
-    // pi-agent-core supports "xhigh"; Moltbot enables it for specific models.
+    // pi-agent-core supports "xhigh"; Poolbot enables it for specific models.
     if (!level)
         return "off";
     return level;

package/dist/agents/pi-embedded-subscribe.js

@@ -6,7 +6,8 @@ import { buildCodeSpanIndex, createInlineCodeState } from "../markdown/code-span
 import { EmbeddedBlockChunker } from "./pi-embedded-block-chunker.js";
 import { isMessagingToolDuplicateNormalized, normalizeTextForComparison, } from "./pi-embedded-helpers.js";
 import { createEmbeddedPiSessionEventHandler } from "./pi-embedded-subscribe.handlers.js";
-import { formatReasoningMessage } from "./pi-embedded-utils.js";
+import { formatReasoningMessage, stripDowngradedToolCallText } from "./pi-embedded-utils.js";
+import { hasNonzeroUsage, normalizeUsage } from "./usage.js";
 const THINKING_TAG_SCAN_RE = /<\s*(\/?)\s*(?:think(?:ing)?|thought|antthinking)\s*>/gi;
 const FINAL_TAG_SCAN_RE = /<\s*(\/?)\s*final\s*>/gi;
 const log = createSubsystemLogger("agent/embedded");
@@ -29,7 +30,10 @@ export function subscribeEmbeddedPiSession(params) {
         blockBuffer: "",
         // Track if a streamed chunk opened a <think> block (stateful across chunks).
         blockState: { thinking: false, final: false, inlineCode: createInlineCodeState() },
+        partialBlockState: { thinking: false, final: false, inlineCode: createInlineCodeState() },
         lastStreamedAssistant: undefined,
+        lastStreamedAssistantCleaned: undefined,
+        emittedAssistantUpdate: false,
         lastStreamedReasoning: undefined,
         lastBlockReplyText: undefined,
         assistantMessageIndex: 0,
@@ -49,6 +53,14 @@ export function subscribeEmbeddedPiSession(params) {
         pendingMessagingTexts: new Map(),
         pendingMessagingTargets: new Map(),
     };
+    const usageTotals = {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        total: 0,
+    };
+    let compactionCount = 0;
     const assistantTexts = state.assistantTexts;
     const toolMetas = state.toolMetas;
     const toolMetaById = state.toolMetaById;
@@ -59,15 +71,22 @@ export function subscribeEmbeddedPiSession(params) {
     const pendingMessagingTexts = state.pendingMessagingTexts;
     const pendingMessagingTargets = state.pendingMessagingTargets;
     const replyDirectiveAccumulator = createStreamingDirectiveAccumulator();
+    const partialReplyDirectiveAccumulator = createStreamingDirectiveAccumulator();
     const resetAssistantMessageState = (nextAssistantTextBaseline) => {
         state.deltaBuffer = "";
         state.blockBuffer = "";
         blockChunker?.reset();
         replyDirectiveAccumulator.reset();
+        partialReplyDirectiveAccumulator.reset();
         state.blockState.thinking = false;
         state.blockState.final = false;
         state.blockState.inlineCode = createInlineCodeState();
+        state.partialBlockState.thinking = false;
+        state.partialBlockState.final = false;
+        state.partialBlockState.inlineCode = createInlineCodeState();
         state.lastStreamedAssistant = undefined;
+        state.lastStreamedAssistantCleaned = undefined;
+        state.emittedAssistantUpdate = false;
         state.lastBlockReplyText = undefined;
         state.lastStreamedReasoning = undefined;
         state.lastReasoningSent = undefined;
@@ -85,21 +104,26 @@ export function subscribeEmbeddedPiSession(params) {
         state.lastAssistantTextNormalized = normalized.length > 0 ? normalized : undefined;
     };
     const shouldSkipAssistantText = (text) => {
-        if (state.lastAssistantTextMessageIndex !== state.assistantMessageIndex)
+        if (state.lastAssistantTextMessageIndex !== state.assistantMessageIndex) {
             return false;
+        }
         const trimmed = text.trimEnd();
-        if (trimmed && trimmed === state.lastAssistantTextTrimmed)
+        if (trimmed && trimmed === state.lastAssistantTextTrimmed) {
             return true;
+        }
         const normalized = normalizeTextForComparison(text);
-        if (normalized.length > 0 && normalized === state.lastAssistantTextNormalized)
+        if (normalized.length > 0 && normalized === state.lastAssistantTextNormalized) {
             return true;
+        }
         return false;
     };
     const pushAssistantText = (text) => {
-        if (!text)
+        if (!text) {
             return;
-        if (shouldSkipAssistantText(text))
+        }
+        if (shouldSkipAssistantText(text)) {
             return;
+        }
         assistantTexts.push(text);
         rememberAssistantText(text);
     };
@@ -154,8 +178,9 @@ export function subscribeEmbeddedPiSession(params) {
         ensureCompactionPromise();
     };
     const resolveCompactionRetry = () => {
-        if (state.pendingCompactionRetry <= 0)
+        if (state.pendingCompactionRetry <= 0) {
             return;
+        }
         state.pendingCompactionRetry -= 1;
         if (state.pendingCompactionRetry === 0 && !state.compactionInFlight) {
             state.compactionRetryResolve?.();
@@ -170,6 +195,40 @@ export function subscribeEmbeddedPiSession(params) {
             state.compactionRetryPromise = null;
         }
     };
+    const recordAssistantUsage = (usageLike) => {
+        const usage = normalizeUsage((usageLike ?? undefined));
+        if (!hasNonzeroUsage(usage)) {
+            return;
+        }
+        usageTotals.input += usage.input ?? 0;
+        usageTotals.output += usage.output ?? 0;
+        usageTotals.cacheRead += usage.cacheRead ?? 0;
+        usageTotals.cacheWrite += usage.cacheWrite ?? 0;
+        const usageTotal = usage.total ??
+            (usage.input ?? 0) + (usage.output ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
+        usageTotals.total += usageTotal;
+    };
+    const getUsageTotals = () => {
+        const hasUsage = usageTotals.input > 0 ||
+            usageTotals.output > 0 ||
+            usageTotals.cacheRead > 0 ||
+            usageTotals.cacheWrite > 0 ||
+            usageTotals.total > 0;
+        if (!hasUsage) {
+            return undefined;
+        }
+        const derivedTotal = usageTotals.input + usageTotals.output + usageTotals.cacheRead + usageTotals.cacheWrite;
+        return {
+            input: usageTotals.input || undefined,
+            output: usageTotals.output || undefined,
+            cacheRead: usageTotals.cacheRead || undefined,
+            cacheWrite: usageTotals.cacheWrite || undefined,
+            total: usageTotals.total || derivedTotal || undefined,
+        };
+    };
+    const incrementCompactionCount = () => {
+        compactionCount += 1;
+    };
     const blockChunking = params.blockReplyChunking;
     const blockChunker = blockChunking ? new EmbeddedBlockChunker(blockChunking) : null;
     // KNOWN: Provider streams are not strictly once-only or perfectly ordered.
@@ -183,21 +242,25 @@ export function subscribeEmbeddedPiSession(params) {
         : params.verboseLevel === "full";
     const formatToolOutputBlock = (text) => {
         const trimmed = text.trim();
-        if (!trimmed)
+        if (!trimmed) {
             return "(no output)";
-        if (!useMarkdown)
+        }
+        if (!useMarkdown) {
             return trimmed;
+        }
         return `\`\`\`txt\n${trimmed}\n\`\`\``;
     };
     const emitToolSummary = (toolName, meta) => {
-        if (!params.onToolResult)
+        if (!params.onToolResult) {
             return;
+        }
         const agg = formatToolAggregate(toolName, meta ? [meta] : undefined, {
             markdown: useMarkdown,
         });
         const { text: cleanedText, mediaUrls } = parseReplyDirectives(agg);
-        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0))
+        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0)) {
             return;
+        }
         try {
             void params.onToolResult({
                 text: cleanedText,
@@ -209,15 +272,17 @@ export function subscribeEmbeddedPiSession(params) {
         }
     };
     const emitToolOutput = (toolName, meta, output) => {
-        if (!params.onToolResult || !output)
+        if (!params.onToolResult || !output) {
             return;
+        }
         const agg = formatToolAggregate(toolName, meta ? [meta] : undefined, {
             markdown: useMarkdown,
         });
         const message = `${agg}\n${formatToolOutputBlock(output)}`;
         const { text: cleanedText, mediaUrls } = parseReplyDirectives(message);
-        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0))
+        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0)) {
             return;
+        }
         try {
             void params.onToolResult({
                 text: cleanedText,
@@ -229,8 +294,9 @@ export function subscribeEmbeddedPiSession(params) {
         }
     };
     const stripBlockTags = (text, state) => {
-        if (!text)
+        if (!text) {
             return text;
+        }
         const inlineStateStart = state.inlineCode ?? createInlineCodeState();
         const codeSpans = buildCodeSpanIndex(text, inlineStateStart);
         // 1. Handle <think> blocks (stateful, strip content inside)
@@ -240,8 +306,9 @@ export function subscribeEmbeddedPiSession(params) {
         let inThinking = state.thinking;
         for (const match of text.matchAll(THINKING_TAG_SCAN_RE)) {
             const idx = match.index ?? 0;
-            if (codeSpans.isInside(idx))
+            if (codeSpans.isInside(idx)) {
                 continue;
+            }
             if (!inThinking) {
                 processed += text.slice(lastIndex, idx);
             }
@@ -271,8 +338,9 @@ export function subscribeEmbeddedPiSession(params) {
         let everInFinal = state.final;
         for (const match of processed.matchAll(FINAL_TAG_SCAN_RE)) {
             const idx = match.index ?? 0;
-            if (finalCodeSpans.isInside(idx))
+            if (finalCodeSpans.isInside(idx)) {
                 continue;
+            }
             const isClose = match[1] === "/";
             if (!inFinal && !isClose) {
                 // Found <final> start tag.
@@ -309,8 +377,9 @@ export function subscribeEmbeddedPiSession(params) {
         pattern.lastIndex = 0;
         for (const match of text.matchAll(pattern)) {
             const idx = match.index ?? 0;
-            if (isInside(idx))
+            if (isInside(idx)) {
                 continue;
+            }
             output += text.slice(lastIndex, idx);
             lastIndex = idx + match[0].length;
         }
@@ -318,14 +387,18 @@ export function subscribeEmbeddedPiSession(params) {
         return output;
     };
     const emitBlockChunk = (text) => {
-        if (state.suppressBlockChunks)
+        if (state.suppressBlockChunks) {
             return;
+        }
         // Strip <think> and <final> blocks across chunk boundaries to avoid leaking reasoning.
-        const chunk = stripBlockTags(text, state.blockState).trimEnd();
-        if (!chunk)
+        // Also strip downgraded tool call text ([Tool Call: ...], [Historical context: ...], etc.).
+        const chunk = stripDowngradedToolCallText(stripBlockTags(text, state.blockState)).trimEnd();
+        if (!chunk) {
             return;
-        if (chunk === state.lastBlockReplyText)
+        }
+        if (chunk === state.lastBlockReplyText) {
             return;
+        }
         // Only check committed (successful) messaging tool texts - checking pending texts
         // is risky because if the tool fails after suppression, the user gets no response
         const normalizedChunk = normalizeTextForComparison(chunk);
@@ -333,20 +406,24 @@ export function subscribeEmbeddedPiSession(params) {
             log.debug(`Skipping block reply - already sent via messaging tool: ${chunk.slice(0, 50)}...`);
             return;
         }
-        if (shouldSkipAssistantText(chunk))
+        if (shouldSkipAssistantText(chunk)) {
             return;
+        }
         state.lastBlockReplyText = chunk;
         assistantTexts.push(chunk);
         rememberAssistantText(chunk);
-        if (!params.onBlockReply)
+        if (!params.onBlockReply) {
             return;
+        }
         const splitResult = replyDirectiveAccumulator.consume(chunk);
-        if (!splitResult)
+        if (!splitResult) {
             return;
+        }
         const { text: cleanedText, mediaUrls, audioAsVoice, replyToId, replyToTag, replyToCurrent, } = splitResult;
         // Skip empty payloads, but always emit if audioAsVoice is set (to propagate the flag)
-        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0) && !audioAsVoice)
+        if (!cleanedText && (!mediaUrls || mediaUrls.length === 0) && !audioAsVoice) {
             return;
+        }
         void params.onBlockReply({
             text: cleanedText,
             mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
@@ -357,9 +434,11 @@ export function subscribeEmbeddedPiSession(params) {
         });
     };
     const consumeReplyDirectives = (text, options) => replyDirectiveAccumulator.consume(text, options);
+    const consumePartialReplyDirectives = (text, options) => partialReplyDirectiveAccumulator.consume(text, options);
     const flushBlockReplyBuffer = () => {
-        if (!params.onBlockReply)
+        if (!params.onBlockReply) {
             return;
+        }
         if (blockChunker?.hasBuffered()) {
             blockChunker.drain({ force: true, emit: emitBlockChunk });
             blockChunker.reset();
@@ -371,13 +450,16 @@ export function subscribeEmbeddedPiSession(params) {
         }
     };
     const emitReasoningStream = (text) => {
-        if (!state.streamReasoning || !params.onReasoningStream)
+        if (!state.streamReasoning || !params.onReasoningStream) {
             return;
+        }
         const formatted = formatReasoningMessage(text);
-        if (!formatted)
+        if (!formatted) {
             return;
-        if (formatted === state.lastStreamedReasoning)
+        }
+        if (formatted === state.lastStreamedReasoning) {
             return;
+        }
         state.lastStreamedReasoning = formatted;
         void params.onReasoningStream({
             text: formatted,
@@ -411,6 +493,7 @@ export function subscribeEmbeddedPiSession(params) {
         flushBlockReplyBuffer,
         emitReasoningStream,
         consumeReplyDirectives,
+        consumePartialReplyDirectives,
         resetAssistantMessageState,
         resetForCompactionRetry,
         finalizeAssistantTexts,
@@ -419,6 +502,10 @@ export function subscribeEmbeddedPiSession(params) {
         noteCompactionRetry,
         resolveCompactionRetry,
         maybeResolveCompactionWait,
+        recordAssistantUsage,
+        incrementCompactionCount,
+        getUsageTotals,
+        getCompactionCount: () => compactionCount,
     };
     const unsubscribe = params.session.subscribe(createEmbeddedPiSessionEventHandler(ctx));
     return {
@@ -433,6 +520,8 @@ export function subscribeEmbeddedPiSession(params) {
         // which is generated AFTER the tool sends the actual answer.
         didSendViaMessagingTool: () => messagingToolSentTexts.length > 0,
         getLastToolError: () => (state.lastToolError ? { ...state.lastToolError } : undefined),
+        getUsageTotals,
+        getCompactionCount: () => compactionCount,
         waitForCompactionRetry: () => {
             if (state.compactionInFlight || state.pendingCompactionRetry > 0) {
                 ensureCompactionPromise();

package/dist/agents/pi-model-discovery.js

@@ -0,0 +1,10 @@
+import { AuthStorage, ModelRegistry } from "@mariozechner/pi-coding-agent";
+import path from "node:path";
+export { AuthStorage, ModelRegistry } from "@mariozechner/pi-coding-agent";
+// Compatibility helpers for pi-coding-agent 0.50+ (discover* helpers removed).
+export function discoverAuthStorage(agentDir) {
+    return new AuthStorage(path.join(agentDir, "auth.json"));
+}
+export function discoverModels(authStorage, agentDir) {
+    return new ModelRegistry(authStorage, path.join(agentDir, "models.json"));
+}