feat: Add context-aware memory enhancements

Adds provider-agnostic enhancements to memory system to support advanced features like token budgeting, message editing, caching, and priority-based selection.

## New Features

### Enhanced Message Metadata (Optional)
- status: 'active' | 'archived' | 'deleted' for filtering
- importance: 0-1 score for priority-based selection
- tokenCount: Provider-calculated token count
- cacheControl: Cache configuration with expiry
- version: Message versioning for edit tracking
- editedFrom: Original message ID for edit history

### Context Window Management
- ContextWindowConfig: maxTokens, strategy, retainRecentMessages
- Strategies: fifo, priority-based, token-aware
- Token budget-aware message retrieval

### New Storage Methods (Optional)
- updateMessage: Edit existing messages
- getMessageHistory: Get message version history
- getThreadTokenCount: Calculate total tokens
- getMessagesByTokenBudget: Retrieve within token limit
- getCachedMessages: Get messages with active cache
- invalidateCache: Expire cached messages
- getMessagesByStatus: Filter by status
- getMessagesByImportance: Filter by importance score

### Memory Class Enhancements
- updateMessage: Edit messages with versioning
- getMessageVersions: Get edit history
- getThreadTokenCount: Get token usage
- getMessagesWithinBudget: Token-aware retrieval
- getCachedMessages: Cache-aware retrieval
- invalidateCache: Cache management
- getMessagesByStatus: Status filtering
- getMessagesByImportance: Priority filtering

## Design Principles
- Backward compatible: All new fields/methods are optional
- Provider agnostic: No specific API dependencies
- Graceful degradation: Falls back when storage doesn't support features
- Type safe: Full TypeScript support with proper types

## Implementation
- All optional methods implemented in InMemoryStorage
- Memory class checks for method availability before calling
- New types exported: CacheControl, ContextWindowConfig
- All existing tests pass (10/10)

Author: mattapperson

src/index.ts

@@ -13,17 +13,19 @@ export { ReusableReadableStream } from "./lib/reusable-stream.js";
 // Memory system exports
 export { Memory, InMemoryStorage } from "./lib/memory/index.js";
 export type {
-  MemoryStorage,
+  CacheControl,
+  ContextWindowConfig,
+  GetMessagesOptions,
   MemoryConfig,
-  Thread,
-  Resource,
   MemoryMessage,
-  ThreadWorkingMemory,
+  MemoryStorage,
+  Resource,
   ResourceWorkingMemory,
-  WorkingMemoryData,
   SerializedMemoryState,
   SerializedThreadState,
-  GetMessagesOptions,
+  Thread,
+  ThreadWorkingMemory,
+  WorkingMemoryData,
 } from "./lib/memory/index.js";
 // #endregion
 export * from "./sdk/sdk.js";

src/lib/memory/index.ts

@@ -12,6 +12,8 @@ export type { MemoryStorage } from "./storage/interface.js";
 
 // Types
 export type {
+  CacheControl,
+  ContextWindowConfig,
   GetMessagesOptions,
   MemoryConfig,
   MemoryMessage,

src/lib/memory/memory.ts

@@ -16,12 +16,17 @@ import type {
   WorkingMemoryData,
 } from "./types.js";
 
+/**
+ * Resolved configuration with all defaults applied
+ */
+type ResolvedMemoryConfig = Required<Pick<MemoryConfig, 'maxHistoryMessages' | 'autoInject' | 'autoSave' | 'trackTokenUsage'>> & Pick<MemoryConfig, 'contextWindow'>;
+
 /**
  * Memory class for managing conversation history, threads, and working memory
  */
 export class Memory {
   private storage: MemoryStorage;
-  private config: Required<MemoryConfig>;
+  private config: ResolvedMemoryConfig;
 
   /**
    * Create a new Memory instance
@@ -34,13 +39,15 @@ export class Memory {
       maxHistoryMessages: config.maxHistoryMessages ?? 10,
       autoInject: config.autoInject ?? true,
       autoSave: config.autoSave ?? true,
+      ...(config.contextWindow !== undefined && { contextWindow: config.contextWindow }),
+      trackTokenUsage: config.trackTokenUsage ?? false,
     };
   }
 
   /**
    * Get the current memory configuration
    */
-  getConfig(): Required<MemoryConfig> {
+  getConfig(): ResolvedMemoryConfig {
     return { ...this.config };
   }
 
@@ -259,6 +266,152 @@ export class Memory {
     return await this.storage.getResourceWorkingMemory(resourceId);
   }
 
+  // ===== Enhanced Message Operations =====
+
+  /**
+   * Update an existing message
+   * @param messageId The message ID
+   * @param updates Partial message updates
+   * @returns The updated message, or null if storage doesn't support updates
+   */
+  async updateMessage(
+    messageId: string,
+    updates: Partial<Message>,
+  ): Promise<MemoryMessage | null> {
+    if (!this.storage.updateMessage) {
+      return null;
+    }
+
+    return await this.storage.updateMessage(messageId, {
+      message: updates as Message,
+    } as Partial<MemoryMessage>);
+  }
+
+  /**
+   * Get edit history for a message
+   * @param messageId The message ID
+   * @returns Array of message versions, or empty if storage doesn't support history
+   */
+  async getMessageVersions(messageId: string): Promise<MemoryMessage[]> {
+    if (!this.storage.getMessageHistory) {
+      return [];
+    }
+
+    return await this.storage.getMessageHistory(messageId);
+  }
+
+  // ===== Token-Aware Operations =====
+
+  /**
+   * Get total token count for a thread
+   * @param threadId The thread ID
+   * @returns Token count, or 0 if storage doesn't support token counting
+   */
+  async getThreadTokenCount(threadId: string): Promise<number> {
+    if (!this.storage.getThreadTokenCount) {
+      return 0;
+    }
+
+    return await this.storage.getThreadTokenCount(threadId);
+  }
+
+  /**
+   * Get messages within a token budget
+   * Uses contextWindow config if available, otherwise falls back to maxHistoryMessages
+   * @param threadId The thread ID
+   * @param maxTokens Optional max tokens (uses config if not provided)
+   * @returns Array of messages within token budget
+   */
+  async getMessagesWithinBudget(
+    threadId: string,
+    maxTokens?: number,
+  ): Promise<Message[]> {
+    const tokenLimit =
+      maxTokens || this.config.contextWindow?.maxTokens;
+
+    if (!tokenLimit || !this.storage.getMessagesByTokenBudget) {
+      // Fall back to regular getRecentMessages
+      return await this.getRecentMessages(threadId);
+    }
+
+    const memoryMessages = await this.storage.getMessagesByTokenBudget(
+      threadId,
+      tokenLimit,
+    );
+
+    return memoryMessages.map((mm) => mm.message);
+  }
+
+  // ===== Cache Management =====
+
+  /**
+   * Get messages with active cache
+   * @param threadId The thread ID
+   * @returns Array of cached messages, or empty if storage doesn't support caching
+   */
+  async getCachedMessages(threadId: string): Promise<MemoryMessage[]> {
+    if (!this.storage.getCachedMessages) {
+      return [];
+    }
+
+    return await this.storage.getCachedMessages(threadId);
+  }
+
+  /**
+   * Invalidate cache for messages
+   * @param threadId The thread ID
+   * @param beforeDate Optional date - invalidate cache before this date
+   */
+  async invalidateCache(threadId: string, beforeDate?: Date): Promise<void> {
+    if (!this.storage.invalidateCache) {
+      return;
+    }
+
+    await this.storage.invalidateCache(threadId, beforeDate);
+  }
+
+  // ===== Filtered Retrieval =====
+
+  /**
+   * Get messages by status
+   * @param threadId The thread ID
+   * @param status The status to filter by
+   * @returns Array of messages with matching status
+   */
+  async getMessagesByStatus(
+    threadId: string,
+    status: "active" | "archived" | "deleted",
+  ): Promise<MemoryMessage[]> {
+    if (!this.storage.getMessagesByStatus) {
+      // Fall back to getting all messages and filtering
+      const all = await this.storage.getMessages(threadId);
+      return all.filter((m) => m.status === status);
+    }
+
+    return await this.storage.getMessagesByStatus(threadId, status);
+  }
+
+  /**
+   * Get messages by importance threshold
+   * @param threadId The thread ID
+   * @param minImportance Minimum importance score (0-1)
+   * @returns Array of messages meeting importance threshold
+   */
+  async getMessagesByImportance(
+    threadId: string,
+    minImportance: number,
+  ): Promise<MemoryMessage[]> {
+    if (!this.storage.getMessagesByImportance) {
+      // Fall back to getting all messages and filtering
+      const all = await this.storage.getMessages(threadId);
+      return all.filter(
+        (m) => m.importance !== undefined && m.importance >= minImportance,
+      );
+    }
+
+    return await this.storage.getMessagesByImportance(threadId, minImportance);
+  }
+
   // ===== Utility Methods =====
 
   /**

src/lib/memory/storage/in-memory.ts

@@ -28,6 +28,8 @@ export class InMemoryStorage implements MemoryStorage {
   private threadWorkingMemories: Map<string, ThreadWorkingMemory> = new Map();
   private resourceWorkingMemories: Map<string, ResourceWorkingMemory> =
     new Map();
+  // Message version history tracking
+  private messageHistory: Map<string, MemoryMessage[]> = new Map();
 
   // ===== Thread Operations =====
 
@@ -277,4 +279,139 @@ export class InMemoryStorage implements MemoryStorage {
       );
     }
   }
+
+  // ===== Enhanced Message Operations =====
+
+  async updateMessage(
+    messageId: string,
+    updates: Partial<MemoryMessage>,
+  ): Promise<MemoryMessage> {
+    const existing = this.messages.get(messageId);
+    if (!existing) {
+      throw new Error(`Message ${messageId} not found`);
+    }
+
+    // Save current version to history
+    if (!this.messageHistory.has(messageId)) {
+      this.messageHistory.set(messageId, []);
+    }
+    this.messageHistory.get(messageId)!.push({ ...existing });
+
+    // Create updated message with incremented version
+    const updated: MemoryMessage = {
+      ...existing,
+      ...updates,
+      version: (existing.version || 1) + 1,
+      editedFrom: existing.editedFrom || existing.id,
+    };
+
+    this.messages.set(messageId, updated);
+    return updated;
+  }
+
+  async getMessageHistory(messageId: string): Promise<MemoryMessage[]> {
+    const history = this.messageHistory.get(messageId) || [];
+    const current = this.messages.get(messageId);
+
+    // Return history + current (oldest to newest)
+    return current ? [...history, current] : history;
+  }
+
+  // ===== Token-Aware Operations =====
+
+  async getThreadTokenCount(threadId: string): Promise<number> {
+    const messages = await this.getMessages(threadId);
+    return messages.reduce((sum, msg) => sum + (msg.tokenCount || 0), 0);
+  }
+
+  async getMessagesByTokenBudget(
+    threadId: string,
+    maxTokens: number,
+    options: GetMessagesOptions = {},
+  ): Promise<MemoryMessage[]> {
+    const messages = await this.getMessages(threadId, options);
+
+    // Filter to only active messages with token counts
+    const activeMessages = messages.filter(
+      (m) => (!m.status || m.status === "active") && m.tokenCount !== undefined,
+    );
+
+    // Sort by importance (desc) then recency (desc)
+    activeMessages.sort((a, b) => {
+      const importanceA = a.importance || 0;
+      const importanceB = b.importance || 0;
+      if (importanceA !== importanceB) {
+        return importanceB - importanceA; // Higher importance first
+      }
+      return b.createdAt.getTime() - a.createdAt.getTime(); // More recent first
+    });
+
+    // Select messages within token budget
+    const selected: MemoryMessage[] = [];
+    let tokenCount = 0;
+
+    for (const message of activeMessages) {
+      const messageTokens = message.tokenCount || 0;
+      if (tokenCount + messageTokens <= maxTokens) {
+        selected.push(message);
+        tokenCount += messageTokens;
+      }
+    }
+
+    // Sort selected messages chronologically for conversation flow
+    return selected.sort((a, b) => a.createdAt.getTime() - b.createdAt.getTime());
+  }
+
+  // ===== Cache Management Operations =====
+
+  async getCachedMessages(threadId: string): Promise<MemoryMessage[]> {
+    const messages = await this.getMessages(threadId);
+    const now = new Date();
+
+    return messages.filter((msg) => {
+      if (!msg.cacheControl?.enabled) return false;
+      if (!msg.cacheControl.expiresAt) return true;
+      return msg.cacheControl.expiresAt > now;
+    });
+  }
+
+  async invalidateCache(threadId: string, beforeDate?: Date): Promise<void> {
+    const messages = await this.getMessages(threadId);
+    const cutoff = beforeDate || new Date();
+
+    for (const message of messages) {
+      if (message.cacheControl?.enabled) {
+        // Update cache control to mark as expired
+        const updated = {
+          ...message,
+          cacheControl: {
+            ...message.cacheControl,
+            enabled: false,
+            expiresAt: cutoff,
+          },
+        };
+        this.messages.set(message.id, updated);
+      }
+    }
+  }
+
+  // ===== Filtered Retrieval Operations =====
+
+  async getMessagesByStatus(
+    threadId: string,
+    status: string,
+  ): Promise<MemoryMessage[]> {
+    const messages = await this.getMessages(threadId);
+    return messages.filter((m) => m.status === status);
+  }
+
+  async getMessagesByImportance(
+    threadId: string,
+    minImportance: number,
+  ): Promise<MemoryMessage[]> {
+    const messages = await this.getMessages(threadId);
+    return messages.filter(
+      (m) => m.importance !== undefined && m.importance >= minImportance,
+    );
+  }
 }