refactor(session): make message ordering explicit

This commit is contained in:
Aiden Cline 2026-05-23 17:53:32 -05:00
commit 75144fedf6
2 changed files with 35 additions and 21 deletions

View file

@ -579,9 +579,10 @@ export const cursor = {
}, },
} }
const messageOrder = new WeakMap<Info, number>() const chronologicalOrder = Symbol("chronologicalOrder")
const messageRowID = sql<number>`rowid` const messageRowID = sql<number>`rowid`
type MessageRow = typeof MessageTable.$inferSelect & { sequence?: number } type MessageRow = typeof MessageTable.$inferSelect & { sequence?: number }
type Chronological = WithParts & { [chronologicalOrder]?: number }
const info = (row: MessageRow) => const info = (row: MessageRow) =>
({ ({
@ -1045,7 +1046,7 @@ export function filterCompacted(msgs: Iterable<WithParts>) {
completed.add(msg.info.parentID) completed.add(msg.info.parentID)
} }
result.reverse() result.reverse()
result.forEach((msg, index) => messageOrder.set(msg.info, index)) result.forEach((msg, index) => ((msg as Chronological)[chronologicalOrder] = index))
const compactionIndex = result.findLastIndex( const compactionIndex = result.findLastIndex(
(msg) => (msg) =>
msg.info.role === "user" && msg.info.role === "user" &&
@ -1079,10 +1080,10 @@ export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: Ses
return filterCompacted(stream(sessionID)) return filterCompacted(stream(sessionID))
}) })
export function compare(a: Info, b: Info, indexA = -1, indexB = -1) { export function compare(a: WithParts, b: WithParts, indexA = -1, indexB = -1) {
if (a.time.created !== b.time.created) return a.time.created - b.time.created if (a.info.time.created !== b.info.time.created) return a.info.time.created - b.info.time.created
const sequenceA = messageOrder.get(a) const sequenceA = (a as Chronological)[chronologicalOrder]
const sequenceB = messageOrder.get(b) const sequenceB = (b as Chronological)[chronologicalOrder]
if (sequenceA !== undefined && sequenceB !== undefined && sequenceA !== sequenceB) return sequenceA - sequenceB if (sequenceA !== undefined && sequenceB !== undefined && sequenceA !== sequenceB) return sequenceA - sequenceB
return indexA - indexB return indexA - indexB
} }
@ -1095,33 +1096,41 @@ export function compare(a: Info, b: Info, indexA = -1, indexB = -1) {
// tasks are compaction/subtask parts attached to user messages newer than the // tasks are compaction/subtask parts attached to user messages newer than the
// latest finished assistant — i.e. unprocessed work. // latest finished assistant — i.e. unprocessed work.
export function latest(msgs: WithParts[]) { export function latest(msgs: WithParts[]) {
let user: User | undefined let user: WithParts | undefined
let assistant: Assistant | undefined let assistant: WithParts | undefined
let finished: Assistant | undefined let finished: WithParts | undefined
let userIndex = -1 let userIndex = -1
let assistantIndex = -1 let assistantIndex = -1
let finishedIndex = -1 let finishedIndex = -1
for (const [index, msg] of msgs.entries()) { for (const [index, msg] of msgs.entries()) {
const info = msg.info const info = msg.info
if (info.role === "user" && (!user || compare(info, user, index, userIndex) > 0)) { if (info.role === "user" && (!user || compare(msg, user, index, userIndex) > 0)) {
user = info user = msg
userIndex = index userIndex = index
} }
if (info.role === "assistant" && (!assistant || compare(info, assistant, index, assistantIndex) > 0)) { if (info.role === "assistant" && (!assistant || compare(msg, assistant, index, assistantIndex) > 0)) {
assistant = info assistant = msg
assistantIndex = index assistantIndex = index
} }
if (info.role === "assistant" && info.finish && (!finished || compare(info, finished, index, finishedIndex) > 0)) { if (info.role === "assistant" && info.finish && (!finished || compare(msg, finished, index, finishedIndex) > 0)) {
finished = info finished = msg
finishedIndex = index finishedIndex = index
} }
} }
const tasks = msgs.flatMap((m, index) => const tasks = msgs.flatMap((m, index) =>
finished && compare(m.info, finished, index, finishedIndex) <= 0 finished && compare(m, finished, index, finishedIndex) <= 0
? [] ? []
: m.parts.filter((p): p is CompactionPart | SubtaskPart => p.type === "compaction" || p.type === "subtask"), : m.parts.filter((p): p is CompactionPart | SubtaskPart => p.type === "compaction" || p.type === "subtask"),
) )
return { user, assistant, finished, tasks } return {
user: user?.info.role === "user" ? user.info : undefined,
assistant: assistant?.info.role === "assistant" ? assistant.info : undefined,
finished: finished?.info.role === "assistant" ? finished.info : undefined,
userMessage: user,
assistantMessage: assistant,
finishedMessage: finished,
tasks,
}
} }
export function fromError( export function fromError(

View file

@ -1250,14 +1250,18 @@ export const layer = Layer.effect(
let msgs = yield* MessageV2.filterCompactedEffect(sessionID) let msgs = yield* MessageV2.filterCompactedEffect(sessionID)
const { user: lastUser, assistant: lastAssistant, finished: lastFinished, tasks } = MessageV2.latest(msgs) const latest = MessageV2.latest(msgs)
const before = (a: MessageV2.Info, b: MessageV2.Info) => MessageV2.compare(a, b) < 0 const { user: lastUser, assistant: lastAssistant, finished: lastFinished, tasks } = latest
if (!lastUser) throw new Error("No user message found in stream. This should never happen.") if (!lastUser) throw new Error("No user message found in stream. This should never happen.")
const lastAssistantMsg = msgs.findLast( const lastAssistantMsg = msgs.findLast(
(msg) => msg.info.role === "assistant" && msg.info.id === lastAssistant?.id, (msg) => msg.info.role === "assistant" && msg.info.id === lastAssistant?.id,
) )
const userBeforeAssistant =
latest.userMessage &&
latest.assistantMessage &&
MessageV2.compare(latest.userMessage, latest.assistantMessage) < 0
// Some providers return "stop" even when the assistant message contains tool calls. // Some providers return "stop" even when the assistant message contains tool calls.
// Keep the loop running so tool results can be sent back to the model. // Keep the loop running so tool results can be sent back to the model.
// Skip provider-executed tool parts — those were fully handled within the // Skip provider-executed tool parts — those were fully handled within the
@ -1269,7 +1273,7 @@ export const layer = Layer.effect(
lastAssistant?.finish && lastAssistant?.finish &&
!["tool-calls"].includes(lastAssistant.finish) && !["tool-calls"].includes(lastAssistant.finish) &&
!hasToolCalls && !hasToolCalls &&
before(lastUser, lastAssistant) userBeforeAssistant
) { ) {
yield* slog.info("exiting loop") yield* slog.info("exiting loop")
break break
@ -1399,7 +1403,8 @@ export const layer = Layer.effect(
if (step > 1 && lastFinished) { if (step > 1 && lastFinished) {
for (const m of msgs) { for (const m of msgs) {
if (m.info.role !== "user" || !before(lastFinished, m.info)) continue const finishedBeforeMessage = latest.finishedMessage && MessageV2.compare(latest.finishedMessage, m) < 0
if (m.info.role !== "user" || !finishedBeforeMessage) continue
for (const p of m.parts) { for (const p of m.parts) {
if (p.type !== "text" || p.ignored || p.synthetic) continue if (p.type !== "text" || p.ignored || p.synthetic) continue
if (!p.text.trim()) continue if (!p.text.trim()) continue