From d97fc715998bb32b8570f1892e97091608ec3783 Mon Sep 17 00:00:00 2001 From: imsherr Date: Tue, 24 Mar 2026 17:43:38 -0400 Subject: [PATCH 1/2] feat(ai): add clientOutput to tool definitions for server-side result filtering Tools that return sensitive data (PII, internal scores, credentials) often need the LLM to see the full result for reasoning while keeping that data off the wire to the client. `clientOutput` is an optional transform on the tool definition that splits the result: the full output feeds back to the LLM as a tool-role message, while only the transformed output is streamed to the client via TOOL_CALL_END. --- .../ai/src/activities/chat/index.ts | 13 +- .../src/activities/chat/tools/tool-calls.ts | 10 +- .../activities/chat/tools/tool-definition.ts | 2 + packages/typescript/ai/src/types.ts | 3 + .../ai/tests/client-output-types.test.ts | 130 ++++++++++++++ .../ai/tests/tool-call-manager.test.ts | 166 ++++++++++++++++++ .../ai/tests/tool-definition.test.ts | 36 ++++ 7 files changed, 356 insertions(+), 4 deletions(-) create mode 100644 packages/typescript/ai/tests/client-output-types.test.ts 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..555b9e467 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,118 @@ 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() + }) }) From ef97fcfe8ff34b5fb507a61210a7571c42ecf84f Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 24 Mar 2026 21:57:59 +0000 Subject: [PATCH 2/2] ci: apply automated fixes --- .../ai/tests/tool-call-manager.test.ts | 21 +++---------------- 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/packages/typescript/ai/tests/tool-call-manager.test.ts b/packages/typescript/ai/tests/tool-call-manager.test.ts index 555b9e467..bc706223f 100644 --- a/packages/typescript/ai/tests/tool-call-manager.test.ts +++ b/packages/typescript/ai/tests/tool-call-manager.test.ts @@ -792,12 +792,7 @@ describe('executeToolCalls', () => { const toolCalls = [makeToolCall('call_1', 'lookup_user', '{}')] - const gen = executeToolCalls( - toolCalls, - [tool], - new Map(), - new Map(), - ) + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) const chunks: Array = [] let next = await gen.next() @@ -829,12 +824,7 @@ describe('executeToolCalls', () => { const toolCalls = [makeToolCall('call_1', 'failing_tool', '{}')] - const gen = executeToolCalls( - toolCalls, - [tool], - new Map(), - new Map(), - ) + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) const chunks: Array = [] let next = await gen.next() @@ -864,12 +854,7 @@ describe('executeToolCalls', () => { const toolCalls = [makeToolCall('call_1', 'open_tool', '{}')] - const gen = executeToolCalls( - toolCalls, - [tool], - new Map(), - new Map(), - ) + const gen = executeToolCalls(toolCalls, [tool], new Map(), new Map()) const chunks: Array = [] let next = await gen.next()