fix: hide provider errors after session revert

2026-08-28 19:11:03 +08:00 · 2026-06-19 11:51:50 +02:00
parent 6f86c95279
commit 4d0933371c
9 changed files with 218 additions and 29 deletions
@@ -0,0 +1,6 @@
+---
+"@kilocode/cli": patch
+"kilo-code": patch
+---
+
+Hide reverted provider errors so Redo controls remain visible after rewinding a session.
@@ -6,8 +6,10 @@ import {
  queuedUserMessageIDs,
  stableMessageTurns,
  visibleMessages,
+  visibleParts,
+  type RevertBoundary,
 } from "../../webview-ui/src/context/session-queue"
-import type { Message, SessionStatusInfo } from "../../webview-ui/src/types/messages"
+import type { Message, Part, SessionStatusInfo } from "../../webview-ui/src/types/messages"

 const base = {
  sessionID: "session",
@@ -30,10 +32,12 @@ const assistant = (id: string, parentID: string, opts: Partial<Message> = {}): M
  ...opts,
 })

-const layout = (messages: Message[], status: SessionStatusInfo, boundary?: string) => {
+const part = (id: string, messageID: string): Part => ({ id, messageID, type: "text", text: id })
+
+const layout = (messages: Message[], status: SessionStatusInfo, revert?: RevertBoundary) => {
  const active = activeUserMessageID(messages, status)
  return partitionTurns(
-    messageTurns(messages, boundary),
+    messageTurns(messages, revert),
    new Set(active ? [active] : []),
    new Set(queuedUserMessageIDs(messages, status)),
  )
@@ -292,7 +296,7 @@ describe("partitionTurns", () => {

  it("does not render an active turn hidden by a revert boundary", () => {
    const messages = [user("message_1"), assistant("message_2", "message_1", { finish: "stop" }), user("message_3")]
-    const result = layout(messages, { type: "busy" }, "message_3")
+    const result = layout(messages, { type: "busy" }, { messageID: "message_3" })

    expect(result.virtual.map((turn) => turn.user.id)).toEqual(["message_1"])
    expect(result.direct).toEqual([])
@@ -394,7 +398,45 @@ describe("messageTurns", () => {
      assistant("message_4", "message_3"),
    ]

-    expect(messageTurns(messages, "message_3").map((turn) => turn.user.id)).toEqual(["message_1"])
+    expect(messageTurns(messages, { messageID: "message_3" }).map((turn) => turn.user.id)).toEqual(["message_1"])
+  })
+
+  it("keeps the part-boundary assistant and hides later provider errors", () => {
+    const messages = [
+      user("message_1"),
+      assistant("message_2", "message_1"),
+      assistant("message_3", "message_1", { error: { name: "ProviderError" } }),
+      assistant("message_4", "message_1", { error: { name: "ProviderError" } }),
+    ]
+    const turns = messageTurns(messages, { messageID: "message_2", partID: "part_2" })
+
+    expect(turns).toHaveLength(1)
+    expect(turns[0]?.assistant.map((msg) => msg.id)).toEqual(["message_2"])
+  })
+
+  it("applies assistant boundaries by id when messages arrive out of order", () => {
+    const messages = [
+      user("message_1"),
+      assistant("message_4", "message_1", { error: { name: "ProviderError" } }),
+      assistant("message_2", "message_1"),
+    ]
+    const turns = messageTurns(messages, { messageID: "message_2", partID: "part_2" })
+
+    expect(turns[0]?.assistant.map((msg) => msg.id)).toEqual(["message_2"])
+  })
+})
+
+describe("visibleParts", () => {
+  it("keeps only parts before the active part boundary", () => {
+    const parts = [part("part_1", "message_2"), part("part_2", "message_2"), part("part_3", "message_2")]
+
+    expect(visibleParts("message_2", parts, { messageID: "message_2", partID: "part_2" })).toEqual([parts[0]])
+  })
+
+  it("fails closed when the boundary part is unavailable", () => {
+    const parts = [part("part_1", "message_2")]
+
+    expect(visibleParts("message_2", parts, { messageID: "message_2", partID: "part_missing" })).toEqual([])
  })
 })

@@ -407,7 +449,10 @@ describe("visibleMessages", () => {
      assistant("message_4", "message_3"),
    ]

-    expect(visibleMessages(messages, "message_3").map((msg) => msg.id)).toEqual(["message_1", "message_2"])
+    expect(visibleMessages(messages, { messageID: "message_3" }).map((msg) => msg.id)).toEqual([
+      "message_1",
+      "message_2",
+    ])
  })

  it("keeps leading partial assistant output", () => {
@@ -84,6 +84,32 @@ describe("transcriptRows", () => {
    expect(rows.at(-1)).toMatchObject({ type: "error", message: a3, error: a3.error })
  })

+  it("hides provider errors at and after an assistant part boundary", () => {
+    const revert = { messageID: "message_2", partID: "part_2" }
+    const u1 = user("message_1")
+    const a1 = assistant("message_2", "message_1", { error: { name: "ProviderError" } })
+    const a2 = assistant("message_3", "message_1", { error: { name: "ProviderError" } })
+    const parts = { message_2: [part("part_1", "message_2"), part("part_2", "message_2")] }
+    const rows = transcriptRows(messageTurns([u1, a1, a2], revert), lookup(parts), { revert })
+
+    expect(rows.map((row) => row.type)).toEqual(["user", "assistant"])
+    expect(rows.filter((row) => row.type === "assistant").flatMap((row) => row.parts.map((item) => item.id))).toEqual([
+      "part_1",
+    ])
+    expect(rows.some((row) => row.message.id === "message_3")).toBe(false)
+  })
+
+  it("keeps provider errors before an assistant part boundary", () => {
+    const revert = { messageID: "message_3", partID: "part_2" }
+    const u1 = user("message_1")
+    const a1 = assistant("message_2", "message_1", { error: { name: "ProviderError" } })
+    const a2 = assistant("message_3", "message_1")
+    const parts = { message_3: [part("part_1", "message_3"), part("part_2", "message_3")] }
+    const rows = transcriptRows(messageTurns([u1, a1, a2], revert), lookup(parts), { revert })
+
+    expect(rows.at(-1)).toMatchObject({ type: "error", message: a1 })
+  })
+
  it("keeps keys stable when older turns are prepended and parts are appended", () => {
    const u1 = user("u1")
    const a1 = assistant("a1", "u1")
@@ -131,7 +157,7 @@ describe("transcriptRows", () => {
    })
    const a2 = assistant("a2", "u1")
    const u3 = user("u3")
-    const turns = messageTurns([u1, a1, u2, a2, u3], "u3")
+    const turns = messageTurns([u1, a1, u2, a2, u3], { messageID: "u3" })
    const rows = transcriptRows(turns, (id) => (id === "u2" ? (u2.parts ?? []) : []))

    expect(rows.map((row) => `${row.turn}:${row.message.id}`)).toEqual(["u1:u1", "u1:a1", "u2:u2", "u2:a2"])
@@ -91,14 +91,14 @@ export const MessageList: Component<MessageListProps> = (props) => {
  const [virtualizer, setVirtualizer] = createSignal<VirtualizerHandle>()
  const [layout, setLayout] = createSignal("")

-  const boundary = () => session.revert()?.messageID
+  const revert = () => session.revert() ?? undefined
  const turns = createMemo((prev: MessageTurn[] | undefined) =>
    stableMessageTurns(
-      messageTurns(session.messages(), boundary(), (msg) => session.getParts(msg.id)),
+      messageTurns(session.messages(), revert(), (msg) => session.getParts(msg.id)),
      prev,
    ),
  )
-  const isEmpty = () => turns().length === 0 && !session.loading() && !boundary()
+  const isEmpty = () => turns().length === 0 && !session.loading() && !revert()

  const activeUserID = createMemo(() =>
    getActiveUserMessageID(session.messages(), session.statusInfo(), (msg) => session.getParts(msg.id)),
@@ -115,6 +115,7 @@ export const MessageList: Component<MessageListProps> = (props) => {
        queued: queuedIDs(),
        live: new Set(active ? [active] : []),
        hidden: session.isErrorHidden,
+        revert: revert(),
      },
      prev,
    )
@@ -326,7 +327,7 @@ export const MessageList: Component<MessageListProps> = (props) => {
                </For>
              </div>
            </Show>
-            <Show when={boundary()}>
+            <Show when={revert()}>
              <RevertBanner />
            </Show>
            <For each={partition().queued}>{(row) => <TranscriptRowView row={row} />}</For>
@@ -1,4 +1,6 @@
-import type { Message, SessionStatusInfo } from "../types/messages"
+import type { Message, Part, SessionInfo, SessionStatusInfo } from "../types/messages"
+
+export type RevertBoundary = Pick<NonNullable<SessionInfo["revert"]>, "messageID" | "partID">

 export interface MessageTurn {
  id: string
@@ -55,24 +57,37 @@ function target(messages: Message[], index: number, id: string, parts?: (msg: Me
  return id
 }

+function visibleMessage(id: string, revert?: RevertBoundary) {
+  if (!revert || id < revert.messageID) return true
+  return id === revert.messageID && !!revert.partID
+}
+
+export function visibleParts(id: string, parts: Part[], revert?: RevertBoundary) {
+  if (!revert || id < revert.messageID) return parts
+  if (id !== revert.messageID || !revert.partID) return []
+  const idx = parts.findIndex((part) => part.id === revert.partID)
+  return idx < 0 ? [] : parts.slice(0, idx)
+}
+
 export function messageTurns(
  messages: Message[],
-  boundary?: string,
+  revert?: RevertBoundary,
  parts?: (msg: Message) => Message["parts"],
 ): MessageTurn[] {
  const result: MessageTurn[] = []
  const lead: Message[] = []
  const by = new Map<string, { turn: MessageTurn; index: number }>()
+  const projected = (msg: Message) => visibleParts(msg.id, parts?.(msg) ?? msg.parts ?? [], revert)
  let compact: { turn: MessageTurn; index: number } | undefined

  for (const msg of messages) {
+    if (!visibleMessage(msg.id, revert)) continue
    if (msg.role === "user") {
-      if (boundary && msg.id >= boundary) break
      const turn = { id: msg.id, user: msg, assistant: [] }
      const item = { turn, index: result.length }
      result.push(turn)
      by.set(msg.id, item)
-      if (isCompact(msg, parts)) compact = item
+      if (isCompact(msg, projected)) compact = item
      continue
    }

@@ -105,10 +120,10 @@ export function messageTurns(

 export function visibleMessages(
  messages: Message[],
-  boundary?: string,
+  revert?: RevertBoundary,
  parts?: (msg: Message) => Message["parts"],
 ): Message[] {
-  return messageTurns(messages, boundary, parts).flatMap((turn) =>
+  return messageTurns(messages, revert, parts).flatMap((turn) =>
    turn.partial ? turn.assistant : [turn.user, ...turn.assistant],
  )
 }
@@ -2550,8 +2550,10 @@ export const SessionProvider: ParentComponent = (props) => {
  const userMessages = createMemo(() => messages().filter((m) => m.role === "user"))

  function visible(sessionID: string) {
-    return filterVisibleMessages(store.messages[sessionID] ?? [], store.sessions[sessionID]?.revert?.messageID, (msg) =>
-      getParts(msg.id),
+    return filterVisibleMessages(
+      store.messages[sessionID] ?? [],
+      store.sessions[sessionID]?.revert ?? undefined,
+      (msg) => getParts(msg.id),
    )
  }

@@ -1,5 +1,5 @@
 import type { Message, Part } from "../types/messages"
-import type { MessageTurn } from "./session-queue"
+import { visibleParts, type MessageTurn, type RevertBoundary } from "./session-queue"

 interface TranscriptMeta {
  turn: string
@@ -46,6 +46,7 @@ export interface TranscriptOptions {
  queued?: ReadonlySet<string>
  live?: ReadonlySet<string>
  hidden?: (id: string) => boolean
+  revert?: RevertBoundary
 }

 export interface TranscriptPartition {
@@ -127,6 +128,8 @@ export function transcriptRows(
 ): TranscriptRow[] {
  const size = Math.max(1, Math.floor(opts.size ?? 8))
  const rows: TranscriptRow[] = []
+  const parts = (id: string) => visibleParts(id, getParts(id), opts.revert)
+  const terminal = (msg: Message) => !(opts.revert?.partID && msg.id === opts.revert.messageID)

  for (const turn of turns) {
    const meta = {
@@ -135,7 +138,7 @@ export function transcriptRows(
      queued: opts.queued?.has(turn.id) === true,
      live: opts.live?.has(turn.id) === true,
    }
-    const copied = copy(turn.assistant, getParts)
+    const copied = copy(turn.assistant, parts)

    if (!turn.partial) {
      rows.push({
@@ -143,27 +146,27 @@ export function transcriptRows(
        type: "user",
        key: `${turn.id}:user`,
        message: turn.user,
-        parts: getParts(turn.user.id),
-        interrupted: turn.assistant.some((msg) => msg.error?.name === "MessageAbortedError"),
+        parts: parts(turn.user.id),
+        interrupted: turn.assistant.some((msg) => terminal(msg) && msg.error?.name === "MessageAbortedError"),
        answered: turn.assistant.length > 0,
      })
    }

    for (const msg of turn.assistant) {
-      const parts = getParts(msg.id)
-      if (parts.length === 0) {
+      const visible = parts(msg.id)
+      if (visible.length === 0) {
        rows.push({
          ...meta,
          type: "assistant",
          key: `${turn.id}:assistant:${msg.id}:empty`,
          message: msg,
-          parts,
+          parts: visible,
          copy: copied,
        })
        continue
      }
-      for (let start = 0; start < parts.length; start += size) {
-        const chunk = parts.slice(start, start + size)
+      for (let start = 0; start < visible.length; start += size) {
+        const chunk = visible.slice(start, start + size)
        rows.push({
          ...meta,
          type: "assistant",
@@ -181,7 +184,7 @@ export function transcriptRows(
    }

    const failed = turn.assistant.find(
-      (msg) => msg.error && msg.error.name !== "MessageAbortedError" && opts.hidden?.(msg.id) !== true,
+      (msg) => terminal(msg) && msg.error && msg.error.name !== "MessageAbortedError" && opts.hidden?.(msg.id) !== true,
    )
    if (failed?.error) {
      rows.push({ ...meta, type: "error", key: `${turn.id}:error:${failed.id}`, message: failed, error: failed.error })
@@ -152,6 +152,12 @@ export const layer = Layer.effect(
              partID: part.id,
            })
          }
+          // kilocode_change start - clear a reverted provider error from the retained assistant message
+          if (target.info.role === "assistant" && target.info.error) {
+            delete target.info.error
+            yield* sessions.updateMessage(target.info)
+          }
+          // kilocode_change end
        }
      }
      yield* sessions.clearRevert(sessionID)
@@ -0,0 +1,85 @@
+import { describe, expect } from "bun:test"
+import { Effect, Layer } from "effect"
+import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
+import { ModelID, ProviderID } from "@/provider/schema"
+import { MessageV2 } from "@/session/message-v2"
+import { SessionRevert } from "@/session/revert"
+import { MessageID, PartID } from "@/session/schema"
+import { Session } from "@/session/session"
+import { Snapshot } from "@/snapshot"
+import { provideTmpdirInstance } from "../../fixture/fixture"
+import { testEffect } from "../../lib/effect"
+
+const env = Layer.mergeAll(
+  Session.defaultLayer,
+  SessionRevert.defaultLayer,
+  Snapshot.defaultLayer,
+  CrossSpawnSpawner.defaultLayer,
+)
+const it = testEffect(env)
+
+describe("partial assistant revert", () => {
+  it.live(
+    "clears provider errors when the revert becomes permanent",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const sessions = yield* Session.Service
+          const revert = yield* SessionRevert.Service
+          const session = yield* sessions.create({})
+          const providerID = ProviderID.make("test")
+          const user = yield* sessions.updateMessage({
+            id: MessageID.ascending(),
+            sessionID: session.id,
+            role: "user",
+            agent: "default",
+            model: { providerID, modelID: ModelID.make("test") },
+            time: { created: Date.now() },
+          })
+          const assistant = yield* sessions.updateMessage({
+            id: MessageID.ascending(),
+            sessionID: session.id,
+            role: "assistant",
+            parentID: user.id,
+            mode: "default",
+            agent: "default",
+            path: { cwd: dir, root: dir },
+            cost: 1,
+            tokens: { input: 1, output: 1, reasoning: 0, cache: { read: 0, write: 0 } },
+            modelID: ModelID.make("test"),
+            providerID,
+            time: { created: Date.now(), completed: Date.now() },
+            finish: "error",
+            error: MessageV2.fromError(new Error("Provider returned error"), { providerID }),
+          })
+          const kept = yield* sessions.updatePart({
+            id: PartID.ascending(),
+            messageID: assistant.id,
+            sessionID: session.id,
+            type: "text",
+            text: "keep",
+          })
+          const boundary = yield* sessions.updatePart({
+            id: PartID.ascending(),
+            messageID: assistant.id,
+            sessionID: session.id,
+            type: "text",
+            text: "remove",
+          })
+
+          yield* sessions.setRevert({
+            sessionID: session.id,
+            revert: { messageID: assistant.id, partID: boundary.id },
+            summary: { additions: 0, deletions: 0, files: 0 },
+          })
+          yield* revert.cleanup(yield* sessions.get(session.id))
+
+          const messages = yield* sessions.messages({ sessionID: session.id })
+          const result = messages.find((message) => message.info.id === assistant.id)
+          expect(result?.parts.map((part) => part.id)).toEqual([kept.id])
+          expect(result?.info).not.toHaveProperty("error")
+        }),
+      { git: true },
+    ),
+  )
+})