From f20e45175e8b16357266758416606a30a146a507 Mon Sep 17 00:00:00 2001
From: Dev Randalpura <devrandalpura@google.com>
Date: Thu, 5 Mar 2026 09:30:20 -0800
Subject: [PATCH] fix(ui): add partial output to cancelled shell UI (#21178)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../components/messages/ShellToolMessage.test.tsx |  8 ++++++++
 .../__snapshots__/ShellToolMessage.test.tsx.snap  |  8 ++++++++
 packages/core/src/scheduler/state-manager.ts      | 15 +++++++++++++--
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/packages/cli/src/ui/components/messages/ShellToolMessage.test.tsx b/packages/cli/src/ui/components/messages/ShellToolMessage.test.tsx
index 233f905760..76b8f95ce7 100644
--- a/packages/cli/src/ui/components/messages/ShellToolMessage.test.tsx
+++ b/packages/cli/src/ui/components/messages/ShellToolMessage.test.tsx
@@ -137,6 +137,14 @@ describe('<ShellToolMessage />', () => {
         { status: CoreToolCallStatus.Error, resultDisplay: 'Error output' },
         undefined,
       ],
+      [
+        'renders in Cancelled state with partial output',
+        {
+          status: CoreToolCallStatus.Cancelled,
+          resultDisplay: 'Partial output before cancellation',
+        },
+        undefined,
+      ],
       [
         'renders in Alternate Buffer mode while focused',
         {
diff --git a/packages/cli/src/ui/components/messages/__snapshots__/ShellToolMessage.test.tsx.snap b/packages/cli/src/ui/components/messages/__snapshots__/ShellToolMessage.test.tsx.snap
index b51d7c435b..3fa8a62bf8 100644
--- a/packages/cli/src/ui/components/messages/__snapshots__/ShellToolMessage.test.tsx.snap
+++ b/packages/cli/src/ui/components/messages/__snapshots__/ShellToolMessage.test.tsx.snap
@@ -309,6 +309,14 @@ exports[`<ShellToolMessage /> > Snapshots > renders in Alternate Buffer mode whi
 "
 `;
 
+exports[`<ShellToolMessage /> > Snapshots > renders in Cancelled state with partial output 1`] = `
+"╭──────────────────────────────────────────────────────────────────────────────╮
+│ -  Shell Command A shell command                                             │
+│                                                                              │
+│ Partial output before cancellation                                           │
+"
+`;
+
 exports[`<ShellToolMessage /> > Snapshots > renders in Error state 1`] = `
 "╭──────────────────────────────────────────────────────────────────────────────╮
 │ x  Shell Command A shell command                                             │
diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts
index 005f3004d6..428b7f87a8 100644
--- a/packages/core/src/scheduler/state-manager.ts
+++ b/packages/core/src/scheduler/state-manager.ts
@@ -477,13 +477,24 @@ export class SchedulerStateManager {
       }
     }
 
+    // Capture any existing live output so it isn't lost when forcing cancellation.
+    let existingOutput: ToolResultDisplay | undefined = undefined;
+    if (call.status === CoreToolCallStatus.Executing && call.liveOutput) {
+      existingOutput = call.liveOutput;
+    }
+
     if (isToolCallResponseInfo(reason)) {
+      const finalResponse = { ...reason };
+      if (!finalResponse.resultDisplay) {
+        finalResponse.resultDisplay = resultDisplay ?? existingOutput;
+      }
+
       return {
         request: call.request,
         tool: call.tool,
         invocation: call.invocation,
         status: CoreToolCallStatus.Cancelled,
-        response: reason,
+        response: finalResponse,
         durationMs: startTime ? Date.now() - startTime : undefined,
         outcome: call.outcome,
         schedulerId: call.schedulerId,
@@ -508,7 +519,7 @@ export class SchedulerStateManager {
             },
           },
         ],
-        resultDisplay,
+        resultDisplay: resultDisplay ?? existingOutput,
         error: undefined,
         errorType: undefined,
         contentLength: errorMessage.length,