refactor(session): make message ordering explicit

This commit is contained in:
Aiden Cline
2026-05-23 17:53:32 -05:00
parent 46ea66112e
commit 75144fedf6
2 changed files with 35 additions and 21 deletions
+26 -17
View File
@@ -579,9 +579,10 @@ export const cursor = {
},
}
const messageOrder = new WeakMap<Info, number>()
const chronologicalOrder = Symbol("chronologicalOrder")
const messageRowID = sql<number>`rowid`
type MessageRow = typeof MessageTable.$inferSelect & { sequence?: number }
type Chronological = WithParts & { [chronologicalOrder]?: number }
const info = (row: MessageRow) =>
({
@@ -1045,7 +1046,7 @@ export function filterCompacted(msgs: Iterable<WithParts>) {
completed.add(msg.info.parentID)
}
result.reverse()
result.forEach((msg, index) => messageOrder.set(msg.info, index))
result.forEach((msg, index) => ((msg as Chronological)[chronologicalOrder] = index))
const compactionIndex = result.findLastIndex(
(msg) =>
msg.info.role === "user" &&
@@ -1079,10 +1080,10 @@ export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: Ses
return filterCompacted(stream(sessionID))
})
export function compare(a: Info, b: Info, indexA = -1, indexB = -1) {
if (a.time.created !== b.time.created) return a.time.created - b.time.created
const sequenceA = messageOrder.get(a)
const sequenceB = messageOrder.get(b)
export function compare(a: WithParts, b: WithParts, indexA = -1, indexB = -1) {
if (a.info.time.created !== b.info.time.created) return a.info.time.created - b.info.time.created
const sequenceA = (a as Chronological)[chronologicalOrder]
const sequenceB = (b as Chronological)[chronologicalOrder]
if (sequenceA !== undefined && sequenceB !== undefined && sequenceA !== sequenceB) return sequenceA - sequenceB
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
// latest finished assistant — i.e. unprocessed work.
export function latest(msgs: WithParts[]) {
let user: User | undefined
let assistant: Assistant | undefined
let finished: Assistant | undefined
let user: WithParts | undefined
let assistant: WithParts | undefined
let finished: WithParts | undefined
let userIndex = -1
let assistantIndex = -1
let finishedIndex = -1
for (const [index, msg] of msgs.entries()) {
const info = msg.info
if (info.role === "user" && (!user || compare(info, user, index, userIndex) > 0)) {
user = info
if (info.role === "user" && (!user || compare(msg, user, index, userIndex) > 0)) {
user = msg
userIndex = index
}
if (info.role === "assistant" && (!assistant || compare(info, assistant, index, assistantIndex) > 0)) {
assistant = info
if (info.role === "assistant" && (!assistant || compare(msg, assistant, index, assistantIndex) > 0)) {
assistant = msg
assistantIndex = index
}
if (info.role === "assistant" && info.finish && (!finished || compare(info, finished, index, finishedIndex) > 0)) {
finished = info
if (info.role === "assistant" && info.finish && (!finished || compare(msg, finished, index, finishedIndex) > 0)) {
finished = msg
finishedIndex = 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"),
)
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(
+9 -4
View File
@@ -1250,14 +1250,18 @@ export const layer = Layer.effect(
let msgs = yield* MessageV2.filterCompactedEffect(sessionID)
const { user: lastUser, assistant: lastAssistant, finished: lastFinished, tasks } = MessageV2.latest(msgs)
const before = (a: MessageV2.Info, b: MessageV2.Info) => MessageV2.compare(a, b) < 0
const latest = MessageV2.latest(msgs)
const { user: lastUser, assistant: lastAssistant, finished: lastFinished, tasks } = latest
if (!lastUser) throw new Error("No user message found in stream. This should never happen.")
const lastAssistantMsg = msgs.findLast(
(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.
// 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
@@ -1269,7 +1273,7 @@ export const layer = Layer.effect(
lastAssistant?.finish &&
!["tool-calls"].includes(lastAssistant.finish) &&
!hasToolCalls &&
before(lastUser, lastAssistant)
userBeforeAssistant
) {
yield* slog.info("exiting loop")
break
@@ -1399,7 +1403,8 @@ export const layer = Layer.effect(
if (step > 1 && lastFinished) {
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) {
if (p.type !== "text" || p.ignored || p.synthetic) continue
if (!p.text.trim()) continue