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
         ? [