feat(chat): add support for client tool execution after approval #228
Conversation
- refactor(stream): improve tool call part updates and message handling - test(chat): add tests for approval flow and full message sending Resolves TanStack#225 by ensuring client-side tools execute immediately after approval and fixing a state overwrite bug in the stream processor.
📝 WalkthroughWalkthroughThis PR implements client-side tool execution upon approval. The chat client now parses tool call arguments when approval is confirmed, executes matching client-side tools with the parsed input, and records execution results (success or error states). A new Changes
Sequence DiagramsequenceDiagram
participant Client as Chat Client
participant Stream as Message Stream
participant Processor as Message Processor
participant Tool as Client Tool
participant Result as Output Recorder
Client->>Stream: Receive tool_call chunk (approval-requested)
Stream->>Client: Update message with approval request
Client->>Client: User triggers addToolApprovalResponse()
Client->>Processor: Parse toolCallPart.arguments
Processor-->>Client: foundToolInput (JSON parsed)
Client->>Tool: Execute client tool with foundToolInput
alt Tool succeeds
Tool-->>Client: Return tool result
Client->>Result: Record output with state: 'output-available'
else Tool fails
Tool-->>Client: Throw error
Client->>Result: Record output with state: 'output-error' + errorText
end
Result-->>Client: Return early, skip continuation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/typescript/ai-client/src/chat-client.ts (1)
419-434: Handle JSON parse failure more explicitly.When
JSON.parse(toolCallPart.arguments)fails,foundToolInputremainsundefined, but the tool still executes on line 455 with potentially undefined input. This may cause unexpected behavior in client tools that expect specific input shapes.Consider setting a fallback or logging a warning:
♻️ Suggested improvement
foundToolName = toolCallPart.name try { foundToolInput = JSON.parse(toolCallPart.arguments) } catch { - // Ignore parse errors + // Log warning but continue - tool will receive undefined input + console.warn(`Failed to parse arguments for tool ${toolCallPart.name}`) }Alternatively, skip tool execution if input parsing fails, as the tool likely cannot function correctly without its expected input.
🧹 Nitpick comments (5)
packages/typescript/ai-client/src/connection-adapters.ts (1)
312-324: Consider addingsendFullMessagessupport tostream()andrpcStream()adapters for consistency.The
stream()andrpcStream()adapters always convert toModelMessages, whilefetchServerSentEventsandfetchHttpStreamnow support thesendFullMessagesoption. If UIMessage metadata (like tool approval state) needs to flow through these adapters in the future, they would need similar updates.This may be intentional if these adapters are only used in contexts where ModelMessages suffice, but worth noting for future extensibility.
Also applies to: 342-356
packages/typescript/ai-client/tests/chat-client-approval.test.ts (3)
67-67: Type cast suggests incomplete type definitions.The
as anycast indicates that theApprovalRequestedStreamChunktype may not fully align with what's being constructed here. Consider updating the type definitions in@tanstack/aito avoid the cast, or add a comment explaining why it's necessary.
82-132: Consider adding test cases for rejection and error scenarios.The current test validates the happy path (approved = true). For comprehensive coverage of the linked issue
#225, consider adding tests for:
- Rejection flow:
approved: falseshould not execute the tool- Tool execution error: Mock
executeto throw and verifyoutput-errorstate is recorded- Missing tool: Approve a tool that isn't registered in
toolsThese scenarios ensure the fix is robust against edge cases.
Would you like me to generate additional test cases for these scenarios?
116-116: Usevi.waitFor()to wait for assertions instead of arbitrary timeouts.Hardcoded 100ms delays are unreliable across environments and can cause test flakiness. Vitest's
vi.waitFor()polls the callback until the condition is met (or timeout at 1000ms), making tests more robust:Example refactoring
// Line 128: Instead of: await new Promise((resolve) => setTimeout(resolve, 100)) expect(execute).toHaveBeenCalledWith(input) // Use: await vi.waitFor(() => { expect(execute).toHaveBeenCalledWith(input) })This approach also applies to line 116 where a condition-based wait would be more reliable than a fixed delay.
packages/typescript/ai-client/src/chat-client.ts (1)
463-471: Typeerrormore precisely thanany.The
error: anycatch clause loses type information. Consider usingunknownwith type narrowing:♻️ Suggested improvement
- } catch (error: any) { + } catch (error: unknown) { await this.addToolResult({ toolCallId: foundToolCallId, tool: foundToolName, output: null, state: 'output-error', - errorText: error.message, + errorText: error instanceof Error ? error.message : String(error), }) return }This provides safer error handling and aligns with TypeScript best practices.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
packages/typescript/ai-client/src/chat-client.tspackages/typescript/ai-client/src/connection-adapters.tspackages/typescript/ai-client/tests/chat-client-approval.test.tspackages/typescript/ai-client/tests/connection-adapters.test.tspackages/typescript/ai/src/activities/chat/index.tspackages/typescript/ai/src/activities/chat/stream/message-updaters.tspackages/typescript/ai/tests/ai-text.test.tstesting/panel/src/routes/api.chat.tstesting/panel/src/routes/index.tsx
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.{ts,tsx}: Use tree-shakeable adapter architecture for provider implementations - export specialized adapters (text, embedding, summarize, image) as separate imports from/adapterssubpath rather than monolithic adapters
Use Zod for runtime schema validation and type inference, particularly for tool input/output definitions withtoolDefinition()and Zod schema inference
Implement isomorphic tool system usingtoolDefinition()with.server()and.client()implementations for dual-environment execution
Use type-safe per-model configuration with provider options typed based on selected model to ensure compile-time safety
Implement stream processing with StreamProcessor for handling chunked responses and support partial JSON parsing for streaming AI responses
Files:
testing/panel/src/routes/api.chat.tspackages/typescript/ai/tests/ai-text.test.tspackages/typescript/ai-client/src/connection-adapters.tspackages/typescript/ai-client/tests/chat-client-approval.test.tspackages/typescript/ai-client/tests/connection-adapters.test.tstesting/panel/src/routes/index.tsxpackages/typescript/ai-client/src/chat-client.tspackages/typescript/ai/src/activities/chat/stream/message-updaters.tspackages/typescript/ai/src/activities/chat/index.ts
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (CLAUDE.md)
Use camelCase for function and variable names throughout the codebase
Files:
testing/panel/src/routes/api.chat.tspackages/typescript/ai/tests/ai-text.test.tspackages/typescript/ai-client/src/connection-adapters.tspackages/typescript/ai-client/tests/chat-client-approval.test.tspackages/typescript/ai-client/tests/connection-adapters.test.tstesting/panel/src/routes/index.tsxpackages/typescript/ai-client/src/chat-client.tspackages/typescript/ai/src/activities/chat/stream/message-updaters.tspackages/typescript/ai/src/activities/chat/index.ts
**/*.test.ts
📄 CodeRabbit inference engine (CLAUDE.md)
Write unit tests using Vitest alongside source files with
.test.tsnaming convention
Files:
packages/typescript/ai/tests/ai-text.test.tspackages/typescript/ai-client/tests/chat-client-approval.test.tspackages/typescript/ai-client/tests/connection-adapters.test.ts
🧠 Learnings (4)
📚 Learning: 2025-12-13T17:09:09.794Z
Learnt from: CR
Repo: TanStack/ai PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-13T17:09:09.794Z
Learning: Applies to packages/typescript/*/src/adapters/*.ts : Create individual adapter implementations for each provider capability (text, embed, summarize, image) with separate exports to enable tree-shaking
Applied to files:
packages/typescript/ai-client/src/connection-adapters.tspackages/typescript/ai-client/tests/connection-adapters.test.tspackages/typescript/ai/src/activities/chat/index.ts
📚 Learning: 2025-12-13T17:09:09.794Z
Learnt from: CR
Repo: TanStack/ai PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-13T17:09:09.794Z
Learning: Applies to **/*.{ts,tsx} : Implement isomorphic tool system using `toolDefinition()` with `.server()` and `.client()` implementations for dual-environment execution
Applied to files:
testing/panel/src/routes/index.tsx
📚 Learning: 2025-12-13T17:09:09.794Z
Learnt from: CR
Repo: TanStack/ai PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-13T17:09:09.794Z
Learning: Applies to **/*.{ts,tsx} : Implement stream processing with StreamProcessor for handling chunked responses and support partial JSON parsing for streaming AI responses
Applied to files:
packages/typescript/ai-client/src/chat-client.tspackages/typescript/ai/src/activities/chat/index.ts
📚 Learning: 2025-12-13T17:09:09.794Z
Learnt from: CR
Repo: TanStack/ai PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-13T17:09:09.794Z
Learning: Applies to packages/typescript/*/src/index.ts : Export tree-shakeable adapters with clear subpath exports in package.json (e.g., `tanstack/ai/adapters`, `tanstack/ai-openai/adapters`) to minimize bundle size
Applied to files:
packages/typescript/ai/src/activities/chat/index.ts
🧬 Code graph analysis (7)
testing/panel/src/routes/api.chat.ts (3)
examples/ts-solid-chat/src/lib/guitar-tools.ts (1)
addToCartToolServer(109-117)examples/ts-react-chat/src/lib/guitar-tools.ts (1)
addToCartToolDef(69-84)testing/panel/src/lib/guitar-tools.ts (1)
addToCartToolDef(69-84)
packages/typescript/ai/tests/ai-text.test.ts (3)
packages/typescript/ai/src/types.ts (2)
TextOptions(565-650)StreamChunk(740-748)packages/typescript/ai-client/src/connection-adapters.ts (1)
stream(312-324)packages/typescript/ai/src/activities/chat/index.ts (1)
chat(975-998)
packages/typescript/ai-client/tests/chat-client-approval.test.ts (3)
packages/typescript/ai/src/types.ts (1)
StreamChunk(740-748)packages/typescript/ai-client/src/connection-adapters.ts (1)
stream(312-324)packages/typescript/ai-client/src/chat-client.ts (1)
ChatClient(17-583)
packages/typescript/ai-client/tests/connection-adapters.test.ts (2)
packages/typescript/ai-client/src/connection-adapters.ts (1)
fetchServerSentEvents(134-200)packages/typescript/ai-client/src/index.ts (1)
fetchServerSentEvents(23-23)
testing/panel/src/routes/index.tsx (1)
packages/typescript/ai-client/src/connection-adapters.ts (1)
fetchServerSentEvents(134-200)
packages/typescript/ai-client/src/chat-client.ts (1)
packages/typescript/ai-svelte/src/create-chat.svelte.ts (1)
messages(121-123)
packages/typescript/ai/src/activities/chat/index.ts (3)
packages/typescript/ai/src/types.ts (4)
ModelMessage(232-243)UIMessage(293-298)ToolCall(87-94)StreamChunk(740-748)packages/typescript/ai/src/index.ts (1)
convertMessagesToModelMessages(81-81)packages/typescript/ai-client/src/connection-adapters.ts (1)
stream(312-324)
🔇 Additional comments (17)
testing/panel/src/routes/api.chat.ts (1)
47-56: LGTM! Consider removing debug logging before production.The refactor from arrow function to block-bodied function is clean and preserves the same return structure. The
console.logis useful for debugging the tool approval flow during development.Consider removing or conditionalizing the
console.logstatement before production deployment to avoid noisy server logs.testing/panel/src/routes/index.tsx (1)
448-450: LGTM! Enables full UIMessage transmission for tool approval workflow.The
sendFullMessages: trueoption is essential for the tool approval flow. This ensures that when messages contain tool-call parts with approval state (e.g.,approval-responded), the fullUIMessageobjects—including thepartsarray—are sent to the server, allowing it to process pending tool executions correctly.packages/typescript/ai/src/activities/chat/stream/message-updaters.ts (2)
73-80: Critical fix: Preserves existing tool call fields during updates.The spread operator
...existingensures that fields likeapprovalandoutputare preserved when updating tool call parts. This is essential for the tool approval workflow—without this fix, updating thestatewould have inadvertently wiped out the approval metadata, causing the approval flow to break.
83-90: LGTM!Clean refactor to create the
toolCallPartobject explicitly before pushing. The explicit type annotation improves code clarity.packages/typescript/ai-client/tests/connection-adapters.test.ts (2)
321-364: LGTM! Good test coverage for default message conversion behavior.This test correctly validates that without
sendFullMessages, the connection adapter convertsUIMessageobjects toModelMessageformat, stripping thepartsarray. The assertion at line 360-361 properly verifies the conversion.
366-420: LGTM! Validates full UIMessage transmission.This test correctly validates that with
sendFullMessages: true, the fullUIMessageobjects are sent including thepartsarray with tool-call approval metadata. This is essential for the tool approval workflow.packages/typescript/ai/tests/ai-text.test.ts (1)
1476-1554: LGTM! Good test coverage for UIMessage-based tool execution.This test validates the new flow where approved tool calls are extracted from
UIMessage.partsrather than the legacytoolCallsarray. This is essential for the tool approval workflow where the client sends back UIMessages withapproval-respondedstate.The test correctly verifies:
- Tool execution with parsed arguments from the
partsarray- Tool result message is passed to the adapter
- Only one adapter call (tool executed before streaming)
packages/typescript/ai-client/src/connection-adapters.ts (2)
92-98: LGTM! Well-documented option for advanced message handling.The
sendFullMessagesoption is clearly documented and provides the flexibility needed for tool approval flows that require UIMessage metadata.
147-149: LGTM! Clean conditional message transformation.The ternary logic correctly chooses between passing messages as-is or converting them based on the
sendFullMessagesflag.packages/typescript/ai-client/tests/chat-client-approval.test.ts (1)
6-12: LGTM! Clean mock adapter factory.The helper correctly wraps the
stream()adapter to yield predetermined chunks for testing.packages/typescript/ai-client/src/chat-client.ts (2)
296-312: LGTM! Correctly passes messages to the connection adapter.The change to use
this.processor.getMessages()and pass messages directly aligns with the newsendFullMessagescapability in connection adapters, allowing UIMessage metadata to flow through when needed.
450-474: LGTM! Core fix for issue#225- client tool execution on approval.This block correctly:
- Checks if the response is approved and tool details are available
- Looks up the client tool by name
- Executes the tool and records the result via
addToolResult- Uses early returns to avoid duplicate continuation checks
The structure mirrors the
onToolCallhandler (lines 130-155), maintaining consistency.packages/typescript/ai/src/activities/chat/index.ts (5)
211-211: LGTM! Type expansion enables UIMessage support in the engine.Expanding the internal
messagestype to accept bothModelMessageandUIMessageallows the engine to work with UI-originated messages while maintaining proper conversion at API boundaries.
248-250: LGTM! Consistent conversion at public API boundary.The
getMessages()method correctly converts toModelMessage[]to maintain a stable public API while the internal state can hold either message type.
723-731: Good defensive checks for approval and output properties.The added null checks (
part.approval &&andpart.output !== undefined) prevent potential runtime errors when processing messages that may not have these optional properties populated.
847-868: LGTM! Consistent message conversion for tool call detection.Converting messages at the start of
getPendingToolCallsFromMessages()ensures the filtering and iteration logic operates on a uniformModelMessage[]structure.
1003-1017: LGTM! Helper functions updated for flexible message handling.The function signatures and internal logic properly handle the generic types while providing sensible defaults (
messages ?? []) when messages are not provided.Also applies to: 1023-1030, 1038-1055
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
1 participant
Resolves #225 by ensuring client-side tools execute immediately after approval and fixing a state overwrite bug in the stream processor.
🎯 Changes
✅ Checklist
pnpm run test:pr.🚀 Release Impact
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.