fix(core): persist assistant text state

This commit is contained in:
Aiden Cline 2026-07-23 14:05:13 +00:00 committed by 𝓛𝓲𝓽𝓽𝓵𝓮 𝓕𝓻𝓪𝓷𝓴
commit ce1203ce83
8 changed files with 144 additions and 36 deletions

View file

@ -307,7 +307,9 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
},
"session.text.started": (event) => {
return updateOwnedAssistant(event.data.assistantMessageID, (draft) => {
draft.content.push(castDraft(SessionMessage.AssistantText.make({ type: "text", text: "" })))
draft.content.push(
castDraft(SessionMessage.AssistantText.make({ type: "text", text: "", state: event.data.state })),
)
})
},
"session.text.delta": (event) => {
@ -319,7 +321,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
if (event.data.state !== undefined) match.state = event.data.state
}
})
},
"session.tool.input.started": (event) => {

View file

@ -142,13 +142,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,
@ -317,15 +318,16 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
return
case "text-start":
retryEvidence = 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(),
ordinal: startedTextOrdinal,
state: providerState(event.providerMetadata),
})
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(),
@ -334,7 +336,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":
retryEvidence = true

View file

@ -123,7 +123,14 @@ const assistant = (message: SessionMessage.Assistant, model: ModelV2.Ref, provid
String(message.model.providerID) === String(model.providerID) && 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: reuseProviderMetadata ? providerMetadata(providerMetadataKey, item.state) : undefined,
},
]
if (item.type === "reasoning")
return reuseProviderMetadata
? [

View file

@ -15,6 +15,36 @@ const model = ModelV2.Ref.make({ id: ModelV2.ID.make("model"), providerID: Provi
const build = AgentV2.defaultID
describe("toLLMMessages", () => {
test("restores same-model provider state on assistant text", () => {
const messages = toLLMMessages(
[
SessionMessage.Assistant.make({
id: id("assistant-text-state"),
type: "assistant",
agent: build,
model,
content: [
SessionMessage.AssistantText.make({
type: "text",
text: "Checking.",
state: { itemId: "msg_commentary", phase: "commentary" },
}),
],
time: { created, completed: created },
}),
],
model,
)
expect(messages[0]?.content).toEqual([
{
type: "text",
text: "Checking.",
providerMetadata: { provider: { itemId: "msg_commentary", phase: "commentary" } },
},
])
})
test("omits empty assistant turns", () => {
const assistant = (value: string, content: SessionMessage.Assistant["content"]) =>
SessionMessage.Assistant.make({
@ -661,6 +691,11 @@ Recent work
agent: build,
model: { id: ModelV2.ID.make("old-model"), providerID: ProviderV2.ID.make("provider") },
content: [
SessionMessage.AssistantText.make({
type: "text",
text: "Checking",
state: { itemId: "msg_old", phase: "commentary" },
}),
SessionMessage.AssistantReasoning.make({
type: "reasoning",
text: "Visible thought",
@ -705,6 +740,7 @@ Recent work
)
expect(messages[0]?.content).toEqual([
{ type: "text", text: "Checking", providerMetadata: undefined },
{ type: "text", text: "Visible thought" },
{
type: "tool-call",

View file

@ -2434,6 +2434,51 @@ describe("SessionRunnerLLM", () => {
}),
)
it.effect("preserves assistant text provider state across durable continuation", () =>
Effect.gen(function* () {
const session = yield* setup
yield* admit(session, "Check this")
const providerMetadata = { openai: { itemId: "msg_commentary", phase: "commentary" } }
responses = [
[
LLMEvent.stepStart({ index: 0 }),
LLMEvent.textStart({ id: "msg_commentary", providerMetadata }),
LLMEvent.textDelta({ id: "msg_commentary", text: "Checking.", providerMetadata }),
LLMEvent.textEnd({ id: "msg_commentary", providerMetadata }),
LLMEvent.toolCall({ id: "call-echo", name: "echo", input: { text: "hello" } }),
LLMEvent.stepFinish({ index: 0, reason: "tool-calls" }),
LLMEvent.finish({ reason: "tool-calls" }),
],
reply.text("Done", "text-final"),
]
yield* session.resume(sessionID)
expect(requests).toHaveLength(2)
expect(requests[1]?.messages[1]?.content[0]).toEqual({
type: "text",
text: "Checking.",
providerMetadata,
})
expect(yield* session.context(sessionID)).toMatchObject([
{ type: "user", text: "Check this" },
{
type: "assistant",
content: [
{
type: "text",
text: "Checking.",
state: { itemId: "msg_commentary", phase: "commentary" },
},
{ type: "tool", id: "call-echo" },
],
},
{ type: "assistant", content: [{ type: "text", text: "Done" }] },
])
}),
)
it.effect("reloads a model switch before a tool-driven continuation step", () =>
Effect.gen(function* () {
const session = yield* setup