Make session containers sticky and harden warm pool

Author: stainlu

src/index.ts

@@ -160,6 +160,10 @@ async function main(): Promise<void> {
   const gatewayPort = envInt("OPENCLAW_GATEWAY_PORT", 18789);
   const readyTimeoutMs = envInt("OPENCLAW_READY_TIMEOUT_MS", 60_000);
   const runTimeoutMs = envInt("OPENCLAW_RUN_TIMEOUT_MS", 10 * 60_000);
+  // Idle reap now applies only to sessions the caller explicitly marks
+  // as reapable (today: ephemeral one-shot sessions). Sticky product
+  // sessions keep their owned container alive between turns until
+  // explicit teardown or process restart/adoption.
   const idleTimeoutMs = envInt("OPENCLAW_IDLE_TIMEOUT_MS", 10 * 60_000);
   const sweepIntervalMs = envInt("OPENCLAW_SWEEP_INTERVAL_MS", 60_000);
   // Warm pool is bounded so a host with many agent templates does not
@@ -246,18 +250,21 @@ async function main(): Promise<void> {
 
   // Per-session container pool. isBusy closes over the session store so the
   // sweeper can skip containers whose session currently has a run in flight
-  // — the pool itself has no store dependency. cleanupOnReap closes over
-  // BOTH the store and the JSONL reader so it can tear down ephemeral
-  // sessions (auto-created by keyless POST /v1/chat/completions) along with
-  // their container. Called only on the idle-reap path; manual evictSession
-  // and shutdown paths preserve session data.
+  // — the pool itself has no store dependency. shouldReapSession marks only
+  // ephemeral sessions as eligible for idle reap; sticky sessions keep their
+  // container until explicit teardown. cleanupOnReap closes over BOTH the
+  // store and the JSONL reader so it can tear down those ephemeral sessions
+  // (auto-created by keyless POST /v1/chat/completions) along with their
+  // container. Called only on the idle-reap path; manual evictSession and
+  // shutdown paths preserve session data.
   const pool = new SessionContainerPool(runtime, {
     idleTimeoutMs,
     readyTimeoutMs,
     sweepIntervalMs,
     maxWarmContainers,
     warmIdleTimeoutMs,
     isBusy: (sessionId) => store.sessions.get(sessionId)?.status === "running",
+    shouldReapSession: (sessionId) => store.sessions.get(sessionId)?.ephemeral === true,
     cleanupOnReap: async (sessionId) => {
       const session = store.sessions.get(sessionId);
       if (!session?.ephemeral) return;

src/orchestrator/router.ts

@@ -465,6 +465,20 @@ export class AgentRouter {
     await this.pool.dropWarmForAgent(agentId);
   }
 
+  /**
+   * Tear down any live runtime resources for a session without deleting the
+   * session metadata itself. Used by DELETE /v1/sessions so the container,
+   * queue, pending approvals, and persistent session↔container mapping do
+   * not outlive the session row.
+   */
+  async disposeSessionRuntime(sessionId: string): Promise<void> {
+    this.cancelledDuringAcquire.delete(sessionId);
+    this.queue.clear(sessionId);
+    this.pendingApprovals.delete(sessionId);
+    this.clearApprovalSubscriptions(sessionId);
+    await this.pool.evictSession(sessionId);
+  }
+
   /**
    * Post a user.message to an existing session. Behavior depends on
    * session status:

src/orchestrator/server.test.ts

@@ -12,6 +12,7 @@ function makeApp() {
   const routerCalls = {
     warmForAgent: [] as string[],
     dropWarmForAgent: [] as string[],
+    disposeSessionRuntime: [] as string[],
   };
 
   const events = {
@@ -67,6 +68,10 @@ function makeApp() {
       routerCalls.dropWarmForAgent.push(agentId);
       return;
     },
+    async disposeSessionRuntime(sessionId: string) {
+      routerCalls.disposeSessionRuntime.push(sessionId);
+      return;
+    },
     async runEvent(args: { sessionId: string; content: string }) {
       const started = store.sessions.beginRun(args.sessionId);
       if (!started) {
@@ -338,4 +343,22 @@ describe("session ownership in the HTTP API", () => {
     expect(routerCalls.warmForAgent).toContain(agent.agentId);
     expect(store.agents.get(agent.agentId)?.model).toBe("openai/gpt-5.4");
   });
+
+  it("evicts live runtime state before deleting a session", async () => {
+    const { app, store, routerCalls } = makeApp();
+    const agent = createAgent(store);
+    const session = store.sessions.create({
+      agentId: agent.agentId,
+      userId: null,
+    });
+
+    const res = await req(app, `/v1/sessions/${session.sessionId}`, {
+      method: "DELETE",
+      token: "admin-secret",
+    });
+
+    expect(res.status).toBe(200);
+    expect(routerCalls.disposeSessionRuntime).toEqual([session.sessionId]);
+    expect(store.sessions.get(session.sessionId)).toBeUndefined();
+  });
 });

src/orchestrator/server.ts

@@ -677,6 +677,9 @@ export function buildApp(deps: ServerDeps): Hono {
       outcome: archived ? "ok" : "agent_not_found",
     });
     if (!archived) return c.json({ error: "agent_not_found" }, 404);
+    void deps.router.dropWarmForAgent(agentId).catch((err) => {
+      log.warn({ err, agent_id: agentId }, "drop-warm-for-agent after archive failed (non-fatal)");
+    });
     return c.json(agentResponse(archived));
   });
 
@@ -1036,7 +1039,7 @@ export function buildApp(deps: ServerDeps): Hono {
     return c.json(sessionResponse(session, deps.events, deps.sessionContainers));
   });
 
-  app.delete("/v1/sessions/:sessionId", (c) => {
+  app.delete("/v1/sessions/:sessionId", async (c) => {
     const sessionId = c.req.param("sessionId");
     const session = getScopedSession(c, sessionId);
     if (!session) {
@@ -1047,8 +1050,19 @@ export function buildApp(deps: ServerDeps): Hono {
       });
       return c.json({ error: "session_not_found" }, 404);
     }
-    // Drop the Pi JSONL + sessions.json entry on disk first, then the
-    // orchestrator-side metadata row.
+    // Tear down the live runtime first so sticky session containers do
+    // not outlive their session row. Only once the container is gone do
+    // we delete the persisted JSONL / metadata.
+    try {
+      await deps.router.disposeSessionRuntime(sessionId);
+    } catch (err) {
+      writeAudit(deps.audit, c, {
+        action: "session.delete",
+        target: sessionId,
+        outcome: err instanceof RouterError ? err.code : "error",
+      });
+      return handleRouterError(err, c);
+    }
     deps.events.deleteBySession(session.agentId, session.sessionId);
     deps.sessions.delete(sessionId);
     writeAudit(deps.audit, c, {

src/runtime/gateway-ws.ts

@@ -105,6 +105,10 @@ export class GatewayWebSocketClient {
 
   constructor(private readonly cfg: GatewayWsConfig) {}
 
+  isConnected(): boolean {
+    return this.connected && !this.closed;
+  }
+
   /**
    * Subscribe to gateway broadcast events. Returns an unsubscribe function.
    * Used by the orchestrator to listen for `plugin.approval.requested`

src/runtime/pool.test.ts

@@ -9,15 +9,20 @@ vi.mock("./gateway-ws.js", () => {
   const instances: FakeWs[] = [];
   class FakeWs {
     readonly closeCount = { n: 0 };
+    private connected = false;
     constructor(public readonly cfg: unknown) {
       instances.push(this);
     }
     async connect(): Promise<void> {
-      /* noop */
+      this.connected = true;
     }
     async close(): Promise<void> {
+      this.connected = false;
       this.closeCount.n += 1;
     }
+    isConnected(): boolean {
+      return this.connected;
+    }
     async abort(): Promise<void> {
       /* noop */
     }
@@ -413,6 +418,39 @@ describe("SessionContainerPool.warmForAgent", () => {
     expect(runtime.stopped.has("cnt_1")).toBe(true);
     await pool.shutdown();
   });
+
+  it("does not exceed maxWarmContainers when different agents warm concurrently", async () => {
+    const { pool, runtime } = makePool({ maxWarmContainers: 1 });
+    runtime.readyDelayMs = 25;
+
+    await Promise.all([
+      pool.warmForAgent("agt_a", baseSpawnOptions("warm-a")),
+      pool.warmForAgent("agt_b", baseSpawnOptions("warm-b")),
+    ]);
+
+    expect(runtime.calls.filter((c) => c.kind === "spawn")).toHaveLength(1);
+    await pool.shutdown();
+  });
+
+  it("cancels an inflight warm when dropWarmForAgent is called", async () => {
+    const { pool, runtime } = makePool();
+    runtime.readyDelayMs = 25;
+
+    const warming = pool.warmForAgent("agt_x", baseSpawnOptions());
+    await new Promise((r) => setImmediate(r));
+    await pool.dropWarmForAgent("agt_x");
+    await warming;
+
+    const c = await pool.acquireForSession({
+      sessionId: "ses_after_drop",
+      spawnOptions: baseSpawnOptions(),
+      agentId: "agt_x",
+    });
+
+    expect(c.id).toBe("cnt_2");
+    expect(runtime.stopped.has("cnt_1")).toBe(true);
+    await pool.shutdown();
+  });
 });
 
 describe("SessionContainerPool.evictSession", () => {
@@ -501,6 +539,29 @@ describe("SessionContainerPool.reapIdle", () => {
     await pool.shutdown();
   });
 
+  it("keeps sticky sessions alive when shouldReapSession returns false", async () => {
+    const { pool, runtime } = makePool({
+      idleTimeoutMs: 50,
+      shouldReapSession: (id) => id === "ses_ephemeral",
+    });
+    await pool.acquireForSession({
+      sessionId: "ses_sticky",
+      spawnOptions: baseSpawnOptions(),
+    });
+    await pool.acquireForSession({
+      sessionId: "ses_ephemeral",
+      spawnOptions: baseSpawnOptions(),
+    });
+    await new Promise((r) => setTimeout(r, 80));
+    // @ts-expect-error — private reapIdle
+    await pool.reapIdle();
+    const sessionsLeft = pool.snapshot().map((e) => e.sessionId);
+    expect(sessionsLeft).toEqual(["ses_sticky"]);
+    expect(runtime.stopped.has("cnt_1")).toBe(false);
+    expect(runtime.stopped.has("cnt_2")).toBe(true);
+    await pool.shutdown();
+  });
+
   it("reaps warm containers past warmIdleTimeoutMs", async () => {
     const { pool, runtime } = makePool({
       idleTimeoutMs: 10 * 60_000,
@@ -534,6 +595,18 @@ describe("SessionContainerPool.shutdown", () => {
     expect(runtime.stopped.has("cnt_1")).toBe(true);
     expect(runtime.stopped.has("cnt_2")).toBe(true);
   });
+
+  it("waits for inflight warm boots and stops them during shutdown", async () => {
+    const { pool, runtime } = makePool();
+    runtime.readyDelayMs = 25;
+
+    const warming = pool.warmForAgent("agt_x", baseSpawnOptions());
+    await new Promise((r) => setImmediate(r));
+    await pool.shutdown();
+    await warming;
+
+    expect(runtime.stopped.has("cnt_1")).toBe(true);
+  });
 });
 
 // --------------------------------------------------------------------
@@ -774,4 +847,50 @@ describe("SessionContainerPool — networking: limited", () => {
     expect(spawnsAfter - spawnsBefore).toBe(2);
     await pool.shutdown();
   });
+
+  it("falls back to a cold spawn when a claimed warm container is no longer ready", async () => {
+    const { pool, runtime } = makePool();
+    await pool.warmForAgent("agt_x", baseSpawnOptions());
+    const baseWaitForReady = runtime.waitForReady.bind(runtime);
+    runtime.waitForReady = async (container, timeoutMs) => {
+      if (container.id === "cnt_1") {
+        throw new Error("warm container is dead");
+      }
+      return baseWaitForReady(container, timeoutMs);
+    };
+
+    const c = await pool.acquireForSession({
+      sessionId: "ses_claim",
+      spawnOptions: baseSpawnOptions(),
+      agentId: "agt_x",
+    });
+
+    expect(c.id).toBe("cnt_2");
+    expect(runtime.stopped.has("cnt_1")).toBe(true);
+    await pool.shutdown();
+  });
+
+  it("cold-respawns when the existing active container fails the reuse probe", async () => {
+    const { pool, runtime } = makePool();
+    const first = await pool.acquireForSession({
+      sessionId: "ses_live",
+      spawnOptions: baseSpawnOptions(),
+    });
+    const baseWaitForReady = runtime.waitForReady.bind(runtime);
+    runtime.waitForReady = async (container, timeoutMs) => {
+      if (container.id === first.id) {
+        throw new Error("active container is dead");
+      }
+      return baseWaitForReady(container, timeoutMs);
+    };
+
+    const second = await pool.acquireForSession({
+      sessionId: "ses_live",
+      spawnOptions: baseSpawnOptions(),
+    });
+
+    expect(second.id).toBe("cnt_2");
+    expect(runtime.stopped.has(first.id)).toBe(true);
+    await pool.shutdown();
+  });
 });