fix(ai): finish middleware runs on client tool waits#996
Open
JsonKim wants to merge 3 commits into
Open
Conversation
Call terminal middleware hooks when a server invocation pauses for a client tool, ensuring OpenTelemetry spans and duration metrics are finalized. Defer structured-output finalization until the client result is submitted in the next invocation, with unit and E2E regression coverage.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe chat engine now completes middleware terminal hooks during client-tool waits and postpones structured-output finalization until client results arrive. Unit and end-to-end tests cover lifecycle hooks, telemetry, deferred output, and browser-side tool resolution. ChangesClient-tool wait lifecycle
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Browser
participant TextEngine
participant ClientTool
participant Middleware
Browser->>TextEngine: start structured chat
TextEngine->>ClientTool: request client context
TextEngine->>Middleware: run terminal hook
TextEngine-->>Browser: wait interrupt
Browser->>ClientTool: resolve client tool
ClientTool-->>TextEngine: return context
TextEngine-->>Browser: structured output complete
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://cdn.github.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
packages/ai/tests/chat-client-tool-wait.test.ts (1)
83-174: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding an approval-wait analog to these onFinish-once tests.
The changeset/PR explicitly covers hooks finishing "while waiting for client tools or approvals," but this suite only exercises the client-tool-execution wait path (
needsClientExecution).runTerminalHook()'s dispatch is shared with theneedsApprovalbranch inemitActionableInterruptBoundary, so a matching test (e.g. a tool withneedsApproval: trueawaiting an approval decision) would close the coverage gap for the other half of the documented fix.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai/tests/chat-client-tool-wait.test.ts` around lines 83 - 174, Add a matching approval-wait lifecycle test alongside the client-tool wait tests, using a tool configured with needsApproval: true and an unresolved approval decision. Assert that the terminal middleware’s onFinish is called exactly once and validates the expected approval-wait finish metadata, reusing the existing helpers such as createMockAdapter, collectChunks, and expectOnlyFinish.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/ai/src/activities/chat/index.ts`:
- Around line 1071-1075: Update the pendingPhase guard in the agent loop to
return after runTerminalHook() when checkForPendingToolCalls() yields either
"stop" or "wait". Preserve the existing terminal handling while preventing
streamModelResponse() from running after a "stop" phase.
In `@testing/e2e/tests/client-tool-wait.spec.ts`:
- Around line 49-71: Update the telemetry assertions in the client-tool wait
test to poll the middleware-test capture endpoint until the expected root span,
client span, and duration record are finalized. Replace the one-shot capture
read with retry logic that waits for ended spans and the histogram before
asserting the exact counts, while preserving the existing testId validation and
response checks.
---
Nitpick comments:
In `@packages/ai/tests/chat-client-tool-wait.test.ts`:
- Around line 83-174: Add a matching approval-wait lifecycle test alongside the
client-tool wait tests, using a tool configured with needsApproval: true and an
unresolved approval decision. Assert that the terminal middleware’s onFinish is
called exactly once and validates the expected approval-wait finish metadata,
reusing the existing helpers such as createMockAdapter, collectChunks, and
expectOnlyFinish.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a3d4dc25-112c-4c0b-b05f-b9489b616d29
📒 Files selected for processing (8)
.changeset/fuzzy-tools-wait.mdpackages/ai/src/activities/chat/index.tspackages/ai/tests/chat-client-tool-wait.test.tstesting/e2e/fixtures/middleware-test/structured-client-tool-wait.jsontesting/e2e/src/lib/middleware-test-tools.tstesting/e2e/src/routes/api.middleware-test.tstesting/e2e/src/routes/middleware-test.tsxtesting/e2e/tests/client-tool-wait.spec.ts
A restored or retried conversation can start with an unresolved approval or client tool call in its message history. If middleware fails while processing the resulting RUN_FINISHED interrupt boundary, the engine emits RUN_ERROR and returns the pending "stop" phase. Because run() previously handled only "wait", it entered the agent loop and made another model request after the terminal error. Reproduction: 1. Start chat() with an unresolved tool call in the initial messages. 2. Configure the tool to require approval or client execution. 3. Throw from middleware while processing the interrupt RUN_FINISHED event. 4. Consume the stream and observe RUN_ERROR followed by an adapter call. Return immediately for both "stop" and "wait", and add regression coverage verifying that RUN_ERROR invokes onError without another provider call.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Call terminal middleware hooks when a server invocation pauses for a client tool, ensuring OpenTelemetry spans and duration metrics are finalized.
Defer structured-output finalization until the client result is submitted in the next invocation, with unit and E2E regression coverage.
🎯 Changes
onAbortandonErrorlifecycle hooks.Tests
pnpm --filter @tanstack/ai test:lib --run tests/chat-client-tool-wait.test.tspnpm --filter @tanstack/ai-e2e test:e2e -- tests/client-tool-wait.spec.ts✅ Checklist
pnpm run test:pr.🚀 Release Impact
Summary by CodeRabbit
Bug Fixes
Tests