fix: pass client tool schemas during tool continuation#882
Conversation
When a client tool result triggers server auto-continuation via `autoContinueAfterToolResult`, the `onChatMessage` callback was called without `clientTools`. This caused the LLM continuation to fail when trying to call another client tool, since the tool schemas were missing from the `streamText` call, resulting in `tool-output-error`. Fix by storing client tool schemas from the initial chat request on the agent instance (`_lastClientTools`) and passing them to `onChatMessage` during tool continuations. Also add a `toolResolutionTrigger` state counter on the client side to force the tool resolution effect to re-run after completing a batch. This handles the race condition where new tool calls arrive via server auto-continuation while `isResolvingToolsRef` is true, preventing the effect from picking them up.
🦋 Changeset detectedLatest commit: d39ab84 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 |
|
/bonk review this PR |
commit: |
![ask-bonk[bot]](https://avatars.githubusercontent.com/in/2454580?s=60&v=4){kind=small}
There was a problem hiding this comment.
Good fix overall — the root cause analysis is correct and both changes address real bugs. The server-side _lastClientTools fix is the core of the PR and directly resolves the tool-output-error on continuation. The client-side toolResolutionTrigger is a clean solution to the ref-not-triggering-rerender problem.
A few things to consider:
1. _lastClientTools is never cleared
When chat is cleared via CF_AGENT_CHAT_CLEAR (line 424), many internal state fields are reset (_activeStreamId, _activeRequestId, _streamChunkIndex, messages, etc.), but _lastClientTools is not. This means after a chat clear, a subsequent request that omits clientTools will still have the stale schemas available for continuations. Additionally, if a client sends a request without client tools (e.g. tools were removed from the config), the conditional if (clientTools && clientTools.length > 0) means the old schemas persist. Consider resetting _lastClientTools on chat clear and also handling the case where a new request arrives with no client tools:
// On chat clear (line 424 block):
this._lastClientTools = undefined;
// On chat request:
this._lastClientTools = clientTools?.length ? clientTools : undefined;2. Multi-connection staleness
Since AIChatAgent is a Durable Object, multiple WebSocket connections share the same _lastClientTools. If connection A sends tools [foo, bar] and connection B sends tools [baz], a continuation triggered by connection A would use connection B's schemas. This is a pre-existing architectural concern (not introduced by this PR), but worth noting since _lastClientTools makes it more visible. A per-connection map would be more correct, though may be overkill for now.
3. Changeset needed
The changeset bot flagged this — since this is a bug fix in @cloudflare/ai-chat, a patch changeset should be added.
The toolResolutionTrigger approach on the client side is clean. Using a state counter to force re-render after the ref-based guard is reset is a well-established React pattern. No issues there.
| // Store client tools for use during tool continuations | ||
| if (clientTools && clientTools.length > 0) { | ||
| this._lastClientTools = clientTools; | ||
| } |
There was a problem hiding this comment.
_lastClientTools is only ever set, never cleared. This means:
- If a subsequent chat request arrives with no
clientTools(e.g. the client removed tools from its config), the stale schemas from a previous request will persist and be used during continuations. - When chat is cleared (
CF_AGENT_CHAT_CLEARhandler at line 424),_lastClientToolsis not reset alongside_activeStreamId,messages, etc.
Consider:
| } | |
| // Store client tools for use during tool continuations | |
| this._lastClientTools = clientTools?.length ? clientTools : undefined; |
And in the CF_AGENT_CHAT_CLEAR handler, add this._lastClientTools = undefined; alongside the other resets.
|
@alexanderjacobsen can you please add the changeset and address valid bonk review comments? |
- Clear _lastClientTools when client sends no tools (instead of keeping stale schemas) - Reset _lastClientTools on chat clear alongside other state - Add patch changeset for @cloudflare/ai-chat
|
@whoiskatrin Should be done now |
| abortSignal | ||
| abortSignal, | ||
| clientTools: this._lastClientTools | ||
| } |
There was a problem hiding this comment.
This is the core fix and it looks correct. One thing I notice: the initial request at line 403-410 also passes body (the custom body fields) to onChatMessage, but the continuation here only passes abortSignal and clientTools. If a user's onChatMessage implementation relies on options.body for model selection, temperature, or other custom metadata, those would be lost during continuation.
This is arguably a pre-existing issue (the body was never passed during continuation even before this PR), but since you're already storing _lastClientTools for this purpose, it might be worth also storing the custom body and forwarding it here for consistency. Not a blocker — just flagging it as a follow-up consideration.
| } | |
| abortSignal, | |
| clientTools: this._lastClientTools, | |
| // TODO: consider also forwarding stored custom body for continuation parity |
|
Review posted successfully. Here's a summary of what I did: PR #882 Review — fix: pass client tool schemas during tool continuation I posted a top-level review summary and 3 inline comments:
|
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as duplicate.
This comment was marked as duplicate.
|
@alexanderjacobsen can we add some tests for this? otherwise happy to ship in the next release |
- Test that client tools are passed to onChatMessage during auto-continuation - Test that stored client tools are cleared on chat clear - Test that stored client tools are cleared when a new request has no tools
|
@whoiskatrin Added tests now |
…fter reconnect Client tools are sent with each prompt, but _lastClientTools is in-memory only and lost on DO restart. This breaks tool confirmation flows: if the user refreshes before approving a tool, the continuation calls onChatMessage without tool schemas and the AI says it doesn't have the tools. Fix: include clientTools in the CF_AGENT_TOOL_RESULT message. The client is the source of truth for its tools and always has current schemas, so re-sending them is both simpler and more correct than persisting server-side. Related to cloudflare#882.
…fter reconnect Client tools are sent with each prompt, but _lastClientTools is in-memory only and lost on DO restart. This breaks tool confirmation flows: if the user refreshes before approving a tool, the continuation calls onChatMessage without tool schemas and the AI says it doesn't have the tools. Fix: include clientTools in the CF_AGENT_TOOL_RESULT message. The client is the source of truth for its tools and always has current schemas, so re-sending them is both simpler and more correct than persisting server-side. Related to cloudflare#882.
…fter reconnect (#897) Client tools are sent with each prompt, but _lastClientTools is in-memory only and lost on DO restart. This breaks tool confirmation flows: if the user refreshes before approving a tool, the continuation calls onChatMessage without tool schemas and the AI says it doesn't have the tools. Fix: include clientTools in the CF_AGENT_TOOL_RESULT message. The client is the source of truth for its tools and always has current schemas, so re-sending them is both simpler and more correct than persisting server-side. Related to #882.
4 participants
Summary
autoContinueAfterToolResulttriggers a server continuation, theonChatMessagecallback was called withoutclientTools, causing the LLM to fail withtool-output-errorwhen trying to call another client tool_lastClientTools) and pass them during continuationtoolResolutionTriggeron the client side to handle race conditions where new tool calls arrive while the resolution effect is busyTest plan
autoContinueAfterToolResult: trueand multiple client-side toolstool-output-error