Skip to content

Commit 892db1c

Browse files
committed
fix(gastown): fix eviction test assertions and replace as casts with Zod validation
- Replace incorrect test that asserted touchAgentHeartbeat() clears drain flag (it doesn't) with tests for acknowledgeContainerReady() which is the actual mechanism that clears drain state - Add test for nonce mismatch case (wrong nonce keeps draining=true) - Add test verifying touchAgentHeartbeat returns drainNonce but doesn't clear draining flag, and only acknowledgeContainerReady clears it - Replace manual 'as' casts in handleContainerReady with Zod schema validation per project conventions
1 parent ea7b576 commit 892db1c

File tree

2 files changed

+45
-21
lines changed

2 files changed

+45
-21
lines changed

cloudflare-gastown/src/handlers/town-eviction.handler.ts

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { Context } from 'hono';
2+
import { z } from 'zod';
23
import { extractBearerToken } from '@kilocode/worker-utils';
34
import type { GastownEnv } from '../gastown.worker';
45
import { getTownDOStub } from '../dos/Town.do';
@@ -125,24 +126,13 @@ export async function handleContainerReady(
125126
return c.json(resError('Cross-town access denied'), 403);
126127
}
127128

128-
let nonce: string | undefined;
129-
try {
130-
const body: unknown = await c.req.json();
131-
if (
132-
body &&
133-
typeof body === 'object' &&
134-
'nonce' in body &&
135-
typeof (body as { nonce: unknown }).nonce === 'string'
136-
) {
137-
nonce = (body as { nonce: string }).nonce;
138-
}
139-
} catch {
140-
// No body or invalid JSON
141-
}
129+
const ContainerReadyBody = z.object({ nonce: z.string() });
142130

143-
if (!nonce) {
131+
const parsed = ContainerReadyBody.safeParse(await c.req.json().catch(() => null));
132+
if (!parsed.success) {
144133
return c.json(resError('Missing required field: nonce'), 400);
145134
}
135+
const { nonce } = parsed.data;
146136

147137
const town = getTownDOStub(c.env, params.townId);
148138
const cleared = await town.acknowledgeContainerReady(nonce);

cloudflare-gastown/test/integration/town-container.test.ts

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -137,23 +137,57 @@ describe('Town DO — container eviction draining lifecycle', () => {
137137
expect(await town.isDraining()).toBe(true);
138138
});
139139

140-
it('touchAgentHeartbeat() clears draining flag after eviction', async () => {
140+
it('acknowledgeContainerReady() clears draining flag with correct nonce', async () => {
141+
const id = `town-${crypto.randomUUID()}`;
142+
const town = env.TOWN.get(env.TOWN.idFromName(id));
143+
144+
// Set draining flag and capture the nonce
145+
const drainNonce = await town.recordContainerEviction();
146+
expect(await town.isDraining()).toBe(true);
147+
148+
// Acknowledge with the correct nonce clears the drain flag
149+
const cleared = await town.acknowledgeContainerReady(drainNonce);
150+
expect(cleared).toBe(true);
151+
expect(await town.isDraining()).toBe(false);
152+
});
153+
154+
it('acknowledgeContainerReady() rejects wrong nonce and stays draining', async () => {
155+
const id = `town-${crypto.randomUUID()}`;
156+
const town = env.TOWN.get(env.TOWN.idFromName(id));
157+
158+
// Set draining flag
159+
await town.recordContainerEviction();
160+
expect(await town.isDraining()).toBe(true);
161+
162+
// Wrong nonce should not clear the drain flag
163+
const cleared = await town.acknowledgeContainerReady('wrong-nonce');
164+
expect(cleared).toBe(false);
165+
expect(await town.isDraining()).toBe(true);
166+
});
167+
168+
it('touchAgentHeartbeat() returns drainNonce during eviction', async () => {
141169
const id = `town-${crypto.randomUUID()}`;
142170
const town = env.TOWN.get(env.TOWN.idFromName(id));
143171

144172
// Register an agent so heartbeat has a valid target
145173
const agent = await town.registerAgent({
146174
role: 'polecat',
147-
name: 'drain-clear-test',
148-
identity: `drain-clear-${id}`,
175+
name: 'drain-nonce-test',
176+
identity: `drain-nonce-${id}`,
149177
});
150178

151179
// Set draining flag
152-
await town.recordContainerEviction();
180+
const drainNonce = await town.recordContainerEviction();
153181
expect(await town.isDraining()).toBe(true);
154182

155-
// Heartbeat should clear draining flag
156-
await town.touchAgentHeartbeat(agent.id);
183+
// Heartbeat returns the drainNonce (but does NOT clear draining)
184+
const heartbeatResult = await town.touchAgentHeartbeat(agent.id);
185+
expect(heartbeatResult.drainNonce).toBe(drainNonce);
186+
expect(await town.isDraining()).toBe(true);
187+
188+
// Only acknowledgeContainerReady with the nonce clears it
189+
const cleared = await town.acknowledgeContainerReady(drainNonce);
190+
expect(cleared).toBe(true);
157191
expect(await town.isDraining()).toBe(false);
158192
});
159193

0 commit comments

Comments
 (0)