fix(ai): preserve response message phases (#38777)

packages/core/src/session/message-updater.ts

@@ -319,7 +319,10 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
       "session.text.ended": (event) => {
         return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
           const match = latestText(draft)
-          if (match) match.text = event.data.text
+          if (match) {
+            match.text = event.data.text
+            match.state = castDraft(event.data.state)
+          }
         })
       },
       "session.tool.input.started": (event) => {

packages/core/src/session/runner/publish-llm-event.ts

@@ -149,13 +149,14 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
 
   const text = fragments(
     "text",
-    (_textID, value, ordinal) =>
+    (_textID, value, ordinal, state) =>
       Effect.gen(function* () {
         yield* events.publish(SessionEvent.Text.Ended, {
           sessionID: input.sessionID,
           assistantMessageID: yield* currentAssistantMessageID(),
           ordinal,
           text: value,
+          state,
         })
       }),
     true,
@@ -326,7 +327,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
         return
       case "text-start":
         outputStarted = true
-        const startedTextOrdinal = yield* text.start(event.id)
+        const startedTextOrdinal = yield* text.start(event.id, providerState(event.providerMetadata))
         yield* events.publish(SessionEvent.Text.Started, {
           sessionID: input.sessionID,
           assistantMessageID: yield* startAssistant(),
@@ -334,7 +335,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
         })
         return
       case "text-delta":
-        const deltaTextOrdinal = yield* text.append(event.id, event.text)
+        const deltaTextOrdinal = yield* text.append(event.id, event.text, providerState(event.providerMetadata))
         yield* events.publish(SessionEvent.Text.Delta, {
           sessionID: input.sessionID,
           assistantMessageID: yield* currentAssistantMessageID(),
@@ -343,7 +344,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
         })
         return
       case "text-end":
-        yield* text.end(event.id)
+        yield* text.end(event.id, providerState(event.providerMetadata))
         return
       case "reasoning-start":
         outputStarted = true

packages/core/src/session/runner/to-llm-message.ts

@@ -113,7 +113,14 @@ const assistant = (message: SessionMessage.Assistant, model: ModelV2.Ref, provid
   const sameModel = sameProvider && String(message.model.id) === String(model.id)
   const reuseProviderMetadata = sameModel && message.error === undefined
   const content = message.content.flatMap((item): ContentPart[] => {
-    if (item.type === "text") return [{ type: "text", text: item.text }]
+    if (item.type === "text")
+      return [
+        {
+          type: "text",
+          text: item.text,
+          providerMetadata: sameProvider ? providerMetadata(providerMetadataKey, item.state) : undefined,
+        },
+      ]
     if (item.type === "reasoning")
       return reuseProviderMetadata
         ? [

packages/core/test/session-runner-message.test.ts

@@ -767,4 +767,35 @@ Recent work
       },
     ])
   })
+
+  test("preserves assistant text provider state across same-provider model changes and failures", () => {
+    const messages = toLLMMessages(
+      [
+        SessionMessage.Assistant.make({
+          id: id("assistant-phase"),
+          type: "assistant",
+          agent: build,
+          model: { id: ModelV2.ID.make("old"), providerID: ProviderV2.ID.make("provider") },
+          content: [
+            SessionMessage.AssistantText.make({
+              type: "text",
+              text: "Checking.",
+              state: { phase: "commentary" },
+            }),
+          ],
+          error: { type: "provider.unknown", message: "Interrupted after commentary" },
+          time: { created, completed: created },
+        }),
+      ],
+      ModelV2.Ref.make({ id: ModelV2.ID.make("new"), providerID: ProviderV2.ID.make("provider") }),
+    )
+
+    expect(messages[0]?.content).toEqual([
+      {
+        type: "text",
+        text: "Checking.",
+        providerMetadata: { provider: { phase: "commentary" } },
+      },
+    ])
+  })
 })

packages/core/test/session-runner.test.ts

@@ -2672,6 +2672,47 @@ describe("SessionRunnerLLM", () => {
     }),
   )
 
+  it.effect("restores durable text provider metadata in the next request", () =>
+    Effect.gen(function* () {
+      const session = yield* setup
+      yield* admit(session, "Check first")
+
+      response = [
+        LLMEvent.stepStart({ index: 0 }),
+        LLMEvent.textStart({ id: "commentary", providerMetadata: { openai: { phase: "commentary" } } }),
+        LLMEvent.textDelta({ id: "commentary", text: "Checking." }),
+        LLMEvent.textEnd({
+          id: "commentary",
+          providerMetadata: { openai: { phase: "commentary" }, anthropic: { ignored: true } },
+        }),
+        LLMEvent.stepFinish({ index: 0, reason: { normalized: "stop" } }),
+        LLMEvent.finish({ reason: { normalized: "stop" } }),
+      ]
+      yield* session.resume(sessionID)
+      yield* replaySessionProjection(sessionID)
+
+      expect(yield* session.context(sessionID)).toMatchObject([
+        { type: "user", text: "Check first" },
+        {
+          type: "assistant",
+          content: [{ type: "text", text: "Checking.", state: { phase: "commentary" } }],
+        },
+      ])
+
+      yield* admit(session, "Continue")
+      response = []
+      yield* session.resume(sessionID)
+
+      expect(requests[1]?.messages[1]?.content).toEqual([
+        {
+          type: "text",
+          text: "Checking.",
+          providerMetadata: { openai: { phase: "commentary" } },
+        },
+      ])
+    }),
+  )
+
   it.effect("replays durable provider-executed tool results inline in the next request", () =>
     Effect.gen(function* () {
       const session = yield* setup