From 0f3503e4952d52c2b361ab6646ea1061140ad564 Mon Sep 17 00:00:00 2001
From: Dev Randalpura <devrandalpura@google.com>
Date: Wed, 4 Mar 2026 17:40:03 -0500
Subject: [PATCH 1/2] Added partial output to cancelled shell UI

---
 .../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 233f9057603..76b8f95ce76 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 b51d7c435b1..3fa8a62bf86 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 005f3004d67..68b60fbbb27 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 && existingOutput) {
+        finalResponse.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,

From 2f1db4eb05db605fe1aec5b67bc80b2bb4e0784b Mon Sep 17 00:00:00 2001
From: Dev Randalpura <devrandalpura@google.com>
Date: Wed, 4 Mar 2026 17:45:23 -0500
Subject: [PATCH 2/2] Addressed comment related to resultDisplay not being used

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 packages/core/src/scheduler/state-manager.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts
index 68b60fbbb27..428b7f87a8f 100644
--- a/packages/core/src/scheduler/state-manager.ts
+++ b/packages/core/src/scheduler/state-manager.ts
@@ -485,8 +485,8 @@ export class SchedulerStateManager {
 
     if (isToolCallResponseInfo(reason)) {
       const finalResponse = { ...reason };
-      if (!finalResponse.resultDisplay && existingOutput) {
-        finalResponse.resultDisplay = existingOutput;
+      if (!finalResponse.resultDisplay) {
+        finalResponse.resultDisplay = resultDisplay ?? existingOutput;
       }
 
       return {