From 2a8918c72fb155f07e474ebbb7497066ce1fe604 Mon Sep 17 00:00:00 2001
From: Michael Bleigh <bleigh@google.com>
Date: Tue, 24 Mar 2026 12:00:18 -0700
Subject: [PATCH] feat: achieve rich tool call display parity for agent stream

This change ensures that tool execution in the Agent Protocol experiment
reaches visual parity with the legacy implementation by passing rich
metadata and objects (like FileDiff) through the AgentEvent stream.

Key changes:
- Core 'LegacyAgentSession' now builds temporary invocations for tool requests
  to capture rich, argument-aware descriptions (e.g. 'Writing to poem.md').
- Core 'LegacyAgentSession' now attaches the raw 'resultDisplay' object and
  'outputFile' path to the 'tool_response' event metadata.
- Core 'LegacyAgentSession' now passes the rich description through 'tool_update'
  events to ensure dynamic descriptions are updated during execution.
- UI 'useAgentStream' hook now extracts these rich values from event metadata
  to populate the local 'trackedTools' state, allowing 'mapToDisplay' to
  correctly trigger bespoke rendering components (like the diff viewer).
---
 packages/cli/src/ui/hooks/useAgentStream.ts     |  7 ++++++-
 packages/core/src/agent/legacy-agent-session.ts | 11 ++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/packages/cli/src/ui/hooks/useAgentStream.ts b/packages/cli/src/ui/hooks/useAgentStream.ts
index 30fdac3c67..14e816021a 100644
--- a/packages/cli/src/ui/hooks/useAgentStream.ts
+++ b/packages/cli/src/ui/hooks/useAgentStream.ts
@@ -216,6 +216,7 @@ export const useAgentStream = (
               tool: {
                 displayName: (event._meta?.['displayName'] as string) ?? event.name,
                 isOutputMarkdown: (event._meta?.['isOutputMarkdown'] as boolean) ?? false,
+                kind: event._meta?.['kind'] as any,
               },
               invocation: {
                 getDescription: () => (event._meta?.['description'] as string) ?? '',
@@ -234,6 +235,9 @@ export const useAgentStream = (
                     progress: event.data?.['progress'] as number | undefined,
                     progressTotal: event.data?.['progressTotal'] as number | undefined,
                     pid: event.data?.['pid'] as number | undefined,
+                    invocation: {
+                      getDescription: () => (event._meta?.['description'] as string) ?? (tc as any).invocation?.getDescription(),
+                    },
                   } as unknown as TrackedToolCall)
                 : tc,
             ),
@@ -247,7 +251,8 @@ export const useAgentStream = (
                     ...tc,
                     status: event.isError ? 'error' : 'success',
                     response: {
-                      resultDisplay: event.displayContent?.[0]?.type === 'text' ? event.displayContent[0].text : undefined,
+                      resultDisplay: event._meta?.['resultDisplay'] ?? (event.displayContent?.[0]?.type === 'text' ? event.displayContent[0].text : undefined),
+                      outputFile: event._meta?.['outputFile'] as string | undefined,
                     },
                     responseSubmittedToGemini: true,
                   } as unknown as TrackedToolCall)
diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts
index 36760b2629..975134969d 100644
--- a/packages/core/src/agent/legacy-agent-session.ts
+++ b/packages/core/src/agent/legacy-agent-session.ts
@@ -179,6 +179,9 @@ class LegacyAgentProtocol implements AgentProtocol {
                 progressTotal: tc.progressTotal,
                 pid: tc.pid,
               },
+              _meta: {
+                description: tc.invocation.getDescription(),
+              },
             }),
           );
         }
@@ -226,10 +229,12 @@ class LegacyAgentProtocol implements AgentProtocol {
           for (const ev of translatedEvents) {
             if (ev.type === 'tool_request') {
               const tool = this._config.getToolRegistry().getTool(ev.name);
+              const invocation = tool?.build(ev.args);
               ev._meta = {
                 displayName: tool?.displayName ?? ev.name,
-                description: tool?.description ?? '',
+                description: invocation?.getDescription() ?? tool?.description ?? '',
                 isOutputMarkdown: tool?.isOutputMarkdown ?? false,
+                kind: tool?.kind,
               };
             }
           }
@@ -298,6 +303,10 @@ class LegacyAgentProtocol implements AgentProtocol {
               isError: response.error !== undefined,
               ...(displayContent ? { displayContent } : {}),
               ...(data ? { data } : {}),
+              _meta: {
+                resultDisplay: response.resultDisplay,
+                outputFile: response.outputFile,
+              },
             }),
           ]);