feat: context aware tools #93

ahmedrowaihi · 2025-12-05T12:03:52Z

🎯 Changes

Supporting context in tool execution..

to make tools behaviour context aware, the only way is to reconstruct them everytime you want to invoke a new trigger..
this PR suggests supporting context.

✅ Checklist

I have followed the steps in the Contributing guide.
I have tested this code locally with pnpm run test:pr.
I have tested this code locally with pnpm run test:lib.

🚀 Release Impact

This change affects published code, and I have generated a changeset.
This change is docs/CI/dev-only (no release).

Summary by CodeRabbit

Release Notes

New Features
- Tools now support shared context during execution, enabling stateful operations across both client and server environments.
- Chat operations can pass context to tools for enhanced functionality and data access.
Tests
- Added comprehensive test coverage for context propagation in tool execution across client and server scenarios.
- Updated existing tests to validate context parameter handling.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

ahmedrowaihi · 2025-12-05T12:05:41Z

this was my initial thought, feel free to suggest enhancement or if there is a better way

jherr

One small API change, otherwise I love it!

packages/typescript/ai/src/tools/tool-definition.ts

ahmedrowaihi · 2025-12-05T15:32:56Z

hey @jherr , thanks for this suggestion ! yes it's much better to make this prop scalable..
now I think we can later on different PR add Tool Hooks ! that would be cool feature..
for example, compressing tool results hook ? yeah why not

nx-cloud · 2025-12-05T15:35:50Z

View your CI Pipeline Execution ↗ for commit c818617

Command	Status	Duration	Result
`nx affected --targets=test:sherif,test:knip,tes...`	✅ Succeeded	1m 16s	View ↗
`nx run-many --targets=build --exclude=examples/**`	✅ Succeeded	29s	View ↗

☁️ Nx Cloud last updated this comment at 2025-12-05 15:40:42 UTC

packages/typescript/ai-client/src/chat-client.ts

AlemTuzlak · 2025-12-05T15:38:35Z

I have 2 questions on my mind here:

do we want to send the context in between the two
how do we type the tools to have the context they need

pkg-pr-new · 2025-12-05T15:39:26Z

Open in StackBlitz

@tanstack/ai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai@93

@tanstack/ai-anthropic

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-anthropic@93

@tanstack/ai-client

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-client@93

@tanstack/ai-devtools-core

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-devtools-core@93

@tanstack/ai-gemini

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-gemini@93

@tanstack/ai-ollama

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-ollama@93

@tanstack/ai-openai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-openai@93

@tanstack/ai-react

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react@93

@tanstack/ai-react-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react-ui@93

@tanstack/ai-solid

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid@93

@tanstack/ai-solid-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid-ui@93

@tanstack/react-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/react-ai-devtools@93

@tanstack/solid-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/solid-ai-devtools@93

commit: c818617

ahmedrowaihi · 2025-12-05T15:44:12Z

do we want to send the context in between the two

hey @AlemTuzlak ..
was just removing this, seemed stale code...

.server() will have server context
.client() will have client context

isomorphic context :'D

ahmedrowaihi · 2025-12-05T15:56:30Z

@AlemTuzlak , do you think we better also pass current serialisable_only within the request ?

ahmedrowaihi · 2025-12-05T16:09:29Z

how do we type the tools to have the context they need

getToolDef.server<{ userId: string }> will make server.options.context.userId typesafe, same goes to client tools.

do we want to send the context in between the two

I think that was bad initial idea, or maybe not and we can make use of it but only serialisable will be sent

ahmedrowaihi · 2025-12-09T12:01:19Z

@jherr @AlemTuzlak anything is missing here 👀 ? would love to fix or enhance if any

coderabbitai · 2025-12-18T20:57:11Z

Walkthrough

This PR introduces context support throughout the TanStack AI framework by adding a generic context parameter to tool definitions, chat options, and execution flows. Context is threaded from chat initialization through tool execution chains on both server and client sides, enabling context-aware tool behavior while maintaining backward compatibility.

Changes

Cohort / File(s)	Summary
Type System & Tool Infrastructure `packages/typescript/ai/src/types.ts`, `packages/typescript/ai/src/tools/tool-definition.ts`, `packages/typescript/ai-client/src/types.ts`	Introduced `ToolOptions<TContext>` interface and added generic context parameter to `Tool`, `ChatOptions`, `ServerTool`, and `ClientTool` interfaces. Updated tool execute signatures to accept `(args, options: ToolOptions<TContext>)` parameters. Added `context?: TContext` property to `ChatOptions` and `ChatClientOptions`.
Chat Engine & Tool Execution `packages/typescript/ai/src/core/chat.ts`, `packages/typescript/ai/src/tools/tool-calls.ts`, `packages/typescript/ai-client/src/chat-client.ts`	Extended `ChatEngine` with private options field for context. Updated `ToolCallManager.executeTools()` and `executeToolCalls()` to accept and propagate context via options parameter. Made `ChatClient` generic over tools and context, passing context to client-side tool execution via `tool.execute(input, { context })`.
Test Suites `packages/typescript/ai/tests/*.test.ts`, `packages/typescript/ai-client/tests/chat-client.test.ts`	Updated existing tests to pass context parameter to tool execute calls. Added new test suites validating context propagation in tool execution and verifying tools work without context (optional behavior).
Example Application `examples/ts-react-chat/src/routes/api.tanchat.ts`	Updated `addToCartToolServer` handler signature to accept `(args, options)` and expanded return value with `quantity` and `totalItems` fields. Introduced `serverContext` object and passed it to chat stream via `context` parameter.

Sequence Diagram

sequenceDiagram
    participant Client
    participant ChatEngine
    participant ToolCallMgr
    participant Tool

    Client->>ChatEngine: new ChatEngine({<br/>context: myContext<br/>})
    Note over ChatEngine: Store options with context
    
    Client->>ChatEngine: processMessage(msg)
    ChatEngine->>ChatEngine: checkForPendingToolCalls()<br/>(pass this.options)
    
    ChatEngine->>ToolCallMgr: executeToolCalls(tools, options)
    Note over ToolCallMgr: Receive options with context
    
    ToolCallMgr->>Tool: execute(args,<br/>{ context: myContext })
    Tool->>Tool: Access context for execution
    Tool-->>ToolCallMgr: result
    ToolCallMgr-->>ChatEngine: toolResults[]
    ChatEngine-->>Client: response
    
    Note over Client,Tool: Context flows through<br/>entire execution chain

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Areas requiring extra attention:

Generic type parameters: The new TContext generics threaded through ToolDefinition, ChatOptions, ChatClient, and related types create interdependencies; verify consistency across all generic bounds and constraints
Execute signature changes: Tool execute callbacks now require a second options parameter across server and client tools—validate that all tool invocations throughout the codebase pass the correct parameter
Tool type inference: The ClientTool union expansion from 2 to 4 type parameters and the AnyClientTool union updates may impact type checking; confirm no regressions in tool type narrowing
Context propagation paths: Multiple execution paths (normal tools, pre-approved tools, client-executed tools) now thread context; ensure no path is missed or incorrectly typed
Test coverage: While existing tests were updated with context parameters, verify that edge cases (missing context, null context, context type mismatches) are adequately covered

Poem

🐰 Through tools we hop, with context so deep,
Each execution now knows what to keep,
From server to client, the context flows true,
With generics and options, we know what to do,
No more lost parameters—our pathway is clear! 🥕✨

Pre-merge checks and finishing touches

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The description includes the required 'Changes' section explaining context support and completed checklists, but the 'Release Impact' section lacks a changeset despite modifying published code.	Clarify whether this requires a changeset by checking if changes are part of published packages. If so, add a changeset using the provided documentation link.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: context aware tools' directly summarizes the main change: adding context-aware capabilities to the tool execution system across multiple packages.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

feat: introduce ToolOptions and update tool execution context fix: update context handling in ChatClient and API route

coderabbitai

Actionable comments posted: 0

🧹 Nitpick comments (5)

packages/typescript/ai/tests/tool-definition.test.ts (1)
63-71: Consider removing as any type assertion.

The { context: undefined as any } assertion works but bypasses type checking. Since ToolOptions<TContext> accepts unknown by default and context is optional, you could simply use { context: undefined } without the assertion.
🔎 Suggested fix
-      const result = await serverTool.execute(
-        { location: 'Paris' },
-        { context: undefined as any },
-      )
+      const result = await serverTool.execute(
+        { location: 'Paris' },
+        { context: undefined },
+      )
       expect(result).toEqual({ temperature: 72, conditions: 'sunny' })
-      expect(executeFn).toHaveBeenCalledWith(
-        { location: 'Paris' },
-        { context: undefined as any },
-      )
+      expect(executeFn).toHaveBeenCalledWith(
+        { location: 'Paris' },
+        { context: undefined },
+      )
examples/ts-react-chat/src/routes/api.tanchat.ts (1)
39-45: Unused options parameter in server tool handler.

The options parameter is declared but not used in the handler. If context isn't needed here currently, consider using _options to indicate intentional non-use, or remove it until needed.
🔎 Suggested change:
-const addToCartToolServer = addToCartToolDef.server((args, options) => ({
+const addToCartToolServer = addToCartToolDef.server((args, _options) => ({
   success: true,
   cartId: 'CART_' + Date.now(),
   guitarId: args.guitarId,
   quantity: args.quantity,
   totalItems: args.quantity,
 }))
packages/typescript/ai-client/tests/chat-client.test.ts (1)
612-624: Simplify the generic type annotation in the mock function.

The generic <TContext = unknown> on the mock function is unusual since the function is immediately specialized for TestContext. Consider simplifying the type annotation.
🔎 Suggested simplification:
       const executeFn = vi.fn(
-        async <TContext = unknown>(
-          _args: any,
-          options: ToolOptions<TContext>,
-        ) => {
-          const ctx = options.context as TestContext
+        async (_args: any, options: ToolOptions<TestContext>) => {
+          const ctx = options.context!
           ctx.localStorage.setItem(
             `pref_${ctx.userId}_${_args.key}`,
             _args.value,
           )
           return { success: true }
         },
       )
packages/typescript/ai-client/src/chat-client.ts (1)

500-533: Consider adding context to updateOptions for dynamic context updates.

The updateOptions method doesn't currently support updating context. If context needs to change during the client's lifetime (e.g., when user session changes), this would require recreating the client. Consider if this is the intended behavior.
packages/typescript/ai/src/tools/tool-definition.ts (1)
197-223: Consider narrowing the as any cast for better type safety.

The execute as any casts at lines 207 and 221 bypass type checking. While this works correctly at runtime since the execute function passed in has the correct signature, it could mask future type mismatches during refactoring.

A slightly safer alternative would be to use a more targeted type assertion that preserves the function shape:
🔎 Suggested improvement
       return {
         __toolSide: 'server',
         __contextType: undefined as TContext,
         ...config,
-        execute: execute as any,
+        execute: execute as Tool<TInput, TOutput, TName>['execute'],
       }
Similarly for the client method. This maintains the function signature while allowing the TContext variance.
That said, given the complexity of the generic variance involved, the current approach is pragmatic and acceptable for a "Chill" review.

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e011bf8 and 8e92c0e.

📒 Files selected for processing (11)

examples/ts-react-chat/src/routes/api.tanchat.ts (3 hunks)
packages/typescript/ai-client/src/chat-client.ts (4 hunks)
packages/typescript/ai-client/src/types.ts (2 hunks)
packages/typescript/ai-client/tests/chat-client.test.ts (3 hunks)
packages/typescript/ai/src/core/chat.ts (5 hunks)
packages/typescript/ai/src/tools/tool-calls.ts (7 hunks)
packages/typescript/ai/src/tools/tool-definition.ts (7 hunks)
packages/typescript/ai/src/types.ts (4 hunks)
packages/typescript/ai/tests/ai-chat.test.ts (6 hunks)
packages/typescript/ai/tests/tool-call-manager.test.ts (1 hunks)
packages/typescript/ai/tests/tool-definition.test.ts (6 hunks)

🧰 Additional context used

📓 Path-based instructions (4)

**/*.{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 /adapters subpath rather than monolithic adapters
Use Zod for runtime schema validation and type inference, particularly for tool input/output definitions with toolDefinition() and Zod schema inference
Implement isomorphic tool system using toolDefinition() 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:

packages/typescript/ai/tests/tool-call-manager.test.ts
packages/typescript/ai/src/tools/tool-calls.ts
packages/typescript/ai/tests/ai-chat.test.ts
packages/typescript/ai-client/src/chat-client.ts
packages/typescript/ai/src/core/chat.ts
examples/ts-react-chat/src/routes/api.tanchat.ts
packages/typescript/ai-client/src/types.ts
packages/typescript/ai/tests/tool-definition.test.ts
packages/typescript/ai-client/tests/chat-client.test.ts
packages/typescript/ai/src/types.ts
packages/typescript/ai/src/tools/tool-definition.ts

**/*.test.ts

📄 CodeRabbit inference engine (CLAUDE.md)

Write unit tests using Vitest alongside source files with .test.ts naming convention

Files:

packages/typescript/ai/tests/tool-call-manager.test.ts
packages/typescript/ai/tests/ai-chat.test.ts
packages/typescript/ai/tests/tool-definition.test.ts
packages/typescript/ai-client/tests/chat-client.test.ts

**/*.{ts,tsx,js,jsx}

📄 CodeRabbit inference engine (CLAUDE.md)

Use camelCase for function and variable names throughout the codebase

Files:

packages/typescript/ai/tests/tool-call-manager.test.ts
packages/typescript/ai/src/tools/tool-calls.ts
packages/typescript/ai/tests/ai-chat.test.ts
packages/typescript/ai-client/src/chat-client.ts
packages/typescript/ai/src/core/chat.ts
examples/ts-react-chat/src/routes/api.tanchat.ts
packages/typescript/ai-client/src/types.ts
packages/typescript/ai/tests/tool-definition.test.ts
packages/typescript/ai-client/tests/chat-client.test.ts
packages/typescript/ai/src/types.ts
packages/typescript/ai/src/tools/tool-definition.ts

examples/**

📄 CodeRabbit inference engine (CLAUDE.md)

Examples are not built by Nx and should be run independently from their directories with pnpm dev or pnpm install && pnpm dev

Files:

examples/ts-react-chat/src/routes/api.tanchat.ts

🧠 Learnings (10)

📓 Common learnings

Learnt from: CR
Repo: TanStack/ai PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-13T17:09:09.794Z
Learning: Implement framework integrations using the headless `tanstack/ai-client` for state management with framework-specific hooks (useChat) on top