Gateway/Control UI: preserve partial output on abort (#15026)

* Gateway/Control UI: preserve partial output on abort

* fix: finalize abort partial handling and tests (#15026) (thanks @advaitpaliwal)

---------

Co-authored-by: Tyler Yust <TYTYYUST@YAHOO.COM>

src/gateway/server-methods/chat.ts

@@ -15,6 +15,7 @@ import { INTERNAL_MESSAGE_CHANNEL } from "../../utils/message-channel.js";
 import {
   abortChatRunById,
   abortChatRunsForSessionKey,
+  type ChatAbortControllerEntry,
   isChatStopCommandText,
   resolveChatRunExpiresAtMs,
 } from "../chat-abort.js";
@@ -49,6 +50,14 @@ type TranscriptAppendResult = {
 };
 
 type AppendMessageArg = Parameters<SessionManager["appendMessage"]>[0];
+type AbortOrigin = "rpc" | "stop-command";
+
+type AbortedPartialSnapshot = {
+  runId: string;
+  sessionId: string;
+  text: string;
+  abortOrigin: AbortOrigin;
+};
 
 function stripDisallowedChatControlChars(message: string): string {
   let output = "";
@@ -116,6 +125,24 @@ function ensureTranscriptFile(params: { transcriptPath: string; sessionId: strin
   }
 }
 
+function transcriptHasIdempotencyKey(transcriptPath: string, idempotencyKey: string): boolean {
+  try {
+    const lines = fs.readFileSync(transcriptPath, "utf-8").split(/\r?\n/);
+    for (const line of lines) {
+      if (!line.trim()) {
+        continue;
+      }
+      const parsed = JSON.parse(line) as { message?: { idempotencyKey?: unknown } };
+      if (parsed?.message?.idempotencyKey === idempotencyKey) {
+        return true;
+      }
+    }
+    return false;
+  } catch {
+    return false;
+  }
+}
+
 function appendAssistantTranscriptMessage(params: {
   message: string;
   label?: string;
@@ -124,6 +151,12 @@ function appendAssistantTranscriptMessage(params: {
   sessionFile?: string;
   agentId?: string;
   createIfMissing?: boolean;
+  idempotencyKey?: string;
+  abortMeta?: {
+    aborted: true;
+    origin: AbortOrigin;
+    runId: string;
+  };
 }): TranscriptAppendResult {
   const transcriptPath = resolveTranscriptPath({
     sessionId: params.sessionId,
@@ -148,6 +181,10 @@ function appendAssistantTranscriptMessage(params: {
     }
   }
 
+  if (params.idempotencyKey && transcriptHasIdempotencyKey(transcriptPath, params.idempotencyKey)) {
+    return { ok: true };
+  }
+
   const now = Date.now();
   const labelPrefix = params.label ? `[${params.label}]\n\n` : "";
   const usage = {
@@ -176,6 +213,16 @@ function appendAssistantTranscriptMessage(params: {
     api: "openai-responses",
     provider: "openclaw",
     model: "gateway-injected",
+    ...(params.idempotencyKey ? { idempotencyKey: params.idempotencyKey } : {}),
+    ...(params.abortMeta
+      ? {
+          openclawAbort: {
+            aborted: true,
+            origin: params.abortMeta.origin,
+            runId: params.abortMeta.runId,
+          },
+        }
+      : {}),
   };
 
   try {
@@ -189,6 +236,63 @@ function appendAssistantTranscriptMessage(params: {
   }
 }
 
+function collectSessionAbortPartials(params: {
+  chatAbortControllers: Map<string, ChatAbortControllerEntry>;
+  chatRunBuffers: Map<string, string>;
+  sessionKey: string;
+  abortOrigin: AbortOrigin;
+}): AbortedPartialSnapshot[] {
+  const out: AbortedPartialSnapshot[] = [];
+  for (const [runId, active] of params.chatAbortControllers) {
+    if (active.sessionKey !== params.sessionKey) {
+      continue;
+    }
+    const text = params.chatRunBuffers.get(runId);
+    if (!text || !text.trim()) {
+      continue;
+    }
+    out.push({
+      runId,
+      sessionId: active.sessionId,
+      text,
+      abortOrigin: params.abortOrigin,
+    });
+  }
+  return out;
+}
+
+function persistAbortedPartials(params: {
+  context: Pick<GatewayRequestContext, "logGateway">;
+  sessionKey: string;
+  snapshots: AbortedPartialSnapshot[];
+}) {
+  if (params.snapshots.length === 0) {
+    return;
+  }
+  const { storePath, entry } = loadSessionEntry(params.sessionKey);
+  for (const snapshot of params.snapshots) {
+    const sessionId = entry?.sessionId ?? snapshot.sessionId ?? snapshot.runId;
+    const appended = appendAssistantTranscriptMessage({
+      message: snapshot.text,
+      sessionId,
+      storePath,
+      sessionFile: entry?.sessionFile,
+      createIfMissing: true,
+      idempotencyKey: `${snapshot.runId}:assistant`,
+      abortMeta: {
+        aborted: true,
+        origin: snapshot.abortOrigin,
+        runId: snapshot.runId,
+      },
+    });
+    if (!appended.ok) {
+      params.context.logGateway.warn(
+        `chat.abort transcript append failed: ${appended.error ?? "unknown error"}`,
+      );
+    }
+  }
+}
+
 function nextChatSeq(context: { agentRunSeq: Map<string, number> }, runId: string) {
   const next = (context.agentRunSeq.get(runId) ?? 0) + 1;
   context.agentRunSeq.set(runId, next);
@@ -299,7 +403,7 @@ export const chatHandlers: GatewayRequestHandlers = {
       );
       return;
     }
-    const { sessionKey, runId } = params as {
+    const { sessionKey: rawSessionKey, runId } = params as {
       sessionKey: string;
       runId?: string;
     };
@@ -316,10 +420,23 @@ export const chatHandlers: GatewayRequestHandlers = {
     };
 
     if (!runId) {
+      const snapshots = collectSessionAbortPartials({
+        chatAbortControllers: context.chatAbortControllers,
+        chatRunBuffers: context.chatRunBuffers,
+        sessionKey: rawSessionKey,
+        abortOrigin: "rpc",
+      });
       const res = abortChatRunsForSessionKey(ops, {
-        sessionKey,
+        sessionKey: rawSessionKey,
         stopReason: "rpc",
       });
+      if (res.aborted) {
+        persistAbortedPartials({
+          context,
+          sessionKey: rawSessionKey,
+          snapshots,
+        });
+      }
       respond(true, { ok: true, aborted: res.aborted, runIds: res.runIds });
       return;
     }
@@ -329,7 +446,7 @@ export const chatHandlers: GatewayRequestHandlers = {
       respond(true, { ok: true, aborted: false, runIds: [] });
       return;
     }
-    if (active.sessionKey !== sessionKey) {
+    if (active.sessionKey !== rawSessionKey) {
       respond(
         false,
         undefined,
@@ -338,11 +455,26 @@ export const chatHandlers: GatewayRequestHandlers = {
       return;
     }
 
+    const partialText = context.chatRunBuffers.get(runId);
     const res = abortChatRunById(ops, {
       runId,
-      sessionKey,
+      sessionKey: rawSessionKey,
       stopReason: "rpc",
     });
+    if (res.aborted && partialText && partialText.trim()) {
+      persistAbortedPartials({
+        context,
+        sessionKey: rawSessionKey,
+        snapshots: [
+          {
+            runId,
+            sessionId: active.sessionId,
+            text: partialText,
+            abortOrigin: "rpc",
+          },
+        ],
+      });
+    }
     respond(true, {
       ok: true,
       aborted: res.aborted,
@@ -437,6 +569,12 @@ export const chatHandlers: GatewayRequestHandlers = {
     }
 
     if (stopCommand) {
+      const snapshots = collectSessionAbortPartials({
+        chatAbortControllers: context.chatAbortControllers,
+        chatRunBuffers: context.chatRunBuffers,
+        sessionKey: rawSessionKey,
+        abortOrigin: "stop-command",
+      });
       const res = abortChatRunsForSessionKey(
         {
           chatAbortControllers: context.chatAbortControllers,
@@ -450,6 +588,13 @@ export const chatHandlers: GatewayRequestHandlers = {
         },
         { sessionKey: rawSessionKey, stopReason: "stop" },
       );
+      if (res.aborted) {
+        persistAbortedPartials({
+          context,
+          sessionKey: rawSessionKey,
+          snapshots,
+        });
+      }
       respond(true, { ok: true, aborted: res.aborted, runIds: res.runIds });
       return;
     }