Repository navigation
Fix RPC resource leaks in workflows - #1668
Conversation
The Workers runtime warns about un-disposed Cap'n Proto resources when RPC methods return references. Failing to dispose these references causes memory leaks and noisy runtime warnings. This ensures all RPC results and Durable Object stubs created by AgentWorkflow and ThinkWorkflow are properly disposed when they are no longer needed. This includes wait events in the approval flow, submission and wait events during prompts, and the Agent Durable Object stub itself. Regression tests are included to verify disposal and ensure payloads can still be parsed correctly after disposal.
🦋 Changeset detectedLatest commit: 511478b The changes in this PR will be included in the next version bump. This PR includes changesets to release 2 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
agents
@cloudflare/ai-chat
@cloudflare/codemode
hono-agents
@cloudflare/shell
@cloudflare/think
@cloudflare/voice
@cloudflare/worker-bundler
commit: |
|
/bonk do a review on this, look for edge cases and things we might not have thought of, is the ux/dx alright, do we have enough test coverage |
![ask-bonk[bot]](https://avatars.githubusercontent.com/in/2454580?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review: RPC Resource Leak Fixes in Workflows
I went through the full diff, traced every dispose path, and analyzed the test coverage. Overall this is a solid, well-scoped fix. Here are my findings:
What's good
-
Correct execution ordering for
_disposeAgent()— Thefinallyblock in the wrappedrun()only fires after_runWithErrorReportingcompletes (including the_autoReportErrorcall inside its catch). So the agent stub is alive for error reporting and disposed after. This is the trickiest part of the change and it's handled correctly. -
waitForApprovaltry/finally — The event fromstep.waitForEvent()is properly disposed on both the approval and rejection paths. Thefinallyensures cleanup even whenWorkflowRejectedErroris thrown. -
ThinkWorkflow
submitMessagesdisposal — Stripping the RPC stub down to{ submissionId }inside thestep.do()callback before disposing is the right pattern. The serializable plain object survives workflow replay; the RPC stub doesn't leak. -
ThinkWorkflow
waitForEventresult disposal — Same try/finally pattern as the agents package. The event stays alive during validation/parsing and gets disposed after. -
_disposeAgentis idempotent — Setting_agent = undefinedbefore callingdisposeIfPresentmeans double-calls are safe. The test confirms this.
Issues and edge cases
1. broadcastToClients() after _disposeAgent() — user-code footgun
User code that stores a reference to this (the workflow instance) and calls this.broadcastToClients() from an async callback that resolves after run() returns will hit the "Agent not initialized" error. This was always possible in theory, but now that _disposeAgent() actively clears the stub, it's more likely to surface.
This is arguably correct behavior (you shouldn't use the agent after run() completes), but it may surprise users. A brief note in the broadcastToClients JSDoc that it's only valid during run() would help. Not a blocker.
2. Duplicated disposeIfPresent / isDisposableResource utilities
The exact same DisposableResource type, isDisposableResource type guard, and disposeIfPresent function are copy-pasted in both packages/agents/src/workflows.ts and packages/think/src/workflows.ts. Since ThinkWorkflow already imports from agents/workflows, it would be cleaner to export disposeIfPresent from agents/workflows and reuse it.
The existing DisposableStore in packages/agents/src/core/events.ts uses a different convention (.dispose() method vs Symbol.dispose), so it can't be reused directly, but the utility functions could be shared.
3. No test for _disposeAgent() being called when _initAgent itself throws
If _initAgent throws (e.g., missing binding, bad agent name), the try/finally in the wrapped run() will call _disposeAgent() with _agent still undefined (since the assignment at the end of _initAgent never ran). This is safe because disposeIfPresent(undefined) is a no-op, but there's no test covering this path.
4. No test for the super.run() path with disposal
The wrapped run() has a second code path (lines 161-167) for when __agentInitCalled is already true (i.e., called via super.run()). This path does NOT run _disposeAgent() — only the outer caller does. This is correct, but there's no test validating that the agent stub survives through a super.run() call and is only disposed once at the top level.
5. ThinkWorkflow cancelSubmission in timeout path — agent stub may be disposed before cancellation completes
In _waitForPromptEvent (think/workflows.ts:204-221), when waitForEvent times out, the catch block calls this.agent.cancelSubmission(...). This RPC call goes through the agent stub. If the timeout error propagates up and the outer finally in _promptStep tries to dispose the event (which doesn't exist since waitForEvent threw), and then the error continues up to the AgentWorkflow.run() wrapper which calls _disposeAgent() — the question is whether cancelSubmission has completed by then.
Looking at the flow: _waitForPromptEvent throws ThinkPromptTimeoutError → caught by _promptStep? No, it's not caught there, it propagates up → caught by _runWithErrorReporting → _autoReportError runs → then finally { _disposeAgent() }. Since cancelSubmission is awaited inside _waitForPromptEvent, it completes before the throw. So this is actually fine. But it's worth noting that cancelSubmission uses the agent stub that will be disposed shortly after — if cancellation itself triggers a long-running operation, the stub gets disposed while that operation is in-flight. This is probably fine in practice since cancelSubmission is awaited.
Test coverage assessment
The tests are well-structured unit tests that validate the disposal mechanics using Object.create on the prototype (avoiding the need for a real ExecutionContext). Coverage includes:
- Agent stub disposal + idempotency
waitForApprovalapproval path + rejection path disposal- ThinkWorkflow submission result disposal
- ThinkWorkflow wait event disposal
- ThinkWorkflow output read-before-dispose ordering
- ThinkWorkflow timeout path with cancellation
Missing coverage:
_initAgentfailure path (not a big deal, it's a no-op dispose)super.run()inheritance path (the prototype tests check inheritance but not the disposal-during-super.run scenario)- Error in
_disposeAgentitself (what ifSymbol.disposethrows?) — currently swallowed silently which is fine, but untested - Integration test that actually verifies "RPC stub was not disposed" warnings are gone (though this may be impractical to test in vitest)
Summary
The fix is correct and well-targeted. The execution ordering of dispose-vs-error-reporting is the most critical part and it's handled properly. The duplication of helper functions is minor but worth cleaning up. Test coverage is good for the new paths but a couple of edge cases (init failure, super.run disposal) are untested. None of the gaps are blockers.
| @@ -34,7 +34,11 @@ | |||
| */ | |||
There was a problem hiding this comment.
Nit: This utility triad (DisposableResource, isDisposableResource, disposeIfPresent) is duplicated verbatim in packages/think/src/workflows.ts. Since ThinkWorkflow already imports from agents/workflows, consider exporting disposeIfPresent from the agents package and reusing it here. Would reduce the maintenance surface and keep the dispose semantics in one place.
| this._agent = undefined; | ||
| this.__agentInitCalled = false; | ||
| disposeIfPresent(agent); | ||
| } |
There was a problem hiding this comment.
Minor: If Symbol.dispose on the agent stub throws, the error propagates out of the finally block and masks the original error from run(). Unlikely in practice (CF stubs dispose cleanly), but wrapping in try/catch for consistency with the best-effort pattern used elsewhere would make it bulletproof:
private _disposeAgent(): void {
const agent = this._agent;
this._agent = undefined;
this.__agentInitCalled = false;
try {
disposeIfPresent(agent);
} catch {
// best-effort cleanup
}
}| } | ||
|
|
||
| function disposeIfPresent(value: unknown): void { | ||
| if (isDisposableResource(value)) { |
There was a problem hiding this comment.
These three helper functions (DisposableResource, isDisposableResource, disposeIfPresent) are identical to the ones in packages/agents/src/workflows.ts. Since ThinkWorkflow already imports from agents/workflows (line 10), consider exporting disposeIfPresent from there and importing it here to avoid the duplication.
|
Review posted on PR #1668. Here's a summary of the findings: Overall verdict: Solid fix, no blockers. The PR correctly fixes RPC resource leaks in
Key edge cases I flagged:
|
AgentWorkflow and ThinkWorkflow were not disposing the Durable Object references returned by waitForEvent(), submitMessages(), and getAgentByName(). These references hold an open connection to the remote object and must be explicitly disposed; without it, the runtime logs "RPC stub was not disposed" warnings and the connections leak for the duration of the workflow run. Because workflows can be long-lived and may retry, the leaks compound over time.
Regression tests are included for the approval flow, the Think prompt flow, and the agent stub cleanup.