Merge pull request #262 from getcompanion-ai/memory-staging

Add first-class memory management plumbing

packages/coding-agent/src/core/agent-session.ts

@@ -61,6 +61,15 @@ import {
   type ToolHtmlRenderer,
 } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
+import {
+  RuntimeMemoryManager,
+  type RuntimeMemoryForgetInput,
+  type RuntimeMemoryRebuildResult,
+  type RuntimeMemoryRememberInput,
+  type RuntimeMemorySearchResult,
+  type RuntimeMemoryStatus,
+  type RuntimeMemoryRecord,
+} from "./memory/runtime-memory.js";
 import {
   type ContextUsage,
   type ExtensionCommandContextActions,
@@ -337,6 +346,9 @@ export class AgentSession {
 
   // Base system prompt (without extension appends) - used to apply fresh appends each turn
   private _baseSystemPrompt = "";
+  private _memoryManager: RuntimeMemoryManager;
+  private _memoryWriteQueue: Promise<void> = Promise.resolve();
+  private _memoryDisposePromise: Promise<void> | null = null;
 
   constructor(config: AgentSessionConfig) {
     this.agent = config.agent;
@@ -350,6 +362,10 @@ export class AgentSession {
     this._extensionRunnerRef = config.extensionRunnerRef;
     this._initialActiveToolNames = config.initialActiveToolNames;
     this._baseToolsOverride = config.baseToolsOverride;
+    this._memoryManager = new RuntimeMemoryManager({
+      sessionManager: this.sessionManager,
+      settingsManager: this.settingsManager,
+    });
 
     // Always subscribe to agent events for internal handling
     // (session persistence, extensions, auto-compaction, retry logic)
@@ -499,6 +515,16 @@ export class AgentSession {
           this._resolveRetry();
         }
       }
+
+      if (event.message.role === "user" || event.message.role === "assistant") {
+        try {
+          this._memoryManager.recordMessage(event.message);
+        } catch (error) {
+          const message =
+            error instanceof Error ? error.message : String(error);
+          console.error(`[memory] episode write failed: ${message}`);
+        }
+      }
     }
 
     // Check auto-retry and auto-compaction after agent completes
@@ -513,6 +539,10 @@ export class AgentSession {
       }
 
       await this._checkCompaction(msg);
+
+      if (msg.stopReason !== "error") {
+        this._enqueueMemoryPromotion([...this.agent.state.messages]);
+      }
     }
   }
 
@@ -667,6 +697,7 @@ export class AgentSession {
   dispose(): void {
     this._disconnectFromAgent();
     this._eventListeners = [];
+    void this._disposeMemoryManager();
   }
 
   // =========================================================================
@@ -804,6 +835,107 @@ export class AgentSession {
     return this._resourceLoader.getPrompts().prompts;
   }
 
+  async transformRuntimeContext(
+    messages: AgentMessage[],
+    signal?: AbortSignal,
+  ): Promise<AgentMessage[]> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.injectContext(messages, { signal });
+  }
+
+  async getMemoryStatus(): Promise<RuntimeMemoryStatus> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.getStatus();
+  }
+
+  async getCoreMemories(): Promise<RuntimeMemoryRecord[]> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.listCoreMemories();
+  }
+
+  async searchMemory(
+    query: string,
+    limit?: number,
+  ): Promise<RuntimeMemorySearchResult> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.search(query, limit);
+  }
+
+  async rememberMemory(
+    input: RuntimeMemoryRememberInput,
+  ): Promise<RuntimeMemoryRecord | null> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.remember(input);
+  }
+
+  async forgetMemory(
+    input: RuntimeMemoryForgetInput,
+  ): Promise<{ ok: true; forgotten: boolean }> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.forget(input);
+  }
+
+  async rebuildMemory(): Promise<RuntimeMemoryRebuildResult> {
+    await this._awaitMemoryWrites();
+    return this._memoryManager.rebuild();
+  }
+
+  private async _awaitMemoryWrites(): Promise<void> {
+    try {
+      await this._memoryWriteQueue;
+    } catch {
+      // Memory writes are best-effort; failures should not block chat.
+    }
+  }
+
+  private async _disposeMemoryManager(): Promise<void> {
+    if (this._memoryDisposePromise) {
+      await this._memoryDisposePromise;
+      return;
+    }
+
+    this._memoryDisposePromise = (async () => {
+      try {
+        await this._agentEventQueue;
+      } catch {
+        // Event processing failures should not block shutdown.
+      }
+
+      try {
+        await this._memoryWriteQueue;
+      } catch {
+        // Memory writes are best-effort during shutdown too.
+      }
+
+      this._memoryManager.dispose();
+    })();
+
+    await this._memoryDisposePromise;
+  }
+
+  private _enqueueMemoryPromotion(messages: AgentMessage[]): void {
+    this._memoryWriteQueue = this._memoryWriteQueue
+      .catch(() => undefined)
+      .then(async () => {
+        if (!this.model) {
+          return;
+        }
+        const apiKey = await this._modelRegistry.getApiKey(this.model);
+        if (!apiKey) {
+          return;
+        }
+        await this._memoryManager.promoteTurn({
+          model: this.model,
+          apiKey,
+          messages,
+        });
+      })
+      .catch((error: unknown) => {
+        const message = error instanceof Error ? error.message : String(error);
+        console.error(`[memory] promotion failed: ${message}`);
+      });
+  }
+
   private _normalizePromptSnippet(
     text: string | undefined,
   ): string | undefined {

packages/coding-agent/src/core/gateway/runtime.ts

@@ -297,6 +297,16 @@ export class GatewayRuntime {
     return managedSession;
   }
 
+  private async resolveMemorySession(
+    sessionKey: string | null | undefined,
+  ): Promise<AgentSession> {
+    if (!sessionKey || sessionKey === this.primarySessionKey) {
+      return this.primarySession;
+    }
+    const managedSession = await this.ensureSession(sessionKey);
+    return managedSession.session;
+  }
+
   private async processNext(
     managedSession: ManagedGatewaySession,
   ): Promise<void> {
@@ -624,6 +634,102 @@ export class GatewayRuntime {
       return;
     }
 
+    if (method === "GET" && path === "/memory/status") {
+      const sessionKey = url.searchParams.get("sessionKey");
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const memory = await memorySession.getMemoryStatus();
+      this.writeJson(response, 200, { memory });
+      return;
+    }
+
+    if (method === "GET" && path === "/memory/core") {
+      const sessionKey = url.searchParams.get("sessionKey");
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const memories = await memorySession.getCoreMemories();
+      this.writeJson(response, 200, { memories });
+      return;
+    }
+
+    if (method === "POST" && path === "/memory/search") {
+      const body = await this.readJsonBody(request);
+      const query = typeof body.query === "string" ? body.query : "";
+      const limit =
+        typeof body.limit === "number" && Number.isFinite(body.limit)
+          ? Math.max(1, Math.floor(body.limit))
+          : undefined;
+      const sessionKey =
+        typeof body.sessionKey === "string" ? body.sessionKey : undefined;
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const result = await memorySession.searchMemory(query, limit);
+      this.writeJson(response, 200, result);
+      return;
+    }
+
+    if (method === "POST" && path === "/memory/remember") {
+      const body = await this.readJsonBody(request);
+      const content = typeof body.content === "string" ? body.content : "";
+      if (!content.trim()) {
+        this.writeJson(response, 400, { error: "Missing memory content" });
+        return;
+      }
+      const sessionKey =
+        typeof body.sessionKey === "string" ? body.sessionKey : undefined;
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const memory = await memorySession.rememberMemory({
+        bucket:
+          body.bucket === "core" || body.bucket === "archival"
+            ? body.bucket
+            : undefined,
+        kind:
+          body.kind === "profile" ||
+          body.kind === "preference" ||
+          body.kind === "relationship" ||
+          body.kind === "fact" ||
+          body.kind === "secret"
+            ? body.kind
+            : undefined,
+        key: typeof body.key === "string" ? body.key : undefined,
+        content,
+        source: "manual",
+      });
+      this.writeJson(response, 200, { ok: true, memory });
+      return;
+    }
+
+    if (method === "POST" && path === "/memory/forget") {
+      const body = await this.readJsonBody(request);
+      const id =
+        typeof body.id === "number" && Number.isFinite(body.id)
+          ? Math.floor(body.id)
+          : undefined;
+      const key = typeof body.key === "string" ? body.key : undefined;
+      if (id === undefined && !key) {
+        this.writeJson(response, 400, {
+          error: "Memory forget requires an id or key",
+        });
+        return;
+      }
+      const sessionKey =
+        typeof body.sessionKey === "string" ? body.sessionKey : undefined;
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const result = await memorySession.forgetMemory({
+        id,
+        key,
+      });
+      this.writeJson(response, 200, result);
+      return;
+    }
+
+    if (method === "POST" && path === "/memory/rebuild") {
+      const body = await this.readJsonBody(request);
+      const sessionKey =
+        typeof body.sessionKey === "string" ? body.sessionKey : undefined;
+      const memorySession = await this.resolveMemorySession(sessionKey);
+      const result = await memorySession.rebuildMemory();
+      this.writeJson(response, 200, result);
+      return;
+    }
+
     const sessionMatch = path.match(
       /^\/sessions\/([^/]+)(?:\/(events|messages|abort|reset|chat|history|model|reload))?$/,
     );

packages/coding-agent/src/core/sdk.ts

@@ -320,6 +320,7 @@ export async function createAgentSession(
   };
 
   const extensionRunnerRef: { current?: ExtensionRunner } = {};
+  const sessionRef: { current?: AgentSession } = {};
 
   agent = new Agent({
     initialState: {
@@ -331,9 +332,15 @@ export async function createAgentSession(
     convertToLlm: convertToLlmWithBlockImages,
     sessionId: sessionManager.getSessionId(),
     transformContext: async (messages) => {
+      const currentSession = sessionRef.current;
+      let transformedMessages = messages;
+      if (currentSession) {
+        transformedMessages =
+          await currentSession.transformRuntimeContext(transformedMessages);
+      }
       const runner = extensionRunnerRef.current;
-      if (!runner) return messages;
-      return runner.emitContext(messages);
+      if (!runner) return transformedMessages;
+      return runner.emitContext(transformedMessages);
     },
     steeringMode: settingsManager.getSteeringMode(),
     followUpMode: settingsManager.getFollowUpMode(),
@@ -393,6 +400,7 @@ export async function createAgentSession(
     initialActiveToolNames,
     extensionRunnerRef,
   });
+  sessionRef.current = session;
   const extensionsResult = resourceLoader.getExtensions();
 
   return {

packages/coding-agent/src/core/settings-manager.ts

@@ -63,6 +63,17 @@ export interface GatewaySettings {
   webhook?: GatewayWebhookSettings;
 }
 
+export interface CompanionMemorySettings {
+  enabled?: boolean;
+  storageDir?: string;
+  maxCoreTokens?: number;
+  maxRecallResults?: number;
+  writer?: {
+    enabled?: boolean;
+    maxTokens?: number;
+  };
+}
+
 export type TransportSetting = Transport;
 
 /**
@@ -125,6 +136,7 @@ export interface Settings {
   showHardwareCursor?: boolean; // Show terminal cursor while still positioning it for IME
   markdown?: MarkdownSettings;
   gateway?: GatewaySettings;
+  companionMemory?: CompanionMemorySettings;
 }
 
 /** Deep merge settings: project/overrides take precedence, nested objects merge recursively */