@animalabs/context-manager 0.3.0 → 0.5.0

package/src/context-manager.ts

@@ -15,8 +15,13 @@ import type {
   MessageQueryResult,
   ContextInjection,
   CompileResult,
+  ProtectedRange,
+  SearchQuery,
+  SearchResult,
+  SummaryEntry,
 } from './types/index.js';
-import { isResettableStrategy } from './types/index.js';
+import { isResettableStrategy, isPinnableStrategy, isSearchableStrategy, isRenderStatsCapable } from './types/index.js';
+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';
@@ -93,6 +98,8 @@ export class ContextManager {
   /** Whether we own the store (created it) vs app owns it (passed in) */
   private ownsStore: boolean;
   private debugLogContext: boolean;
+  /** Namespace passed to strategies for scoping their persistent state slots. */
+  private strategyNamespace: string;
 
   private constructor(
     store: JsStore,
@@ -100,6 +107,7 @@ export class ContextManager {
     contextLog: ContextLog,
     strategy: ContextStrategy,
     ownsStore: boolean,
+    strategyNamespace: string,
     membrane?: Membrane,
     debugLogContext = false,
   ) {
@@ -108,6 +116,7 @@ export class ContextManager {
     this.contextLog = contextLog;
     this.strategy = strategy;
     this.ownsStore = ownsStore;
+    this.strategyNamespace = strategyNamespace;
     this.membrane = membrane;
     this.debugLogContext = debugLogContext;
 
@@ -173,12 +182,18 @@ export class ContextManager {
     });
     const strategy = config.strategy ?? new PassthroughStrategy();
 
+    // Namespace passed to strategies. Falls back to a stable per-store value
+    // so strategies always have something to scope state IDs by, even when
+    // the caller didn't supply a namespace.
+    const strategyNamespace = config.namespace ?? 'default';
+
     const manager = new ContextManager(
       store,
       messageStore,
       contextLog,
       strategy,
       ownsStore,
+      strategyNamespace,
       config.membrane,
       config.debugLogContext ?? false,
     );
@@ -196,6 +211,11 @@ export class ContextManager {
 
   /**
    * Add a message to the store.
+   *
+   * If the configured strategy implements `chunkIngressMessage` and returns
+   * a non-null sharding decision, the message is stored as multiple records
+   * sharing a `bodyGroupId` (per the adaptive-resolution design §3.6). The
+   * returned MessageId is the first shard's id.
    */
   addMessage(
     participant: string,
@@ -203,6 +223,33 @@ export class ContextManager {
     metadata?: MessageMetadata,
     causedBy?: MessageId[]
   ): MessageId {
+    // Optional strategy-driven ingestion-time chunking
+    const strategyAny = this.strategy as unknown as {
+      chunkIngressMessage?: (
+        participant: string,
+        content: ContentBlock[]
+      ) => { bodyGroupId: string; shards: Array<{ content: ContentBlock[]; shardIndex: number }> } | null;
+    };
+    if (typeof strategyAny.chunkIngressMessage === 'function') {
+      const decision = strategyAny.chunkIngressMessage(participant, content);
+      if (decision && decision.shards.length > 1) {
+        let firstId: MessageId | null = null;
+        for (const shard of decision.shards) {
+          const message = this.messageStore.append(
+            participant,
+            shard.content,
+            metadata,
+            causedBy,
+            {
+              bodyGroupId: decision.bodyGroupId,
+              shardIndex: shard.shardIndex,
+            }
+          );
+          if (firstId === null) firstId = message.id;
+        }
+        return firstId!;
+      }
+    }
     const message = this.messageStore.append(participant, content, metadata, causedBy);
     return message.id;
   }
@@ -294,6 +341,10 @@ export class ContextManager {
   /**
    * Create a branch from a specific message.
    * The new branch will have state as of that message's sequence (time-travel branching).
+   *
+   * Returns the new branch *name*, which is what `switchBranch` and `forkAt`
+   * expect. (Chronicle's branch APIs are name-keyed; the numeric `id` field
+   * on JsBranch is an internal identifier and isn't accepted by switchBranch.)
    */
   branchAt(messageId: MessageId, name?: string): string {
     const message = this.messageStore.get(messageId);
@@ -303,21 +354,50 @@ export class ContextManager {
 
     // Create branch name if not provided
     const branchName = name ?? `branch-${Date.now()}`;
-    
+
     // Get current branch name to branch from
     const currentBranch = this.store.currentBranch();
-    
+
     // Use createBranchAt to branch at the message's sequence (time-travel)
     const branch = this.store.createBranchAt(branchName, currentBranch.name, message.sequence);
 
-    return branch.id;
+    return branch.name;
   }
 
   /**
    * Switch to a different branch.
+   *
+   * Re-initializes the strategy after switching so any branch-scoped state
+   * stored on Chronicle is reloaded. Strategies that hold derived in-memory
+   * caches (e.g. AutobiographicalStrategy.summaries) need this to avoid
+   * showing the previous branch's state on the new branch.
    */
-  switchBranch(branchId: string): void {
+  async switchBranch(branchId: string): Promise<void> {
     this.store.switchBranch(branchId);
+    await this.initializeStrategy();
+  }
+
+  /**
+   * Fork from the current head: create a new branch at the current sequence
+   * and switch to it. The new branch starts with all current state (messages,
+   * context log, and strategy state) and diverges from there.
+   *
+   * Use this when an agent wants to explore an alternate timeline from
+   * "now" — e.g. trying a different response without committing.
+   *
+   * For time-travel branching at a specific historical message, use
+   * `branchAt(messageId, name?)` instead, then `switchBranch(name)`.
+   *
+   * Returns the new branch's name. The strategy is re-initialized on the
+   * new branch so it picks up the forked state.
+   */
+  async fork(name?: string): Promise<string> {
+    const branchName = name ?? `fork-${Date.now()}`;
+    const currentBranch = this.store.currentBranch();
+    const currentSeq = this.store.currentSequence();
+    const branch = this.store.createBranchAt(branchName, currentBranch.name, currentSeq);
+    await this.switchBranch(branch.name);
+    return branch.name;
   }
 
   /**
@@ -390,11 +470,15 @@ export class ContextManager {
     budget?: TokenBudget,
     injections?: ContextInjection[]
   ): Promise<CompileResult> {
-    // Check readiness and wait if needed
-    const readiness = this.strategy.checkReadiness();
-    if (!readiness.ready && readiness.pendingWork) {
-      await readiness.pendingWork;
-    }
+    // Don't block the agent's turn on speculative compression — let it
+    // run in the background. The strategy renders whatever's available
+    // now; the next compile picks up the freshly-formed L1.
+    //
+    // Old behavior (await pendingWork to fold the latest chunk before
+    // the agent responds) added 30+ seconds of latency per turn whenever
+    // a chunk was forming, which is unacceptable UX for an agent that
+    // streams its responses. We accept "this turn doesn't have the very
+    // latest L1" in exchange for non-blocking compile.
 
     // Default budget
     const effectiveBudget: TokenBudget = budget ?? {
@@ -527,6 +611,84 @@ export class ContextManager {
     return this.strategy;
   }
 
+  // ==========================================================================
+  // Pins / documents (passthrough to the active strategy)
+  // ==========================================================================
+
+  /**
+   * Pin a range of messages so they aren't compressed and render raw at
+   * their original chronological position. Returns the new pin id.
+   *
+   * Throws if the active strategy doesn't support pins.
+   */
+  pinRange(firstMessageId: MessageId, lastMessageId: MessageId, opts?: { name?: string }): string {
+    if (!isPinnableStrategy(this.strategy)) {
+      throw new Error('Active strategy does not support pins');
+    }
+    return this.strategy.pinRange(firstMessageId, lastMessageId, opts);
+  }
+
+  /**
+   * Mark a single message as a "document" (semantically a body of
+   * information to retain in full). Same effect as a single-message pin
+   * with `kind: 'document'`. Returns the new pin id.
+   */
+  markDocument(messageId: MessageId, opts?: { name?: string }): string {
+    if (!isPinnableStrategy(this.strategy)) {
+      throw new Error('Active strategy does not support documents');
+    }
+    return this.strategy.markDocument(messageId, opts);
+  }
+
+  /** Remove a pin or document mark. Returns true if removed. */
+  unpin(pinId: string): boolean {
+    if (!isPinnableStrategy(this.strategy)) {
+      throw new Error('Active strategy does not support pins');
+    }
+    return this.strategy.unpin(pinId);
+  }
+
+  /** List all current pins. Returns empty array if strategy is not pinnable. */
+  listPins(): ReadonlyArray<ProtectedRange> {
+    if (!isPinnableStrategy(this.strategy)) return [];
+    return this.strategy.listPins();
+  }
+
+  // ==========================================================================
+  // Search (passthrough to the active strategy)
+  // ==========================================================================
+
+  /**
+   * Search the strategy's summary archive (substring or regex over content).
+   * Returns empty array if the strategy doesn't support search.
+   *
+   * Suitable for building memory-search agent tools at the framework layer
+   * — see e.g. agent-framework's MCPL host integration.
+   */
+  searchSummaries(query: SearchQuery): SearchResult[] {
+    if (!isSearchableStrategy(this.strategy)) return [];
+    return this.strategy.searchSummaries(query);
+  }
+
+  /** Look up a single summary by id. Returns null if not found / unsupported. */
+  getSummary(id: string): SummaryEntry | null {
+    if (!isSearchableStrategy(this.strategy)) return null;
+    return this.strategy.getSummary(id);
+  }
+
+  /**
+   * Per-render stats from the active strategy: head/tail message + token
+   * counts, plus per-level summary counts and total tokens. Returns null
+   * if the strategy doesn't implement `getRenderStats`.
+   *
+   * Designed for TUIs / dashboards that want to display "how much of the
+   * agent's context is folded vs raw" at a glance.
+   */
+  getRenderStats(): RenderStats | null {
+    if (!isRenderStatsCapable(this.strategy)) return null;
+    return this.strategy.getRenderStats(this.messageStore.createView());
+  }
+
   /**
    * Reset the head window to start from a new position.
    * Old head window messages become compressible.
@@ -583,6 +745,8 @@ export class ContextManager {
       contextLog: this.contextLog.createView(),
       membrane: this.membrane,
       currentSequence: this.store.currentSequence(),
+      store: this.store,
+      namespace: this.strategyNamespace,
     };
   }
 

package/src/index.ts

@@ -10,7 +10,7 @@ export { BlobManager } from './blob-manager.js';
 
 // Strategies
 export { PassthroughStrategy } from './strategies/passthrough.js';
-export { AutobiographicalStrategy } from './strategies/autobiographical.js';
+export { AutobiographicalStrategy, type AutobiographicalProgressSnapshot } from './strategies/autobiographical.js';
 export { KnowledgeStrategy } from './strategies/knowledge.js';
 
 // Types

package/src/message-store.ts

@@ -101,14 +101,31 @@ export class MessageStore {
     }
   }
 
+  /**
+   * Optional extra fields for `append`, used by callers that need to set
+   * adaptive-resolution metadata (bodyGroupId for shards, initial
+   * resolution state, etc.) at ingestion time.
+   */
+  static readonly _appendExtraKeys = ['bodyGroupId', 'shardIndex', 'currentResolution', 'lockedByAgent'] as const;
+
   /**
    * Append a new message to the store.
+   *
+   * `extra` is an optional bag of adaptive-resolution metadata
+   * (bodyGroupId / shardIndex / currentResolution / lockedByAgent) that
+   * callers may set at ingestion. Field semantics match StoredMessage.
    */
   append(
     participant: string,
     content: ContentBlock[],
     metadata?: MessageMetadata,
-    causedBy?: MessageId[]
+    causedBy?: MessageId[],
+    extra?: {
+      bodyGroupId?: string;
+      shardIndex?: number;
+      currentResolution?: number;
+      lockedByAgent?: boolean;
+    }
   ): StoredMessage {
     // Extract blobs from content
     const storedContent = this.blobManager.extractBlobs(content);
@@ -120,28 +137,54 @@ export class MessageStore {
       metadata,
       timestamp: Date.now(),
       causedBy,
+      ...(extra ?? {}),
     };
 
     const record = this.store.appendToStateJson(this.stateId, partialInternal);
     const index = this.length() - 1;
 
-    // Now update the stored item to include the id (for rebuildIndex)
+    // 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: record.sequence,
+      sequence: commitSequence,
       ...partialInternal,
     };
     this.store.editStateItem(this.stateId, index, Buffer.from(JSON.stringify(fullInternal)));
 
-    // Build full message with ID and sequence from record
+    // 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: record.sequence,
+      sequence: commitSequence,
       participant,
       content, // Original content with inline data
       metadata,
       timestamp: new Date(partialInternal.timestamp),
       causedBy,
+      ...(extra ?? {}),
     };
 
     // Update index
@@ -153,6 +196,11 @@ export class MessageStore {
 
   /**
    * Edit a message's content.
+   *
+   * Throws if `messageId` belongs to a bodyGroup (i.e., is a shard of a
+   * larger sharded message). Editing one shard would silently corrupt the
+   * bodyGroup's byte-faithful reassembly invariant. To replace a sharded
+   * message, remove the whole bodyGroup and re-append.
    */
   edit(messageId: MessageId, newContent: ContentBlock[]): void {
     const index = this.idToIndex.get(messageId);
@@ -165,6 +213,13 @@ export class MessageStore {
       throw new Error(`Message not found at index: ${index}`);
     }
 
+    if (oldMessage.bodyGroupId) {
+      throw new Error(
+        `Cannot edit shard ${messageId}: it is part of bodyGroup ${oldMessage.bodyGroupId}. ` +
+          `Sharded messages are immutable — remove the whole bodyGroup and re-append instead.`,
+      );
+    }
+
     const oldContent = this.blobManager.resolveBlobs(oldMessage.content);
     const storedContent = this.blobManager.extractBlobs(newContent);
 
@@ -181,6 +236,11 @@ export class MessageStore {
 
   /**
    * Remove a message from the store.
+   *
+   * For a sharded message (one shard of a bodyGroup), the caller MUST
+   * remove all shards in the group together — removing one shard would
+   * orphan the rest and break byte-faithful reassembly. Use
+   * `removeBodyGroup(id)` for that case.
    */
   remove(messageId: MessageId): void {
     const index = this.idToIndex.get(messageId);
@@ -188,14 +248,61 @@ export class MessageStore {
       throw new Error(`Message not found: ${messageId}`);
     }
 
+    const target = this.getInternal(index);
+    if (target?.bodyGroupId) {
+      throw new Error(
+        `Cannot remove shard ${messageId} in isolation: it is part of bodyGroup ${target.bodyGroupId}. ` +
+          `Use removeBodyGroup(${messageId}) to remove all shards atomically.`,
+      );
+    }
+
     this.store.redactStateItems(this.stateId, index, index + 1);
     this.rebuildIndex();
 
     this.emit({ type: 'remove', messageId });
   }
 
+  /**
+   * Remove every shard of the bodyGroup containing `messageId`. If the
+   * message is not part of a bodyGroup, falls back to single-message
+   * remove. Atomic from the caller's perspective.
+   */
+  removeBodyGroup(messageId: MessageId): void {
+    const index = this.idToIndex.get(messageId);
+    if (index === undefined) {
+      throw new Error(`Message not found: ${messageId}`);
+    }
+    const target = this.getInternal(index);
+    if (!target?.bodyGroupId) {
+      // Not sharded — defer to normal remove path (re-look up via getInternal
+      // since the normal `remove` checks bodyGroupId).
+      this.store.redactStateItems(this.stateId, index, index + 1);
+      this.rebuildIndex();
+      this.emit({ type: 'remove', messageId });
+      return;
+    }
+    // Find the contiguous run of shards with this bodyGroupId. Shards are
+    // stored consecutively (they're appended one after the other at
+    // ingestion), so we can scan outward from `index`.
+    const groupId = target.bodyGroupId;
+    const all = this.getAllInternal();
+    let from = index;
+    while (from > 0 && all[from - 1].bodyGroupId === groupId) from--;
+    let to = index;
+    while (to + 1 < all.length && all[to + 1].bodyGroupId === groupId) to++;
+    const firstId = all[from].id;
+    const lastId = all[to].id;
+    this.store.redactStateItems(this.stateId, from, to + 1);
+    this.rebuildIndex();
+    this.emit({ type: 'removeRange', fromId: firstId, toId: lastId });
+  }
+
   /**
    * Remove a range of messages from the store.
+   *
+   * If the range starts or ends in the middle of a bodyGroup, throws —
+   * removeRange must align to bodyGroup boundaries. (Use `removeBodyGroup`
+   * to remove an entire group, then call `removeRange` over plain messages.)
    */
   removeRange(fromId: MessageId, toId: MessageId): void {
     const fromIndex = this.idToIndex.get(fromId);
@@ -208,6 +315,21 @@ export class MessageStore {
       throw new Error(`Message not found: ${toId}`);
     }
 
+    // Verify the range doesn't bisect a bodyGroup.
+    const all = this.getAllInternal();
+    const startGroup = all[fromIndex].bodyGroupId;
+    const endGroup = all[toIndex].bodyGroupId;
+    if (startGroup && fromIndex > 0 && all[fromIndex - 1].bodyGroupId === startGroup) {
+      throw new Error(
+        `removeRange would bisect bodyGroup ${startGroup} at start. Use removeBodyGroup(${fromId}) first.`,
+      );
+    }
+    if (endGroup && toIndex + 1 < all.length && all[toIndex + 1].bodyGroupId === endGroup) {
+      throw new Error(
+        `removeRange would bisect bodyGroup ${endGroup} at end. Use removeBodyGroup(${toId}) first.`,
+      );
+    }
+
     this.store.redactStateItems(this.stateId, fromIndex, toIndex + 1);
     this.rebuildIndex();
 
@@ -457,17 +579,34 @@ export class MessageStore {
   private internalToStored(
     internal: StoredMessageInternal,
     id: MessageId,
-    index: number
+    _index: number,
   ): StoredMessage {
-    return {
+    const stored: StoredMessage = {
       id,
-      sequence: index, // Use index as sequence for now
+      // chronicle record sequence captured when the message was appended
+      // (see `add()` line 131: `sequence: record.sequence`). The previous
+      // implementation returned the slot index ("// Use index as sequence
+      // for now"), which silently corrupted any downstream code that
+      // forwarded this number to chronicle APIs expecting a real sequence
+      // — most notably `ContextManager.branchAt`, which would fork the
+      // chronicle at the index-mistaken-for-sequence and lose every
+      // record between the intended fork point and the actual one. For
+      // typical autobio sessions the index-vs-sequence ratio is ~1:2
+      // (each message append is accompanied by ~1 autobio state update),
+      // so /undo of a 800-message conversation forked ~400 messages back.
+      sequence: internal.sequence,
       participant: internal.participant,
       content: this.blobManager.resolveBlobs(internal.content),
       metadata: internal.metadata,
       timestamp: new Date(internal.timestamp),
       causedBy: internal.causedBy,
     };
+    // Carry adaptive-resolution fields through unchanged.
+    if (internal.bodyGroupId !== undefined) stored.bodyGroupId = internal.bodyGroupId;
+    if (internal.shardIndex !== undefined) stored.shardIndex = internal.shardIndex;
+    if (internal.currentResolution !== undefined) stored.currentResolution = internal.currentResolution;
+    if (internal.lockedByAgent !== undefined) stored.lockedByAgent = internal.lockedByAgent;
+    return stored;
   }
 }