@animalabs/context-manager 0.5.1 → 0.5.2

package/src/context-manager.ts

@@ -25,6 +25,7 @@ import type { RenderStats } from './types/index.js';
 import { MessageStore, MessageStoreEvent } from './message-store.js';
 import { ContextLog } from './context-log.js';
 import { PassthroughStrategy } from './strategies/passthrough.js';
+import { splitMixedToolMessages } from './normalize-tool-messages.js';
 
 /**
  * Base configuration for ContextManager.
@@ -493,12 +494,30 @@ export class ContextManager {
       effectiveBudget
     );
 
-    // Convert to NormalizedMessage[]
-    const messages: NormalizedMessage[] = entries.map((entry) => ({
-      participant: entry.participant,
-      content: entry.content,
-      ...(entry.cacheMarker ? { cacheBreakpoint: true } : {}),
-    }));
+    // Convert to NormalizedMessage[]. We split each entry individually
+    // so we know the output-count per input and can re-attach cache
+    // markers to the last output of each (matching the marker's
+    // "cache up to here" intent).
+    //
+    // Splitting handles the claude.ai bundled-tool-cycle artifact: a
+    // non-user message containing `tool_result` blocks becomes a sequence
+    // of agent/user/agent turns so the API accepts it. Already-API-shape
+    // messages pass through untouched. See `src/normalize-tool-messages.ts`.
+    const messages: NormalizedMessage[] = [];
+    for (const entry of entries) {
+      const splitParts = splitMixedToolMessages([
+        { participant: entry.participant, content: entry.content },
+      ]);
+      for (let i = 0; i < splitParts.length; i++) {
+        const part = splitParts[i];
+        const isLast = i === splitParts.length - 1;
+        messages.push({
+          participant: part.participant,
+          content: part.content,
+          ...(entry.cacheMarker && isLast ? { cacheBreakpoint: true } : {}),
+        });
+      }
+    }
 
     // If no injections, log and return early
     if (!injections || injections.length === 0) {

package/src/index.ts

@@ -13,6 +13,9 @@ export { PassthroughStrategy } from './strategies/passthrough.js';
 export { AutobiographicalStrategy, type AutobiographicalProgressSnapshot } from './strategies/autobiographical.js';
 export { KnowledgeStrategy } from './strategies/knowledge.js';
 
+// Utilities
+export { splitMixedToolMessages } from './normalize-tool-messages.js';
+
 // Types
 export type {
   // Message types

package/src/message-store.ts

@@ -130,7 +130,6 @@ export class MessageStore {
     // Extract blobs from content
     const storedContent = this.blobManager.extractBlobs(content);
 
-    // Append to Chronicle first to get the record ID
     const partialInternal = {
       participant,
       content: storedContent,
@@ -140,45 +139,22 @@ export class MessageStore {
       ...(extra ?? {}),
     };
 
-    const record = this.store.appendToStateJson(this.stateId, partialInternal);
+    // Single atomic op: chronicle peeks the next record id+sequence under its
+    // write lock, splices them into the payload as `id` and `sequence`, and
+    // writes one record. The reconstructed state sees a fully-populated
+    // StoredMessageInternal, and `branchAt(messageId)` forks at this
+    // message's own sequence — exactly the post-fork-visible point.
+    const record = this.store.appendToStateJsonWithIdentity(
+      this.stateId,
+      partialInternal,
+      'id',
+      'sequence',
+    );
     const index = this.length() - 1;
 
-    // The append's payload doesn't carry the chronicle record's id /
-    // sequence — those are only known once the append returns. We patch
-    // them in via editStateItem below. The patch creates a new chronicle
-    // record at `record.sequence + 1` (the next sequence; nothing else
-    // writes between the append and the edit inside this function).
-    //
-    // We store the EDIT'S sequence (not the original append's) as the
-    // canonical "this message is fully readable at and after this seq"
-    // marker. Why: `ContextManager.branchAt(messageId)` forwards
-    // `message.sequence` to `chronicle.createBranchAt`, which forks
-    // visibility at that exact seq. If `sequence` pointed at the original
-    // append (record.sequence), the new branch would see the pre-patch
-    // payload (id/sequence undefined). Pointing at the edit's seq
-    // includes the patched payload in the inherited records — so a
-    // forked branch sees this message in its fully-established form.
-    //
-    // Invariant: this function does exactly ONE chronicle write between
-    // the append and the edit (the edit itself). If that ever changes,
-    // the +1 expression must be updated to match (or refactored to
-    // capture editStateItem's returned record.sequence — which is
-    // chicken-and-egg here since we need the value inside the payload
-    // we're writing).
-    const commitSequence = record.sequence + 1;
-    const fullInternal: StoredMessageInternal = {
-      id: record.id,
-      sequence: commitSequence,
-      ...partialInternal,
-    };
-    this.store.editStateItem(this.stateId, index, Buffer.from(JSON.stringify(fullInternal)));
-
-    // Build full message with ID and sequence from record. The sequence
-    // we expose is the commit sequence (matches what's stored), so
-    // `branchAt(messageId)` forks at the right place for any caller.
     const message: StoredMessage = {
       id: record.id,
-      sequence: commitSequence,
+      sequence: record.sequence,
       participant,
       content, // Original content with inline data
       metadata,
@@ -187,9 +163,7 @@ export class MessageStore {
       ...(extra ?? {}),
     };
 
-    // Update index
     this.idToIndex.set(message.id, index);
-
     this.emit({ type: 'add', message });
     return message;
   }

package/src/normalize-tool-messages.ts

@@ -0,0 +1,71 @@
+import type { ContentBlock } from '@animalabs/membrane';
+
+/**
+ * Anthropic-API shape requires `tool_use` blocks in assistant turns and
+ * `tool_result` blocks in user turns, with each tool_result in the message
+ * *immediately following* its tool_use. claude.ai exports bundle the entire
+ * tool cycle into a single assistant message, and any downstream consumer
+ * that ferries those messages through unchanged hits a 400.
+ *
+ * This helper walks a list of messages and, for every non-user message
+ * containing one or more `tool_result` blocks, splits it so:
+ *
+ *   - blocks before the first tool_result stay under the original participant
+ *   - each run of consecutive tool_results becomes its own `user` message
+ *   - blocks after the tool_results go back to the original participant
+ *
+ * Messages that are already user-side or that don't contain tool_results
+ * pass through unchanged. The participant-name check is case-insensitive
+ * (matching the convention used elsewhere in CM) so it works regardless
+ * of whether the session uses 'user', 'User', or some custom label.
+ *
+ * Called by:
+ *   - `ContextManager.compile()` — live session path, prevents the
+ *     bundled-cycle 400 for already-imported sessions.
+ *   - `AutobiographicalStrategy.compressChunkHierarchical()` — defense in
+ *     depth at the compression LLM call site.
+ *   - `AutobiographicalStrategy.executeMerge()` — same, at the merge call site.
+ *
+ * For new imports, conhost's `scripts/import-claudeai-export.ts` already
+ * splits at ingest time so Chronicle stores hold API-shape messages. This
+ * runtime pass handles legacy sessions imported before that fix landed,
+ * and acts as a safety net for any future bundling source.
+ */
+export function splitMixedToolMessages<T extends { participant: string; content: ContentBlock[] }>(
+  messages: readonly T[],
+): Array<{ participant: string; content: ContentBlock[] }> {
+  const out: Array<{ participant: string; content: ContentBlock[] }> = [];
+  for (const msg of messages) {
+    const isUser = msg.participant.toLowerCase() === 'user';
+    if (isUser || !msg.content.some((b) => b.type === 'tool_result')) {
+      out.push({ participant: msg.participant, content: msg.content });
+      continue;
+    }
+    let preTool: ContentBlock[] = [];
+    let pendingResults: ContentBlock[] = [];
+    const flushPre = () => {
+      if (preTool.length > 0) {
+        out.push({ participant: msg.participant, content: preTool });
+        preTool = [];
+      }
+    };
+    const flushResults = () => {
+      if (pendingResults.length > 0) {
+        out.push({ participant: 'user', content: pendingResults });
+        pendingResults = [];
+      }
+    };
+    for (const block of msg.content) {
+      if (block.type === 'tool_result') {
+        flushPre();
+        pendingResults.push(block);
+      } else {
+        flushResults();
+        preTool.push(block);
+      }
+    }
+    flushResults();
+    flushPre();
+  }
+  return out;
+}

package/src/strategies/autobiographical.ts

@@ -20,6 +20,7 @@ import type {
 } from '../types/index.js';
 import { DEFAULT_AUTOBIOGRAPHICAL_CONFIG } from '../types/index.js';
 import { getSummaryParentId } from '../types/strategy.js';
+import { splitMixedToolMessages } from '../normalize-tool-messages.js';
 import { appendFileSync, mkdirSync } from 'node:fs';
 import { dirname } from 'node:path';
 import { Picker, OverBudgetError, type PickerChunk } from '../adaptive/picker.js';
@@ -747,6 +748,39 @@ export class AutobiographicalStrategy implements ResettableStrategy {
    * Mark a summary as merged into a higher-level summary, updating the
    * chronicle copy at the same index. Index is the position in `this.summaries`.
    */
+  /**
+   * Token-budget cap for recall-pair summary sets. Walks newest→oldest
+   * keeping each summary that still fits; skips (rather than breaks at)
+   * a summary that would put us over budget, so a heterogeneous set
+   * fills the remaining slots with smaller siblings instead of stopping
+   * at the first oversized one. The kept set is re-sorted chronologically.
+   *
+   * Used by both `compressChunkHierarchical` (the L1 compression prompt
+   * recall pairs) and `executeMerge` (the merge prompt recall pairs).
+   * Without the cap, both sites grow their recall set linearly with
+   * conversation length and overflow the 200k window around the same
+   * point — observed empirically at ~chunk 118 in a 4000-message import.
+   *
+   * Per-summary +50 token overhead accounts for the "[CM] Recall memory
+   * <id>." question turn that wraps each recall body. Rough but defensive.
+   */
+  protected capRecallPairs(
+    summariesChronological: SummaryEntry[],
+    maxTokens: number,
+  ): { kept: SummaryEntry[]; keptTokens: number } {
+    const kept: SummaryEntry[] = [];
+    let total = 0;
+    for (let i = summariesChronological.length - 1; i >= 0; i--) {
+      const s = summariesChronological[i]!;
+      const est = (s.tokens ?? Math.ceil(s.content.length / 4)) + 50;
+      if (total + est > maxTokens) continue;
+      kept.push(s);
+      total += est;
+    }
+    kept.reverse();
+    return { kept, keptTokens: total };
+  }
+
   protected setMergedInto(entry: SummaryEntry, mergedIntoId: string): void {
     entry.mergedInto = mergedIntoId;
     if (!this.store) return;
@@ -1537,33 +1571,22 @@ export class AutobiographicalStrategy implements ResettableStrategy {
       .filter((s) => !s.mergedInto)
       .sort((a, b) => a.sourceRange.first.localeCompare(b.sourceRange.first));
 
-    // Token-budget cap on recall pairs as defense-in-depth. With merged
-    // exclusion the set is bounded by `mergeThreshold^depth`, but a
-    // conversation with thousands of unmerged top-level summaries can
-    // still overflow. Newest-first selection so proximate context is
-    // preferred; the kept set is re-sorted chronologically below. We
-    // always keep at least one summary even if it's larger than the
-    // budget — otherwise a single oversized summary would leave the
-    // model with no compressed context at all.
-    const recallBudget = this.config.compressionRecallBudgetTokens ?? 150_000;
-    const keptSummaries: SummaryEntry[] = [];
-    let recallTokens = 0;
-    for (let i = priorSummaries.length - 1; i >= 0; i--) {
-      const s = priorSummaries[i]!;
-      const approxTokens = Math.ceil(s.content.length / 4);
-      if (recallTokens + approxTokens > recallBudget && keptSummaries.length > 0) break;
-      keptSummaries.push(s);
-      recallTokens += approxTokens;
-    }
-    keptSummaries.reverse();
+    // Token-budget cap (see capRecallPairs). Defense-in-depth: even with
+    // merged exclusion the unmerged frontier can be large at extreme scale.
+    const recallBudget = this.config.compressionRecallBudgetTokens ?? 100_000;
+    const { kept: keptSummaries, keptTokens: recallTokens } = this.capRecallPairs(
+      priorSummaries,
+      recallBudget,
+    );
     if (keptSummaries.length < priorSummaries.length) {
       const dropped = priorSummaries.length - keptSummaries.length;
       console.warn(
-        `autobio: recall-pair budget capped (${keptSummaries.length}/${priorSummaries.length} summaries kept, ` +
+        `autobio: compression recall-pair budget capped (${keptSummaries.length}/${priorSummaries.length} summaries kept, ` +
           `~${recallTokens} tokens, budget ${recallBudget}; ${dropped} oldest dropped this compression).`,
       );
       logCompressionCall({
         event: 'recall-budget-capped',
+        site: 'compression',
         kept: keptSummaries.length,
         total: priorSummaries.length,
         tokens: recallTokens,
@@ -1660,8 +1683,15 @@ export class AutobiographicalStrategy implements ResettableStrategy {
       content: [{ type: 'text', text: instructionText }],
     });
 
+    // Split any bundled tool_use+tool_result cycles in non-user turns into
+    // separate API-shape messages. claude.ai-imported sessions carry these
+    // bundles (a tool_result in an assistant message rejects the request);
+    // for fresh imports the conhost importer splits at ingest time, but
+    // already-warmed sessions hit this path. See `normalize-tool-messages.ts`.
+    const split = splitMixedToolMessages(llmMessages);
+
     // Collapse consecutive same-participant messages for API compliance
-    const collapsed = this.collapseConsecutiveMessages(llmMessages);
+    const collapsed = this.collapseConsecutiveMessages(split);
 
     // NO system prompt. The agent's identity is established by the head
     // (the actual conversation opening — user message + agent reply that
@@ -1955,15 +1985,19 @@ export class AutobiographicalStrategy implements ResettableStrategy {
       llmMessages.push({ participant: m.participant, content: m.content });
     }
 
-    // ---- 1b. PRIOR L1 RECALL PAIRS (chronologically before merge range) ----
-    // L1s whose entire source range is before the merge range and that
-    // aren't part of the merge tree. Sort by source position.
-    const priorL1s = this.summaries
-      .filter((s) => s.level === 1)
+    // ---- 1b. PRIOR RECALL PAIRS (chronologically before merge range) ----
+    // The unmerged frontier of summaries whose source range is before the
+    // merge range and which aren't part of the merge tree. Originally this
+    // was filtered to `level === 1` (the "L1 fidelity for prior content"
+    // intent) but at 4000+ messages that produces hundreds of L1s and
+    // overflows the model window. Switching to the unmerged frontier
+    // (`!mergedInto`) lets a merged L1 drop out in favour of its L2/L3
+    // parent — the same rule used everywhere else and now in
+    // `compressChunkHierarchical`. The cap below is the defense-in-depth.
+    const priorSummariesAll = this.summaries
+      .filter((s) => !s.mergedInto)
       .filter((s) => {
-        // Exclude any L1 that's an ancestor of our merge target
         for (const lid of s.sourceIds) if (sourceLeafIds.has(lid)) return false;
-        // Include only if it starts strictly before the merge range
         const firstIdx = allMessages.findIndex((m) => m.id === s.sourceRange.first);
         return firstIdx >= 0 && (mergeStartIdx < 0 || firstIdx < mergeStartIdx);
       })
@@ -1973,12 +2007,37 @@ export class AutobiographicalStrategy implements ResettableStrategy {
         return ai - bi;
       });
 
-    const priorL1MessageIds = new Set<MessageId>();
-    for (const s of priorL1s) {
-      for (const id of s.sourceIds) priorL1MessageIds.add(id);
+    const mergeRecallBudget = this.config.compressionRecallBudgetTokens ?? 100_000;
+    const { kept: keptPriorSummaries, keptTokens: mergeRecallTokens } = this.capRecallPairs(
+      priorSummariesAll,
+      mergeRecallBudget,
+    );
+    if (keptPriorSummaries.length < priorSummariesAll.length) {
+      const dropped = priorSummariesAll.length - keptPriorSummaries.length;
+      console.warn(
+        `autobio: merge recall-pair budget capped (${keptPriorSummaries.length}/${priorSummariesAll.length} summaries kept, ` +
+          `~${mergeRecallTokens} tokens, budget ${mergeRecallBudget}; ${dropped} oldest dropped this merge).`,
+      );
+      logCompressionCall({
+        event: 'recall-budget-capped',
+        site: 'merge',
+        targetLevel,
+        kept: keptPriorSummaries.length,
+        total: priorSummariesAll.length,
+        tokens: mergeRecallTokens,
+        budgetTokens: mergeRecallBudget,
+      });
+    }
+
+    // The full unmerged-frontier set covers what's "compressed somewhere"
+    // for the raw-middle dedup below — a budget-dropped recall pair
+    // doesn't make its underlying raw messages reappear.
+    const priorSummaryMessageIds = new Set<MessageId>();
+    for (const s of priorSummariesAll) {
+      for (const id of s.sourceIds) priorSummaryMessageIds.add(id);
     }
 
-    for (const s of priorL1s) {
+    for (const s of keptPriorSummaries) {
       llmMessages.push({
         participant: 'Context Manager',
         content: [{ type: 'text', text: `[CM] Recall memory ${s.id}.` }],
@@ -1990,12 +2049,12 @@ export class AutobiographicalStrategy implements ResettableStrategy {
     }
 
     // Raw middle: any messages between the head window and the merge
-    // range that aren't covered by a prior L1 or the merge tree.
+    // range that aren't covered by a prior summary or the merge tree.
     // Usually empty (chunking is contiguous).
     if (mergeStartIdx >= 0) {
       for (let i = headEndIdx; i < mergeStartIdx; i++) {
         const m = allMessages[i];
-        if (priorL1MessageIds.has(m.id)) continue;
+        if (priorSummaryMessageIds.has(m.id)) continue;
         if (sourceLeafIds.has(m.id)) continue;
         llmMessages.push({ participant: m.participant, content: m.content });
       }
@@ -2081,7 +2140,9 @@ export class AutobiographicalStrategy implements ResettableStrategy {
       }],
     });
 
-    const collapsed = this.collapseConsecutiveMessages(llmMessages);
+    // Same bundled-tool-cycle defense as compressChunkHierarchical.
+    const split = splitMixedToolMessages(llmMessages);
+    const collapsed = this.collapseConsecutiveMessages(split);
 
     // NO system prompt — identity is established by the head window
     // (present at the start of llmMessages above) and by the prior

package/src/types/strategy.ts

@@ -339,14 +339,25 @@ export interface AutobiographicalConfig {
 
   /**
    * Cap on the total tokens of recall-pair prior-summary content
-   * included in each chunk compression request (default: 150000).
+   * included in each LLM request that builds recall pairs:
+   *
+   *   - L1 chunk compression (`compressChunkHierarchical`)
+   *   - L_n merges (`executeMerge`)
    *
    * Defends against the case where the unmerged frontier itself is
-   * large enough to overflow the API window. The cap walks summaries
-   * newest-first so proximate context survives; the kept set is then
-   * re-sorted chronologically. Set higher (or to Infinity) if you want
-   * to surface overflow via a 400 rather than silently drop oldest
-   * memories from the compression context.
+   * large enough to overflow the API window. Walks newest-first so
+   * proximate context survives; the kept set is then re-sorted
+   * chronologically. Each summary takes (its content tokens + 50 for
+   * the wrapping "[CM] Recall memory <id>." question).
+   *
+   * Default 100000 — chosen so that even an L_n merge (which packs
+   * recall + head + expanded target + instruction into one request)
+   * fits inside a 200k window. For larger context models or
+   * shallower hierarchies, raise this; on Sonnet/Opus default windows,
+   * the practical ceiling is roughly 130k.
+   *
+   * Set to Infinity to disable the cap and surface overflow as a
+   * 400 from the API rather than silently dropping oldest memories.
    */
   compressionRecallBudgetTokens?: number;