diff --git a/packages/ui/src/components/message-progress.tsx b/packages/ui/src/components/message-progress.tsx
index 7717e5e767..afd7f754ac 100644
--- a/packages/ui/src/components/message-progress.tsx
+++ b/packages/ui/src/components/message-progress.tsx
@@ -1,11 +1,27 @@
-import { For, JSXElement, Match, Show, Switch, createEffect, createMemo, createSignal, onCleanup } from "solid-js"
+import {
+  For,
+  JSXElement,
+  Match,
+  Show,
+  Switch,
+  ValidComponent,
+  createEffect,
+  createMemo,
+  createSignal,
+  onCleanup,
+} from "solid-js"
 import { Part } from "./message-part"
 import { Spinner } from "./spinner"
 import { useData } from "../context/data"
 import type { AssistantMessage as AssistantMessageType, ToolPart } from "@opencode-ai/sdk"
-import { Diff } from "./diff"
 
-export function MessageProgress(props: { assistantMessages: () => AssistantMessageType[]; done?: boolean }) {
+export interface MessageProgressProps {
+  assistantMessages: () => AssistantMessageType[]
+  diffComponent: ValidComponent
+  done?: boolean
+}
+
+export function MessageProgress(props: MessageProgressProps) {
   const data = useData()
   const sanitizer = createMemo(() => (data.directory ? new RegExp(`${data.directory}/`, "g") : undefined))
   const parts = createMemo(() => props.assistantMessages().flatMap((m) => data.store.part[m.id]))
@@ -156,7 +172,12 @@ export function MessageProgress(props: { assistantMessages: () => AssistantMessa
                       )
                       return (
                         <div data-slot="message-progress-item">
-                          <Part message={message()!} part={part} sanitize={sanitizer()} diffComponent={Diff} />
+                          <Part
+                            message={message()!}
+                            part={part}
+                            sanitize={sanitizer()}
+                            diffComponent={props.diffComponent}
+                          />
                         </div>
                       )
                     }}
diff --git a/packages/ui/src/components/session-turn.tsx b/packages/ui/src/components/session-turn.tsx
index 2dc83a9a73..963713a220 100644
--- a/packages/ui/src/components/session-turn.tsx
+++ b/packages/ui/src/components/session-turn.tsx
@@ -206,7 +206,11 @@ export function SessionTurn(
                 <div data-slot="session-turn-response-section">
                   <Switch>
                     <Match when={!completed()}>
-                      <MessageProgress assistantMessages={assistantMessages} done={!messageWorking()} />
+                      <MessageProgress
+                        assistantMessages={assistantMessages}
+                        done={!messageWorking()}
+                        diffComponent={props.diffComponent}
+                      />
                     </Match>
                     <Match when={completed() && hasToolPart()}>
                       <Collapsible variant="ghost" open={detailsExpanded()} onOpenChange={setDetailsExpanded}>