fix: re-read container state after wake so restored sessions accept prompts

VmAgentContainer.proxyHttp/fetch read the container state once before
wakeFromSnapshot(), then applied the stopped -> 410 guard using that stale
pre-wake state. A freshly-woken, restored container was rejected with 410
(surfaced by the Worker as a generic 500), so users could restore state but
not continue chatting after wake. Re-read the state after a successful wake.

Adds a discriminating regression test that fails on the pre-fix code.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: raphaeltm

apps/api/src/durable-objects/vm-agent-container.ts

@@ -109,13 +109,18 @@ export class VmAgentContainer extends Container<Env> {
   }
 
   async proxyHttp(request: Request, port?: number): Promise<Response> {
-    const state = await this.getState();
+    let state = await this.getState();
     const lifecycleStatus = await this.ctx.storage.get<LifecycleStatus>('lifecycleStatus');
     if (lifecycleStatus === 'sleeping') {
       const wake = await this.wakeFromSnapshot();
       if (!wake.ok) {
         return new Response(wake.message || WAKE_DEGRADED_RESPONSE, { status: 503 });
       }
+      // wakeFromSnapshot launched a fresh container and restored the session.
+      // Re-read the container state so the stopped-check below reflects the
+      // now-running container instead of the pre-wake stopped snapshot,
+      // otherwise a successfully-woken session is wrongly rejected with 410.
+      state = await this.getState();
     }
     if (state.status === 'stopped' || state.status === 'stopped_with_code') {
       return new Response('Container is stopped; create a new instant session.', { status: 410 });
@@ -234,13 +239,16 @@ export class VmAgentContainer extends Container<Env> {
   }
 
   override async fetch(request: Request): Promise<Response> {
-    const state = await this.getState();
+    let state = await this.getState();
     const lifecycleStatus = await this.ctx.storage.get<LifecycleStatus>('lifecycleStatus');
     if (lifecycleStatus === 'sleeping') {
       const wake = await this.wakeFromSnapshot();
       if (!wake.ok) {
         return new Response(wake.message || WAKE_DEGRADED_RESPONSE, { status: 503 });
       }
+      // Re-read the container state after wake so the stopped-check reflects
+      // the freshly-launched container, not the pre-wake stopped snapshot.
+      state = await this.getState();
     }
     if (state.status === 'stopped' || state.status === 'stopped_with_code') {
       return new Response('Container is stopped; create a new instant session.', { status: 410 });

apps/api/tests/unit/durable-objects/vm-agent-container-wake-state.test.ts

@@ -0,0 +1,96 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { VmAgentContainer } from '../../../src/durable-objects/vm-agent-container';
+
+// Regression test for the restored-session prompt failure: `proxyHttp` read the
+// container state ONCE before `wakeFromSnapshot()`, then applied the
+// `stopped`/`stopped_with_code` -> 410 guard using that stale pre-wake state.
+// A freshly-woken, restored container was therefore rejected with 410 (surfaced
+// by the Worker as a generic 500), even though restore succeeded. The fix
+// re-reads the container state after a successful wake.
+
+interface FakeState {
+  status: string;
+}
+
+function makeFake(opts: {
+  statuses: string[]; // sequence returned by getState()
+  lifecycleStatus: string;
+  wakeOk: boolean;
+}) {
+  const getState = vi.fn<[], Promise<FakeState>>();
+  for (const s of opts.statuses) {
+    getState.mockResolvedValueOnce({ status: s });
+  }
+  getState.mockResolvedValue({ status: opts.statuses[opts.statuses.length - 1] });
+
+  const containerFetch = vi
+    .fn()
+    .mockResolvedValue(new Response('proxied', { status: 200 }));
+  const wakeFromSnapshot = vi
+    .fn()
+    .mockResolvedValue(opts.wakeOk ? { ok: true } : { ok: false, message: 'degraded' });
+
+  const fake = {
+    getState,
+    containerFetch,
+    wakeFromSnapshot,
+    defaultPort: 8080,
+    ctx: { storage: { get: vi.fn().mockResolvedValue(opts.lifecycleStatus) } },
+  };
+  return { fake, getState, containerFetch, wakeFromSnapshot };
+}
+
+function callProxyHttp(fake: unknown, request: Request): Promise<Response> {
+  return (VmAgentContainer.prototype as unknown as {
+    proxyHttp: (this: unknown, request: Request, port?: number) => Promise<Response>;
+  }).proxyHttp.call(fake, request);
+}
+
+describe('VmAgentContainer.proxyHttp wake state re-read', () => {
+  it('proxies the prompt after a successful wake even though the pre-wake state was stopped', async () => {
+    // Pre-wake getState -> stopped; post-wake getState -> running (fresh container).
+    const { fake, getState, containerFetch, wakeFromSnapshot } = makeFake({
+      statuses: ['stopped', 'running'],
+      lifecycleStatus: 'sleeping',
+      wakeOk: true,
+    });
+
+    const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' }));
+
+    expect(wakeFromSnapshot).toHaveBeenCalledTimes(1);
+    // State must be re-read after wake (once before, once after) so the stopped
+    // guard sees the now-running container.
+    expect(getState).toHaveBeenCalledTimes(2);
+    // The request is proxied to the running container, NOT rejected with 410.
+    expect(containerFetch).toHaveBeenCalledTimes(1);
+    expect(res.status).toBe(200);
+  });
+
+  it('returns 503 (not 410/proxy) when wake fails', async () => {
+    const { fake, containerFetch } = makeFake({
+      statuses: ['stopped', 'stopped'],
+      lifecycleStatus: 'sleeping',
+      wakeOk: false,
+    });
+
+    const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' }));
+
+    expect(res.status).toBe(503);
+    expect(containerFetch).not.toHaveBeenCalled();
+  });
+
+  it('still returns 410 for a genuinely stopped, non-sleeping container', async () => {
+    const { fake, containerFetch, wakeFromSnapshot } = makeFake({
+      statuses: ['stopped'],
+      lifecycleStatus: 'running',
+      wakeOk: true,
+    });
+
+    const res = await callProxyHttp(fake, new Request('http://container/prompt', { method: 'POST' }));
+
+    expect(wakeFromSnapshot).not.toHaveBeenCalled();
+    expect(containerFetch).not.toHaveBeenCalled();
+    expect(res.status).toBe(410);
+  });
+});