From 7f7e69e6c6e3a895bdaccfaa7c4e520c9a40d72c Mon Sep 17 00:00:00 2001 From: Michael Bleigh Date: Fri, 10 Apr 2026 10:51:21 -0700 Subject: [PATCH] feat(agent): implement tool-controlled display protocol (Steps 2-3) --- .../cli/src/nonInteractiveCliAgentSession.ts | 10 ++- packages/cli/src/ui/hooks/useAgentStream.ts | 11 +-- packages/cli/src/ui/types.ts | 2 + packages/core/src/agent/content-utils.test.ts | 22 ------ packages/core/src/agent/content-utils.ts | 18 ----- .../core/src/agent/event-translator.test.ts | 23 +++--- packages/core/src/agent/event-translator.ts | 15 ++-- .../src/agent/legacy-agent-session.test.ts | 7 +- .../core/src/agent/legacy-agent-session.ts | 13 ++-- .../core/src/agent/tool-display-utils.test.ts | 71 ++++++++++++++++++ packages/core/src/agent/tool-display-utils.ts | 72 +++++++++++++++++++ packages/core/src/agent/types.ts | 29 ++++++-- packages/core/src/index.ts | 9 +++ 13 files changed, 229 insertions(+), 73 deletions(-) create mode 100644 packages/core/src/agent/tool-display-utils.test.ts create mode 100644 packages/core/src/agent/tool-display-utils.ts diff --git a/packages/cli/src/nonInteractiveCliAgentSession.ts b/packages/cli/src/nonInteractiveCliAgentSession.ts index 7f36ce6cf5..29830b8e96 100644 --- a/packages/cli/src/nonInteractiveCliAgentSession.ts +++ b/packages/cli/src/nonInteractiveCliAgentSession.ts @@ -470,7 +470,10 @@ export async function runNonInteractive({ case 'tool_response': { textOutput.ensureTrailingNewline(); if (streamFormatter) { - const displayText = getTextContent(event.displayContent); + const displayText = + event.display?.result?.type === 'text' + ? event.display.result.text + : undefined; const errorMsg = getTextContent(event.content) ?? 'Tool error'; streamFormatter.emitEvent({ type: JsonStreamEventType.TOOL_RESULT, @@ -490,7 +493,10 @@ export async function runNonInteractive({ }); } if (event.isError) { - const displayText = getTextContent(event.displayContent); + const displayText = + event.display?.result?.type === 'text' + ? event.display.result.text + : undefined; const errorMsg = getTextContent(event.content) ?? 'Tool error'; if (event.data?.['errorType'] === ToolErrorType.STOP_EXECUTION) { diff --git a/packages/cli/src/ui/hooks/useAgentStream.ts b/packages/cli/src/ui/hooks/useAgentStream.ts index 81dbb1e9e9..e978cead6e 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.ts +++ b/packages/cli/src/ui/hooks/useAgentStream.ts @@ -197,6 +197,7 @@ export const useAgentStream = ({ name: displayName, originalRequestName: event.name, description: desc, + display: event.display, status: CoreToolCallStatus.Scheduled, isClientInitiated: false, renderOutputAsMarkdown: isOutputMarkdown, @@ -223,8 +224,8 @@ export const useAgentStream = ({ status = CoreToolCallStatus.Success; const liveOutput = - event.displayContent?.[0]?.type === 'text' - ? event.displayContent[0].text + event.display?.result?.type === 'text' + ? event.display.result.text : tc.resultDisplay; const progressMessage = legacyState?.progressMessage ?? tc.progressMessage; @@ -237,6 +238,7 @@ export const useAgentStream = ({ return { ...tc, status, + display: event.display ?? tc.display, resultDisplay: liveOutput, progressMessage, progress, @@ -256,8 +258,8 @@ export const useAgentStream = ({ const legacyState = event._meta?.legacyState; const outputFile = legacyState?.outputFile; const resultDisplay = - event.displayContent?.[0]?.type === 'text' - ? event.displayContent[0].text + event.display?.result?.type === 'text' + ? event.display.result.text : tc.resultDisplay; return { @@ -265,6 +267,7 @@ export const useAgentStream = ({ status: event.isError ? CoreToolCallStatus.Error : CoreToolCallStatus.Success, + display: event.display ?? tc.display, resultDisplay, outputFile, }; diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 6fbc3151d8..1ded2ae643 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -11,6 +11,7 @@ import { type ThoughtSummary, type SerializableConfirmationDetails, type ToolResultDisplay, + type ToolDisplay, type RetrieveUserQuotaResponse, type SkillDefinition, type AgentDefinition, @@ -121,6 +122,7 @@ export interface IndividualToolCallDisplay { name: string; args?: Record; description: string; + display?: ToolDisplay; resultDisplay: ToolResultDisplay | undefined; status: CoreToolCallStatus; // True when the tool was initiated directly by the user (slash/@/shell flows). diff --git a/packages/core/src/agent/content-utils.test.ts b/packages/core/src/agent/content-utils.test.ts index 96608c8227..7de54c56fa 100644 --- a/packages/core/src/agent/content-utils.test.ts +++ b/packages/core/src/agent/content-utils.test.ts @@ -8,7 +8,6 @@ import { describe, expect, it } from 'vitest'; import { geminiPartsToContentParts, contentPartsToGeminiParts, - toolResultDisplayToContentParts, buildToolResponseData, } from './content-utils.js'; import type { Part } from '@google/genai'; @@ -200,27 +199,6 @@ describe('contentPartsToGeminiParts', () => { }); }); -describe('toolResultDisplayToContentParts', () => { - it('returns undefined for undefined', () => { - expect(toolResultDisplayToContentParts(undefined)).toBeUndefined(); - }); - - it('returns undefined for null', () => { - expect(toolResultDisplayToContentParts(null)).toBeUndefined(); - }); - - it('handles string resultDisplay as-is', () => { - const result = toolResultDisplayToContentParts('File written'); - expect(result).toEqual([{ type: 'text', text: 'File written' }]); - }); - - it('stringifies object resultDisplay', () => { - const display = { type: 'FileDiff', oldPath: 'a.ts', newPath: 'b.ts' }; - const result = toolResultDisplayToContentParts(display); - expect(result).toEqual([{ type: 'text', text: JSON.stringify(display) }]); - }); -}); - describe('buildToolResponseData', () => { it('preserves outputFile and contentLength', () => { const result = buildToolResponseData({ diff --git a/packages/core/src/agent/content-utils.ts b/packages/core/src/agent/content-utils.ts index b117ab69fc..aaf191fe8e 100644 --- a/packages/core/src/agent/content-utils.ts +++ b/packages/core/src/agent/content-utils.ts @@ -101,24 +101,6 @@ export function contentPartsToGeminiParts(content: ContentPart[]): Part[] { return result; } -/** - * Converts a ToolCallResponseInfo.resultDisplay value into ContentPart[]. - * Handles string, object-valued (FileDiff, SubagentProgress, etc.), - * and undefined resultDisplay consistently. - */ -export function toolResultDisplayToContentParts( - resultDisplay: unknown, -): ContentPart[] | undefined { - if (resultDisplay === undefined || resultDisplay === null) { - return undefined; - } - const text = - typeof resultDisplay === 'string' - ? resultDisplay - : JSON.stringify(resultDisplay); - return [{ type: 'text', text }]; -} - /** * Builds the data record for a tool_response AgentEvent, preserving * all available metadata from the ToolCallResponseInfo. diff --git a/packages/core/src/agent/event-translator.test.ts b/packages/core/src/agent/event-translator.test.ts index be9d8ea40e..cfb5cfe300 100644 --- a/packages/core/src/agent/event-translator.test.ts +++ b/packages/core/src/agent/event-translator.test.ts @@ -155,9 +155,10 @@ describe('translateEvent', () => { expect(resp.content).toEqual([ { type: 'text', text: 'Permission denied to write' }, ]); - expect(resp.displayContent).toEqual([ - { type: 'text', text: 'Permission denied' }, - ]); + expect(resp.display?.result).toEqual({ + type: 'text', + text: 'Permission denied', + }); expect(resp.data).toEqual({ errorType: 'permission_denied' }); }); @@ -200,9 +201,12 @@ describe('translateEvent', () => { }; const result = translateEvent(event, state); const resp = result[0] as AgentEvent<'tool_response'>; - expect(resp.displayContent).toEqual([ - { type: 'text', text: JSON.stringify(objectDisplay) }, - ]); + expect(resp.display?.result).toEqual({ + type: 'diff', + path: '/tmp/test.txt', + beforeText: 'a', + afterText: 'b', + }); }); it('passes through string resultDisplay as-is', () => { @@ -220,9 +224,10 @@ describe('translateEvent', () => { }; const result = translateEvent(event, state); const resp = result[0] as AgentEvent<'tool_response'>; - expect(resp.displayContent).toEqual([ - { type: 'text', text: 'Command output text' }, - ]); + expect(resp.display?.result).toEqual({ + type: 'text', + text: 'Command output text', + }); }); it('preserves outputFile and contentLength in data', () => { diff --git a/packages/core/src/agent/event-translator.ts b/packages/core/src/agent/event-translator.ts index cb299b494c..dee56adbd0 100644 --- a/packages/core/src/agent/event-translator.ts +++ b/packages/core/src/agent/event-translator.ts @@ -25,12 +25,13 @@ import type { ErrorData, Usage, AgentEventType, + ToolDisplay, } from './types.js'; import { geminiPartsToContentParts, - toolResultDisplayToContentParts, buildToolResponseData, } from './content-utils.js'; +import { toolResultDisplayToDisplayContent } from './tool-display-utils.js'; // --------------------------------------------------------------------------- // Translation State @@ -241,10 +242,14 @@ export function translateEvent( case GeminiEventType.ToolCallResponse: { ensureStreamStart(state, out); - const displayContent = toolResultDisplayToContentParts( - event.value.resultDisplay, - ); const data = buildToolResponseData(event.value); + const display: ToolDisplay | undefined = event.value.resultDisplay + ? { + result: toolResultDisplayToDisplayContent( + event.value.resultDisplay, + ), + } + : undefined; out.push( makeEvent('tool_response', state, { requestId: event.value.callId, @@ -253,7 +258,7 @@ export function translateEvent( ? [{ type: 'text', text: event.value.error.message }] : geminiPartsToContentParts(event.value.responseParts), isError: event.value.error !== undefined, - ...(displayContent ? { displayContent } : {}), + ...(display ? { display } : {}), ...(data ? { data } : {}), }), ); diff --git a/packages/core/src/agent/legacy-agent-session.test.ts b/packages/core/src/agent/legacy-agent-session.test.ts index 1de5d90e20..9ee8b032ad 100644 --- a/packages/core/src/agent/legacy-agent-session.test.ts +++ b/packages/core/src/agent/legacy-agent-session.test.ts @@ -489,9 +489,10 @@ describe('LegacyAgentSession', () => { expect(toolResp?.content).toEqual([ { type: 'text', text: 'Permission denied' }, ]); - expect(toolResp?.displayContent).toEqual([ - { type: 'text', text: 'Error display' }, - ]); + expect(toolResp?.display?.result).toEqual({ + type: 'text', + text: 'Error display', + }); }); it('stops on STOP_EXECUTION tool error', async () => { diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts index 94763c7d40..5fb024378e 100644 --- a/packages/core/src/agent/legacy-agent-session.ts +++ b/packages/core/src/agent/legacy-agent-session.ts @@ -23,8 +23,8 @@ import { buildToolResponseData, contentPartsToGeminiParts, geminiPartsToContentParts, - toolResultDisplayToContentParts, } from './content-utils.js'; +import { populateToolDisplay } from './tool-display-utils.js'; import { AgentSession } from './agent-session.js'; import { createTranslationState, @@ -262,9 +262,12 @@ export class LegacyAgentProtocol implements AgentProtocol { const content: ContentPart[] = response.error ? [{ type: 'text', text: response.error.message }] : geminiPartsToContentParts(response.responseParts); - const displayContent = toolResultDisplayToContentParts( - response.resultDisplay, - ); + const display = populateToolDisplay({ + name: request.name, + invocation: 'invocation' in tc ? tc.invocation : undefined, + resultDisplay: response.resultDisplay, + displayName: 'tool' in tc ? tc.tool?.displayName : undefined, + }); const data = buildToolResponseData(response); this._emit([ @@ -273,7 +276,7 @@ export class LegacyAgentProtocol implements AgentProtocol { name: request.name, content, isError: response.error !== undefined, - ...(displayContent ? { displayContent } : {}), + ...(display ? { display } : {}), ...(data ? { data } : {}), }), ]); diff --git a/packages/core/src/agent/tool-display-utils.test.ts b/packages/core/src/agent/tool-display-utils.test.ts new file mode 100644 index 0000000000..de88870a7c --- /dev/null +++ b/packages/core/src/agent/tool-display-utils.test.ts @@ -0,0 +1,71 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import type { + ToolInvocation, + ToolResult, + ToolResultDisplay, +} from '../tools/tools.js'; +import { populateToolDisplay } from './tool-display-utils.js'; + +describe('tool-display-utils', () => { + describe('populateToolDisplay', () => { + it('uses displayName if provided', () => { + const mockInvocation = { + getDescription: () => 'Doing something...', + } as unknown as ToolInvocation; + + const display = populateToolDisplay({ + name: 'raw-name', + invocation: mockInvocation, + displayName: 'Custom Display Name', + }); + expect(display.name).toBe('Custom Display Name'); + expect(display.description).toBe('Doing something...'); + }); + + it('falls back to raw name if no displayName provided', () => { + const mockInvocation = { + getDescription: () => 'Doing something...', + } as unknown as ToolInvocation; + + const display = populateToolDisplay({ + name: 'raw-name', + invocation: mockInvocation, + }); + expect(display.name).toBe('raw-name'); + }); + + it('populates result from resultDisplay', () => { + const display = populateToolDisplay({ + name: 'test', + resultDisplay: 'hello world', + }); + expect(display.result).toEqual({ type: 'text', text: 'hello world' }); + }); + + it('translates FileDiff to DisplayDiff', () => { + const fileDiff = { + fileDiff: '@@ ...', + fileName: 'test.ts', + filePath: 'src/test.ts', + originalContent: 'old', + newContent: 'new', + } as unknown as ToolResultDisplay; + const display = populateToolDisplay({ + name: 'test', + resultDisplay: fileDiff, + }); + expect(display.result).toEqual({ + type: 'diff', + path: 'src/test.ts', + beforeText: 'old', + afterText: 'new', + }); + }); + }); +}); diff --git a/packages/core/src/agent/tool-display-utils.ts b/packages/core/src/agent/tool-display-utils.ts new file mode 100644 index 0000000000..f5327d6a1b --- /dev/null +++ b/packages/core/src/agent/tool-display-utils.ts @@ -0,0 +1,72 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ToolInvocation, + ToolResult, + ToolResultDisplay, +} from '../tools/tools.js'; +import type { ToolDisplay, DisplayContent } from './types.js'; + +/** + * Populates a ToolDisplay object from a tool invocation and its result. + * This serves as a centralized bridge during the migration to tool-controlled display. + */ +export function populateToolDisplay({ + name, + invocation, + resultDisplay, + displayName, +}: { + name: string; + invocation?: ToolInvocation; + resultDisplay?: ToolResultDisplay; + displayName?: string; +}): ToolDisplay { + const display: ToolDisplay = { + name: displayName || name, + description: invocation?.getDescription?.(), + }; + + if (resultDisplay) { + display.result = toolResultDisplayToDisplayContent(resultDisplay); + } + + return display; +} + +/** + * Converts a legacy ToolResultDisplay into the new DisplayContent format. + */ +export function toolResultDisplayToDisplayContent( + resultDisplay: ToolResultDisplay, +): DisplayContent { + if (typeof resultDisplay === 'string') { + return { type: 'text', text: resultDisplay }; + } + + // Handle FileDiff -> DisplayDiff + if ( + typeof resultDisplay === 'object' && + resultDisplay !== null && + 'fileDiff' in resultDisplay && + 'newContent' in resultDisplay + ) { + return { + type: 'diff', + path: resultDisplay.filePath || resultDisplay.fileName, + beforeText: resultDisplay.originalContent ?? '', + afterText: resultDisplay.newContent, + }; + } + + // Fallback for other structured types (LsTool, GrepTool, etc.) + // These will be fully migrated in Step 5. + return { + type: 'text', + text: JSON.stringify(resultDisplay), + }; +} diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index 19837c138e..af48973f8f 100644 --- a/packages/core/src/agent/types.ts +++ b/packages/core/src/agent/types.ts @@ -106,7 +106,7 @@ export interface AgentEvents { /** Updates configuration about the current session/agent. */ session_update: SessionUpdate; /** Message content provided by user, agent, or developer. */ - message: Message; + message: AgentMessage; /** Event indicating the start of agent activity on a stream. */ agent_start: AgentStart; /** Event indicating the end of agent activity on a stream. */ @@ -170,17 +170,35 @@ export type ContentPart = ) & WithMeta; -export interface Message { +export interface AgentMessage { role: 'user' | 'agent' | 'developer'; content: ContentPart[]; } +export type DisplayText = { type: 'text'; text: string }; +export type DisplayDiff = { + type: 'diff'; + path?: string; + beforeText: string; + afterText: string; +}; +export type DisplayContent = DisplayText | DisplayDiff; + +export interface ToolDisplay { + name?: string; + description?: string; + resultSummary?: string; + result?: DisplayContent; +} + export interface ToolRequest { /** A unique identifier for this tool request to be correlated by the response. */ requestId: string; /** The name of the tool being requested. */ name: string; /** The arguments for the tool. */ + /** Tool-controlled display information. */ + display?: ToolDisplay; args: Record; /** UI specific metadata */ _meta?: { @@ -201,7 +219,8 @@ export interface ToolRequest { */ export interface ToolUpdate { requestId: string; - displayContent?: ContentPart[]; + /** Tool-controlled display information. */ + display?: ToolDisplay; content?: ContentPart[]; data?: Record; /** UI specific metadata */ @@ -221,8 +240,8 @@ export interface ToolUpdate { export interface ToolResponse { requestId: string; name: string; - /** Content representing the tool call's outcome to be presented to the user. */ - displayContent?: ContentPart[]; + /** Tool-controlled display information. */ + display?: ToolDisplay; /** Multi-part content to be sent to the model. */ content?: ContentPart[]; /** Structured data to be sent to the model. */ diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 04456a2964..6faf3f0a60 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -193,6 +193,7 @@ export * from './agent/agent-session.js'; export * from './agent/legacy-agent-session.js'; export * from './agent/event-translator.js'; export * from './agent/content-utils.js'; +export * from './agent/tool-display-utils.js'; // Agent event types — namespaced to avoid collisions with existing exports export type { AgentEvent, @@ -204,6 +205,7 @@ export type { AgentProtocol, AgentSend, AgentStart, + AgentMessage, ContentPart, ErrorData, StreamEndReason, @@ -211,6 +213,13 @@ export type { Unsubscribe, Usage as AgentUsage, WithMeta, + ToolRequest, + ToolResponse, + ToolUpdate, + ToolDisplay, + DisplayText, + DisplayDiff, + DisplayContent, } from './agent/types.js'; // Export specific tool logic