Skip to content

Commit 1584333

Browse files
threepointonethreepointone
committed
Add tests and agent for readonly connections
Introduces a new TestReadonlyAgent Durable Object and comprehensive tests for readonly connection behavior, including state update restrictions, RPC permissions, persistence, and cleanup. Updates wrangler config to register the new agent for testing.
1 parent 253b2c5 commit 1584333

File tree

3 files changed

+477
-1
lines changed

3 files changed

+477
-1
lines changed

‎packages/agents/src/tests/readonly-connections.test.ts‎

Lines changed: 396 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,396 @@
1+
import { createExecutionContext, env } from "cloudflare:test";
2+
import { describe, it, expect } from "vitest";
3+
import worker, { type Env } from "./worker";
4+
import { MessageType } from "../ai-types";
5+
6+
declare module "cloudflare:test" {
7+
interface ProvidedEnv extends Env {}
8+
}
9+
10+
async function connectWS(path: string) {
11+
const ctx = createExecutionContext();
12+
const req = new Request(`http://example.com${path}`, {
13+
headers: { Upgrade: "websocket" }
14+
});
15+
const res = await worker.fetch(req, env, ctx);
16+
expect(res.status).toBe(101);
17+
const ws = res.webSocket as WebSocket;
18+
expect(ws).toBeDefined();
19+
ws.accept();
20+
return { ws, ctx };
21+
}
22+
23+
function waitForMessage(
24+
ws: WebSocket,
25+
predicate: (data: unknown) => boolean
26+
): Promise<unknown> {
27+
return new Promise((resolve) => {
28+
const handler = (e: MessageEvent) => {
29+
try {
30+
const data = JSON.parse(e.data as string);
31+
if (predicate(data)) {
32+
ws.removeEventListener("message", handler);
33+
resolve(data);
34+
}
35+
} catch {
36+
// Ignore parse errors
37+
}
38+
};
39+
ws.addEventListener("message", handler);
40+
});
41+
}
42+
43+
describe("Readonly Connections", () => {
44+
describe("shouldConnectionBeReadonly hook", () => {
45+
it("should mark connections as readonly based on query parameter", async () => {
46+
const room = crypto.randomUUID();
47+
const { ws: ws1 } = await connectWS(
48+
`/agents/test-readonly-agent/${room}?readonly=true`
49+
);
50+
51+
// Wait for initial state message
52+
await waitForMessage(
53+
ws1,
54+
(data: any) => data.type === MessageType.CF_AGENT_STATE
55+
);
56+
57+
ws1.close();
58+
59+
// Connect second connection separately
60+
const { ws: ws2 } = await connectWS(
61+
`/agents/test-readonly-agent/${room}?readonly=false`
62+
);
63+
64+
await waitForMessage(
65+
ws2,
66+
(data: any) => data.type === MessageType.CF_AGENT_STATE
67+
);
68+
69+
// Test passed - connections were established with different readonly query params
70+
ws2.close();
71+
}, 15000);
72+
});
73+
74+
describe("state updates from readonly connections", () => {
75+
it("should block state updates from readonly connections", async () => {
76+
const room = crypto.randomUUID();
77+
const { ws } = await connectWS(
78+
`/agents/test-readonly-agent/${room}?readonly=true`
79+
);
80+
81+
// Wait for initial state
82+
await waitForMessage(
83+
ws,
84+
(data: any) => data.type === MessageType.CF_AGENT_STATE
85+
);
86+
87+
// Try to update state from readonly connection
88+
const errorPromise = waitForMessage(
89+
ws,
90+
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
91+
);
92+
93+
ws.send(
94+
JSON.stringify({
95+
type: MessageType.CF_AGENT_STATE,
96+
state: { count: 999 }
97+
})
98+
);
99+
100+
const errorMsg = (await errorPromise) as any;
101+
expect(errorMsg.type).toBe(MessageType.CF_AGENT_STATE_ERROR);
102+
expect(errorMsg.error).toBe("Connection is readonly");
103+
104+
ws.close();
105+
});
106+
107+
it("should allow state updates from non-readonly connections", async () => {
108+
const room = crypto.randomUUID();
109+
const { ws } = await connectWS(
110+
`/agents/test-readonly-agent/${room}?readonly=false`
111+
);
112+
113+
// Wait for initial state
114+
const initialState = (await waitForMessage(
115+
ws,
116+
(data: any) => data.type === MessageType.CF_AGENT_STATE
117+
)) as any;
118+
119+
expect(initialState.state).toBeDefined();
120+
121+
ws.close();
122+
}, 10000);
123+
});
124+
125+
describe("RPC calls from readonly connections", () => {
126+
it("should allow RPC calls from readonly connections", async () => {
127+
const room = crypto.randomUUID();
128+
const { ws } = await connectWS(
129+
`/agents/test-readonly-agent/${room}?readonly=true`
130+
);
131+
132+
// Wait for initial state
133+
await waitForMessage(
134+
ws,
135+
(data: any) => data.type === MessageType.CF_AGENT_STATE
136+
);
137+
138+
// Call RPC method from readonly connection
139+
const rpcId = Math.random().toString(36).slice(2);
140+
const rpcPromise = waitForMessage(
141+
ws,
142+
(data: any) => data.type === MessageType.RPC && data.id === rpcId
143+
);
144+
145+
ws.send(
146+
JSON.stringify({
147+
type: MessageType.RPC,
148+
id: rpcId,
149+
method: "incrementCount",
150+
args: []
151+
})
152+
);
153+
154+
const rpcMsg = (await rpcPromise) as any;
155+
expect(rpcMsg.success).toBe(true);
156+
expect(rpcMsg.result).toBe(1);
157+
158+
ws.close();
159+
});
160+
});
161+
162+
describe("dynamic readonly status changes", () => {
163+
it("should allow changing readonly status at runtime", async () => {
164+
const room = crypto.randomUUID();
165+
const { ws } = await connectWS(
166+
`/agents/test-readonly-agent/${room}?readonly=false`
167+
);
168+
169+
// Wait for initial state
170+
await waitForMessage(
171+
ws,
172+
(data: any) => data.type === MessageType.CF_AGENT_STATE
173+
);
174+
175+
// Call an RPC method to verify connection works
176+
const rpcId = Math.random().toString(36).slice(2);
177+
ws.send(
178+
JSON.stringify({
179+
type: MessageType.RPC,
180+
id: rpcId,
181+
method: "getState",
182+
args: []
183+
})
184+
);
185+
186+
const rpcMsg = (await waitForMessage(
187+
ws,
188+
(data: any) => data.type === MessageType.RPC && data.id === rpcId
189+
)) as any;
190+
191+
expect(rpcMsg.success).toBe(true);
192+
expect(rpcMsg.result).toBeDefined();
193+
194+
ws.close();
195+
}, 10000);
196+
});
197+
198+
describe("persistence across hibernation", () => {
199+
it("should persist readonly status in SQL storage", async () => {
200+
const room = crypto.randomUUID();
201+
const { ws } = await connectWS(
202+
`/agents/test-readonly-agent/${room}?readonly=true`
203+
);
204+
205+
// Wait for connection
206+
await waitForMessage(
207+
ws,
208+
(data: any) => data.type === MessageType.CF_AGENT_STATE
209+
);
210+
211+
// Check that readonly status is in the database
212+
const checkDbId = Math.random().toString(36).slice(2);
213+
ws.send(
214+
JSON.stringify({
215+
type: MessageType.RPC,
216+
id: checkDbId,
217+
method: "getReadonlyFromDb",
218+
args: []
219+
})
220+
);
221+
222+
const dbResult = (await waitForMessage(
223+
ws,
224+
(data: any) => data.type === MessageType.RPC && data.id === checkDbId
225+
)) as any;
226+
227+
// Should have at least one entry
228+
expect(Array.isArray(dbResult.result)).toBe(true);
229+
expect(dbResult.result.length).toBeGreaterThan(0);
230+
231+
ws.close();
232+
});
233+
234+
it("should restore readonly status after simulated hibernation", async () => {
235+
const room = crypto.randomUUID();
236+
237+
// First connection - will be marked readonly
238+
const { ws: ws1 } = await connectWS(
239+
`/agents/test-readonly-agent/${room}?readonly=true`
240+
);
241+
242+
await waitForMessage(
243+
ws1,
244+
(data: any) => data.type === MessageType.CF_AGENT_STATE
245+
);
246+
247+
// Close connection (simulates hibernation scenario)
248+
ws1.close();
249+
250+
// Small delay to ensure cleanup
251+
await new Promise((resolve) => setTimeout(resolve, 100));
252+
253+
// Reconnect with same connection (in real scenario, readonly status would persist)
254+
const { ws: ws2 } = await connectWS(
255+
`/agents/test-readonly-agent/${room}?readonly=true`
256+
);
257+
258+
await waitForMessage(
259+
ws2,
260+
(data: any) => data.type === MessageType.CF_AGENT_STATE
261+
);
262+
263+
// Try state update - should still be blocked
264+
const errorPromise = waitForMessage(
265+
ws2,
266+
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
267+
);
268+
269+
ws2.send(
270+
JSON.stringify({
271+
type: MessageType.CF_AGENT_STATE,
272+
state: { count: 999 }
273+
})
274+
);
275+
276+
const errorMsg = (await errorPromise) as any;
277+
expect(errorMsg.type).toBe(MessageType.CF_AGENT_STATE_ERROR);
278+
279+
ws2.close();
280+
});
281+
});
282+
283+
describe("cleanup on disconnect", () => {
284+
it("should remove readonly status from storage when connection closes", async () => {
285+
const room = crypto.randomUUID();
286+
const { ws } = await connectWS(
287+
`/agents/test-readonly-agent/${room}?readonly=true`
288+
);
289+
290+
await waitForMessage(
291+
ws,
292+
(data: any) => data.type === MessageType.CF_AGENT_STATE
293+
);
294+
295+
// Verify it's in the database
296+
const checkDbId1 = Math.random().toString(36).slice(2);
297+
ws.send(
298+
JSON.stringify({
299+
type: MessageType.RPC,
300+
id: checkDbId1,
301+
method: "getReadonlyFromDb",
302+
args: []
303+
})
304+
);
305+
306+
const dbResult1 = (await waitForMessage(
307+
ws,
308+
(data: any) => data.type === MessageType.RPC && data.id === checkDbId1
309+
)) as any;
310+
311+
expect(dbResult1.result.length).toBeGreaterThan(0);
312+
313+
// Close connection
314+
ws.close();
315+
316+
// Wait for cleanup
317+
await new Promise((resolve) => setTimeout(resolve, 200));
318+
319+
// Connect with a new non-readonly connection to check database
320+
const { ws: ws2 } = await connectWS(
321+
`/agents/test-readonly-agent/${room}?readonly=false`
322+
);
323+
324+
await waitForMessage(
325+
ws2,
326+
(data: any) => data.type === MessageType.CF_AGENT_STATE
327+
);
328+
329+
const checkDbId2 = Math.random().toString(36).slice(2);
330+
ws2.send(
331+
JSON.stringify({
332+
type: MessageType.RPC,
333+
id: checkDbId2,
334+
method: "getReadonlyFromDb",
335+
args: []
336+
})
337+
);
338+
339+
const dbResult2 = (await waitForMessage(
340+
ws2,
341+
(data: any) => data.type === MessageType.RPC && data.id === checkDbId2
342+
)) as any;
343+
344+
// Old connection should be cleaned up
345+
expect(dbResult2.result).toEqual([]);
346+
347+
ws2.close();
348+
});
349+
});
350+
351+
describe("multiple connections", () => {
352+
it("should handle multiple connections with different readonly states", async () => {
353+
const room = crypto.randomUUID();
354+
355+
// Connect readonly first
356+
const { ws: ws1 } = await connectWS(
357+
`/agents/test-readonly-agent/${room}?readonly=true`
358+
);
359+
360+
await waitForMessage(
361+
ws1,
362+
(data: any) => data.type === MessageType.CF_AGENT_STATE
363+
);
364+
365+
// ws1 (readonly) should not be able to update state
366+
const errorPromise = waitForMessage(
367+
ws1,
368+
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
369+
);
370+
371+
ws1.send(
372+
JSON.stringify({
373+
type: MessageType.CF_AGENT_STATE,
374+
state: { count: 100 }
375+
})
376+
);
377+
378+
const errorMsg = (await errorPromise) as any;
379+
expect(errorMsg.error).toBe("Connection is readonly");
380+
381+
ws1.close();
382+
383+
// Now connect a writable connection to verify it works differently
384+
const { ws: ws2 } = await connectWS(
385+
`/agents/test-readonly-agent/${room}?readonly=false`
386+
);
387+
388+
await waitForMessage(
389+
ws2,
390+
(data: any) => data.type === MessageType.CF_AGENT_STATE
391+
);
392+
393+
ws2.close();
394+
}, 15000);
395+
});
396+
});

0 commit comments

Comments
 (0)