diff --git a/packages/typescript/ai/src/activities/chat/index.ts b/packages/typescript/ai/src/activities/chat/index.ts index c7ec866d6..2dce0e9cc 100644 --- a/packages/typescript/ai/src/activities/chat/index.ts +++ b/packages/typescript/ai/src/activities/chat/index.ts @@ -1084,7 +1084,14 @@ class TextEngine< const chunks: Array = [] for (const result of results) { - const content = JSON.stringify(result.result) + const fullContent = JSON.stringify(result.result) + + // Apply clientOutput filter if the tool defines one + const tool = this.tools.find((t) => t.name === result.toolName) + const clientContent = + tool?.clientOutput && result.state !== 'output-error' + ? JSON.stringify(tool.clientOutput(result.result)) + : fullContent chunks.push({ type: 'TOOL_CALL_END', @@ -1092,14 +1099,14 @@ class TextEngine< model: finishEvent.model, toolCallId: result.toolCallId, toolName: result.toolName, - result: content, + result: clientContent, }) this.messages = [ ...this.messages, { role: 'tool', - content, + content: fullContent, toolCallId: result.toolCallId, }, ] diff --git a/packages/typescript/ai/src/activities/chat/tools/tool-calls.ts b/packages/typescript/ai/src/activities/chat/tools/tool-calls.ts index f6fe2060e..60d1acfca 100644 --- a/packages/typescript/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/typescript/ai/src/activities/chat/tools/tool-calls.ts @@ -162,6 +162,7 @@ export class ToolCallManager { const tool = this.tools.find((t) => t.name === toolCall.function.name) let toolResultContent: string + let clientResultContent: string | undefined if (tool?.execute) { try { // Parse arguments (normalize "null" to "{}" for empty tool_use blocks) @@ -215,6 +216,13 @@ export class ToolCallManager { toolResultContent = typeof result === 'string' ? result : JSON.stringify(result) + + // Apply clientOutput filter if the tool defines one + if (tool.clientOutput) { + const parsed = + typeof result === 'string' ? JSON.parse(result) : result + clientResultContent = JSON.stringify(tool.clientOutput(parsed)) + } } catch (error: unknown) { // If tool execution fails, add error message const message = @@ -233,7 +241,7 @@ export class ToolCallManager { toolName: toolCall.function.name, model: finishEvent.model, timestamp: Date.now(), - result: toolResultContent, + result: clientResultContent ?? toolResultContent, } // Add tool result message diff --git a/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts b/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts index c12c61898..63a3a9128 100644 --- a/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts +++ b/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts @@ -34,6 +34,7 @@ export interface ClientTool< needsApproval?: boolean lazy?: boolean metadata?: Record + clientOutput?: (result: InferSchemaType) => unknown execute?: ( args: InferSchemaType, ) => Promise> | InferSchemaType @@ -99,6 +100,7 @@ export interface ToolDefinitionConfig< needsApproval?: boolean lazy?: boolean metadata?: Record + clientOutput?: (result: InferSchemaType) => unknown } /** diff --git a/packages/typescript/ai/src/types.ts b/packages/typescript/ai/src/types.ts index f5ddaee59..3bde887ee 100644 --- a/packages/typescript/ai/src/types.ts +++ b/packages/typescript/ai/src/types.ts @@ -496,6 +496,9 @@ export interface Tool< /** If true, this tool is lazy and will only be sent to the LLM after being discovered via the lazy tool discovery mechanism. Only meaningful when used with chat(). */ lazy?: boolean + /** Optional transform to filter tool output before streaming to the client. The full result is always sent to the LLM. */ + clientOutput?: (result: any) => any + /** Additional metadata for adapters or custom extensions */ metadata?: Record } diff --git a/packages/typescript/ai/tests/client-output-types.test.ts b/packages/typescript/ai/tests/client-output-types.test.ts new file mode 100644 index 000000000..959b8c743 --- /dev/null +++ b/packages/typescript/ai/tests/client-output-types.test.ts @@ -0,0 +1,130 @@ +/** + * Type-level tests for clientOutput on tool definitions. + * + * These tests verify that TypeScript correctly infers the clientOutput + * parameter type from the outputSchema, and that the property propagates + * through .server() and .client() builders. + */ + +import { describe, expectTypeOf, it } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '../src/activities/chat/tools/tool-definition' + +describe('clientOutput type inference', () => { + const outputSchema = z.object({ + id: z.string(), + name: z.string(), + ssn: z.string(), + internalScore: z.number(), + }) + + type ExpectedOutput = z.infer + + it('should type the clientOutput parameter from outputSchema', () => { + toolDefinition({ + name: 'lookup_user', + description: 'Look up a user', + outputSchema, + clientOutput: (result) => { + // result should be typed as the inferred output schema type + expectTypeOf(result).toEqualTypeOf() + return { id: result.id } + }, + }) + }) + + it('should allow clientOutput to return any shape', () => { + // clientOutput return type is unknown — no constraints on what shape the + // filtered result has. This is intentional: the client result doesn't need + // to conform to the outputSchema. + toolDefinition({ + name: 'flexible_return', + description: 'Return anything', + outputSchema, + clientOutput: (result) => { + expectTypeOf(result).toEqualTypeOf() + return { justId: result.id, extra: 42 } + }, + }) + }) + + it('should reject accessing nonexistent properties in clientOutput', () => { + toolDefinition({ + name: 'bad_access', + description: 'Bad property access', + outputSchema, + clientOutput: (result) => ({ + id: result.id, + // @ts-expect-error - nonExistent does not exist on output schema + bad: result.nonExistent, + }), + }) + }) + + it('should propagate clientOutput through .server()', () => { + const tool = toolDefinition({ + name: 'server_propagation', + description: 'Check server propagation', + outputSchema, + clientOutput: (result) => ({ id: result.id }), + }) + + const serverTool = tool.server(async () => ({ + id: '1', + name: 'Alice', + ssn: '000', + internalScore: 99, + })) + + // ServerTool extends Tool which uses (result: any) => any for clientOutput. + // Verify it exists and is callable. + expectTypeOf(serverTool).toHaveProperty('clientOutput') + if (serverTool.clientOutput) { + expectTypeOf(serverTool.clientOutput).toBeFunction() + } + }) + + it('should propagate clientOutput through .client()', () => { + const tool = toolDefinition({ + name: 'client_propagation', + description: 'Check client propagation', + outputSchema, + clientOutput: (result) => ({ id: result.id }), + }) + + const clientTool = tool.client() + + // ClientTool preserves the strongly-typed clientOutput + expectTypeOf(clientTool).toHaveProperty('clientOutput') + if (clientTool.clientOutput) { + // The parameter is typed from outputSchema + type Param = Parameters[0] + expectTypeOf().toEqualTypeOf() + } + }) + + it('should type clientOutput parameter as any when no outputSchema', () => { + // Without outputSchema, SchemaInput defaults to StandardJSONSchemaV1 | JSONSchema, + // so InferSchemaType resolves to `any` — matching the convention of execute's args type. + toolDefinition({ + name: 'no_schema', + description: 'No output schema', + clientOutput: (result) => { + expectTypeOf(result).toBeAny() + return result + }, + }) + }) + + it('should allow omitting clientOutput entirely', () => { + const tool = toolDefinition({ + name: 'no_filter', + description: 'No client filtering', + outputSchema, + }) + + // clientOutput should be optional (undefined at runtime) + expectTypeOf(tool).toHaveProperty('clientOutput') + expectTypeOf(tool.clientOutput).toBeNullable() + }) +}) diff --git a/packages/typescript/ai/tests/tool-call-manager.test.ts b/packages/typescript/ai/tests/tool-call-manager.test.ts index 546fc7a95..bc706223f 100644 --- a/packages/typescript/ai/tests/tool-call-manager.test.ts +++ b/packages/typescript/ai/tests/tool-call-manager.test.ts @@ -375,6 +375,58 @@ describe('ToolCallManager', () => { expect(toolCalls[0]?.function.arguments).toBe('{"location":"New York"}') }) }) + + it('should apply clientOutput filter to TOOL_CALL_END but keep full result in tool message', async () => { + const sensitiveWeatherTool: Tool = { + name: 'get_weather_sensitive', + description: 'Get weather with internal data', + execute: vi.fn(() => ({ + temp: 72, + location: 'Paris', + internalStationId: 'SENSOR-42', + })), + clientOutput: (result: any) => ({ + temp: result.temp, + location: result.location, + }), + } + + const manager = new ToolCallManager([sensitiveWeatherTool]) + + manager.addToolCallStartEvent({ + type: 'TOOL_CALL_START', + toolCallId: 'call_filtered', + toolName: 'get_weather_sensitive', + timestamp: Date.now(), + index: 0, + }) + + manager.addToolCallArgsEvent({ + type: 'TOOL_CALL_ARGS', + toolCallId: 'call_filtered', + timestamp: Date.now(), + delta: '{}', + }) + + const { chunks, result: toolMessages } = await collectGeneratorOutput( + manager.executeTools(mockFinishedEvent), + ) + + // TOOL_CALL_END chunk should contain filtered result (no internalStationId) + expect(chunks).toHaveLength(1) + const clientResult = JSON.parse(chunks[0]!.result!) + expect(clientResult).toEqual({ temp: 72, location: 'Paris' }) + expect(clientResult).not.toHaveProperty('internalStationId') + + // Tool message for LLM should contain full result + expect(toolMessages).toHaveLength(1) + const llmResult = JSON.parse(toolMessages[0]!.content as string) + expect(llmResult).toEqual({ + temp: 72, + location: 'Paris', + internalStationId: 'SENSOR-42', + }) + }) }) describe('executeToolCalls', () => { @@ -720,4 +772,103 @@ describe('executeToolCalls', () => { ) }) }) + + describe('clientOutput filtering', () => { + it('should filter tool result in TOOL_CALL_END but keep full result for LLM', async () => { + const tool: Tool = { + name: 'lookup_user', + description: 'Look up a user', + execute: vi.fn(() => ({ + id: 'u1', + name: 'Alice', + ssn: '123-45-6789', + internalScore: 99, + })), + clientOutput: (result: any) => ({ + id: result.id, + name: result.name, + }), + } + + const toolCalls = [makeToolCall('call_1', 'lookup_user', '{}')] + + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) + + const chunks: Array = [] + let next = await gen.next() + while (!next.done) { + chunks.push(next.value) + next = await gen.next() + } + const result = next.value + + // Full result goes to LLM + expect(result.results).toHaveLength(1) + expect(result.results[0]?.result).toEqual({ + id: 'u1', + name: 'Alice', + ssn: '123-45-6789', + internalScore: 99, + }) + }) + + it('should not apply clientOutput when tool execution errors', async () => { + const tool: Tool = { + name: 'failing_tool', + description: 'A tool that fails', + execute: vi.fn(() => { + throw new Error('DB connection failed') + }), + clientOutput: (result: any) => ({ id: result.id }), + } + + const toolCalls = [makeToolCall('call_1', 'failing_tool', '{}')] + + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) + + const chunks: Array = [] + let next = await gen.next() + while (!next.done) { + chunks.push(next.value) + next = await gen.next() + } + const result = next.value + + // Error result should not be filtered + expect(result.results[0]?.result).toEqual({ + error: 'DB connection failed', + }) + expect(result.results[0]?.state).toBe('output-error') + }) + + it('should send full result to client when no clientOutput is defined', async () => { + const tool: Tool = { + name: 'open_tool', + description: 'No filtering', + execute: vi.fn(() => ({ + id: 'u1', + secret: 'visible', + })), + // No clientOutput + } + + const toolCalls = [makeToolCall('call_1', 'open_tool', '{}')] + + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) + + const chunks: Array = [] + let next = await gen.next() + while (!next.done) { + chunks.push(next.value) + next = await gen.next() + } + const result = next.value + + // Full result in both places + expect(result.results[0]?.result).toEqual({ + id: 'u1', + secret: 'visible', + }) + }) + }) }) diff --git a/packages/typescript/ai/tests/tool-definition.test.ts b/packages/typescript/ai/tests/tool-definition.test.ts index 642992c68..e2a27feab 100644 --- a/packages/typescript/ai/tests/tool-definition.test.ts +++ b/packages/typescript/ai/tests/tool-definition.test.ts @@ -286,4 +286,40 @@ describe('toolDefinition', () => { expect(tool.__toolSide).toBe('definition') expect(tool.inputSchema).toBeDefined() }) + + it('should preserve clientOutput on definition, server, and client tools', () => { + const filterFn = (result: { id: string; secret: string }) => ({ + id: result.id, + }) + + const tool = toolDefinition({ + name: 'lookupUser', + description: 'Look up a user', + outputSchema: z.object({ + id: z.string(), + secret: z.string(), + }), + clientOutput: filterFn, + }) + + expect(tool.clientOutput).toBe(filterFn) + + const serverTool = tool.server(async () => ({ id: '1', secret: 'shhh' })) + expect(serverTool.clientOutput).toBe(filterFn) + + const clientTool = tool.client() + expect(clientTool.clientOutput).toBe(filterFn) + }) + + it('should default clientOutput to undefined when not specified', () => { + const tool = toolDefinition({ + name: 'simpleTool', + description: 'No filtering', + }) + + expect(tool.clientOutput).toBeUndefined() + + const serverTool = tool.server(async () => ({})) + expect(serverTool.clientOutput).toBeUndefined() + }) })