fix(agents): drop late frames from a deleted sub-agent's socket so th… · cloudflare/agents@8aba973 · GitHub
Skip to content

Commit 8aba973

Browse files
fix(agents): drop late frames from a deleted sub-agent's socket so they can't reach a same-name replacement (#2003)
Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 9be7381 commit 8aba973

3 files changed

Lines changed: 76 additions & 0 deletions

File tree

‎packages/agents/src/dynamic-agents/dynamic-agents.ts‎

Lines changed: 16 additions & 0 deletions

‎packages/agents/src/tests/agents/sub-agent.ts‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import { Agent, callable, getCurrentAgent } from "../../index.ts";
22
import type {
3+
Connection,
4+
WSMessage,
35
FiberInspection,
46
FiberRecoveryContext,
57
FiberRecoveryResult,
@@ -8,6 +10,8 @@ import type {
810
import { RpcTarget } from "cloudflare:workers";
911
import { MessageType } from "../../types.ts";
1012

13+
const STALE_FRAME_PROBE = "stale-frame-probe";
14+
1115
// ── SubAgent: Counter ───────────────────────────────────────────────
1216
// A SubAgent with its own SQLite counter table.
1317

@@ -98,6 +102,10 @@ export class CounterSubAgent extends Agent {
98102
return rows.length > 0 ? rows[0].value : 0;
99103
}
100104

105+
onMessage(_connection: Connection, message: WSMessage) {
106+
if (message === STALE_FRAME_PROBE) this.increment(STALE_FRAME_PROBE);
107+
}
108+
101109
ping(): string {
102110
return "pong";
103111
}
@@ -1427,6 +1435,44 @@ export class TestSubAgentParent extends Agent {
14271435
await this.deleteSubAgent(CounterSubAgent, subAgentName);
14281436
}
14291437

1438+
/**
1439+
* Deletes the child, recreates it under the same name, then forwards a
1440+
* frame from the socket the delete closed, as a late event would. Returns
1441+
* how many of those frames the replacement received.
1442+
*/
1443+
async subAgentForwardStaleFrameToReplacement(
1444+
subAgentName: string
1445+
): Promise<number> {
1446+
const sockets = (
1447+
this as unknown as {
1448+
_webSockets: { getConnections(): Iterable<Connection> };
1449+
}
1450+
)._webSockets.getConnections();
1451+
const connection = [...sockets].find((candidate) => {
1452+
const outerUrl = this._unsafe_getConnectionFlag(
1453+
candidate,
1454+
"_cf_subAgentOuterUrl"
1455+
);
1456+
return (
1457+
typeof outerUrl === "string" &&
1458+
outerUrl.includes(`/sub/counter-sub-agent/${subAgentName}`)
1459+
);
1460+
});
1461+
if (!connection) throw new Error("no socket for the sub-agent");
1462+
1463+
await this.deleteSubAgent(CounterSubAgent, subAgentName);
1464+
const replacement = await this.subAgent(CounterSubAgent, subAgentName);
1465+
await (
1466+
this as unknown as {
1467+
_cf_forwardSubAgentWebSocketMessage(
1468+
connection: Connection,
1469+
message: WSMessage
1470+
): Promise<boolean>;
1471+
}
1472+
)._cf_forwardSubAgentWebSocketMessage(connection, STALE_FRAME_PROBE);
1473+
return replacement.get(STALE_FRAME_PROBE);
1474+
}
1475+
14301476
async subAgentScheduleDelayed(
14311477
subAgentName: string,
14321478
delaySeconds: number,

‎packages/agents/src/tests/sub-agent.test.ts‎

Lines changed: 14 additions & 0 deletions

0 commit comments

Comments
 (0)