fix(session): order prompt loop by message creation

This commit is contained in:
Aiden Cline 2026-05-23 17:33:22 -05:00
commit 46ea66112e
4 changed files with 118 additions and 33 deletions

View file

@ -10,9 +10,11 @@ import { NotFoundError } from "@/storage/storage"
import { and } from "drizzle-orm"
import { desc } from "drizzle-orm"
import { eq } from "drizzle-orm"
import { getTableColumns } from "drizzle-orm"
import { inArray } from "drizzle-orm"
import { lt } from "drizzle-orm"
import { or } from "drizzle-orm"
import { sql } from "drizzle-orm"
import { MessageTable, PartTable, SessionTable } from "./session.sql"
import * as ProviderError from "@/provider/error"
import { iife } from "@/util/iife"
@ -561,8 +563,8 @@ export type WithParts = {
}
const Cursor = Schema.Struct({
id: MessageID,
time: Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)),
sequence: Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)),
})
type Cursor = typeof Cursor.Type
@ -577,9 +579,17 @@ export const cursor = {
},
}
const info = (row: typeof MessageTable.$inferSelect) =>
const messageOrder = new WeakMap<Info, number>()
const messageRowID = sql<number>`rowid`
type MessageRow = typeof MessageTable.$inferSelect & { sequence?: number }
const info = (row: MessageRow) =>
({
...row.data,
time: {
...row.data.time,
created: row.time_created,
},
id: row.id,
sessionID: row.session_id,
}) as Info
@ -593,9 +603,9 @@ const part = (row: typeof PartTable.$inferSelect) =>
}) as Part
const older = (row: Cursor) =>
or(lt(MessageTable.time_created, row.time), and(eq(MessageTable.time_created, row.time), lt(MessageTable.id, row.id)))
or(lt(MessageTable.time_created, row.time), and(eq(MessageTable.time_created, row.time), lt(messageRowID, row.sequence)))
function hydrate(rows: (typeof MessageTable.$inferSelect)[]) {
function hydrate(rows: MessageRow[]) {
const ids = rows.map((row) => row.id)
const partByMessage = new Map<string, Part[]>()
if (ids.length > 0) {
@ -931,10 +941,10 @@ export const page = Effect.fn("MessageV2.page")(function* (input: {
: eq(MessageTable.session_id, input.sessionID)
const rows = Database.use((db) =>
db
.select()
.select({ ...getTableColumns(MessageTable), sequence: messageRowID })
.from(MessageTable)
.where(where)
.orderBy(desc(MessageTable.time_created), desc(MessageTable.id))
.orderBy(desc(MessageTable.time_created), desc(messageRowID))
.limit(input.limit + 1)
.all(),
)
@ -957,7 +967,7 @@ export const page = Effect.fn("MessageV2.page")(function* (input: {
return {
items,
more,
cursor: more && tail ? cursor.encode({ id: tail.id, time: tail.time_created }) : undefined,
cursor: more && tail ? cursor.encode({ time: tail.time_created, sequence: tail.sequence ?? 0 }) : undefined,
}
})
@ -1035,6 +1045,7 @@ export function filterCompacted(msgs: Iterable<WithParts>) {
completed.add(msg.info.parentID)
}
result.reverse()
result.forEach((msg, index) => messageOrder.set(msg.info, index))
const compactionIndex = result.findLastIndex(
(msg) =>
msg.info.role === "user" &&
@ -1068,25 +1079,45 @@ 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)
if (sequenceA !== undefined && sequenceB !== undefined && sequenceA !== sequenceB) return sequenceA - sequenceB
return indexA - indexB
}
// filterCompacted reorders messages for model consumption
// ([compaction-user, summary, ...retained tail..., continue-user]), so array
// position is not chronological. Derive each binding by max id (MessageID
// is monotonic via MessageID.ascending) so a pre-compaction overflowing tail
// assistant doesn't get mistaken for the most recent turn. tasks are
// compaction/subtask parts attached to user messages newer than the latest
// finished assistant — i.e. unprocessed work.
// position is not chronological. Derive each binding by DB-created time; user
// message IDs can be allocated by clients, so lexical ID order is not reliable.
// Same-millisecond ties use the DB row order captured before compaction reorder.
// 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
for (const msg of msgs) {
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 || info.id > user.id)) user = info
if (info.role === "assistant" && (!assistant || info.id > assistant.id)) assistant = info
if (info.role === "assistant" && info.finish && (!finished || info.id > finished.id)) finished = info
if (info.role === "user" && (!user || compare(info, user, index, userIndex) > 0)) {
user = info
userIndex = index
}
if (info.role === "assistant" && (!assistant || compare(info, assistant, index, assistantIndex) > 0)) {
assistant = info
assistantIndex = index
}
if (info.role === "assistant" && info.finish && (!finished || compare(info, finished, index, finishedIndex) > 0)) {
finished = info
finishedIndex = index
}
}
const tasks = msgs.flatMap((m) =>
finished && m.info.id <= finished.id
const tasks = msgs.flatMap((m, index) =>
finished && compare(m.info, finished, index, finishedIndex) <= 0
? []
: m.parts.filter((p): p is CompactionPart | SubtaskPart => p.type === "compaction" || p.type === "subtask"),
)

View file

@ -1251,6 +1251,7 @@ 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
if (!lastUser) throw new Error("No user message found in stream. This should never happen.")
@ -1268,7 +1269,7 @@ export const layer = Layer.effect(
lastAssistant?.finish &&
!["tool-calls"].includes(lastAssistant.finish) &&
!hasToolCalls &&
lastUser.id < lastAssistant.id
before(lastUser, lastAssistant)
) {
yield* slog.info("exiting loop")
break
@ -1398,7 +1399,7 @@ export const layer = Layer.effect(
if (step > 1 && lastFinished) {
for (const m of msgs) {
if (m.info.role !== "user" || m.info.id <= lastFinished.id) continue
if (m.info.role !== "user" || !before(lastFinished, m.info)) continue
for (const p of m.parts) {
if (p.type !== "text" || p.ignored || p.synthetic) continue
if (!p.text.trim()) continue

View file

@ -58,12 +58,12 @@ const model: Provider.Model = {
release_date: "2026-01-01",
}
function userInfo(id: string): MessageV2.User {
function userInfo(id: string, created = 0): MessageV2.User {
return {
id,
sessionID,
role: "user",
time: { created: 0 },
time: { created },
agent: "user",
model: { providerID, modelID: ModelID.make("test") },
tools: {},
@ -76,13 +76,14 @@ function assistantInfo(
parentID: string,
error?: MessageV2.Assistant["error"],
meta?: { providerID: string; modelID: string },
created = 0,
): MessageV2.Assistant {
const infoModel = meta ?? { providerID: model.providerID, modelID: model.api.id }
return {
id,
sessionID,
role: "assistant",
time: { created: 0 },
time: { created },
error,
parentID,
modelID: infoModel.modelID,
@ -1557,13 +1558,13 @@ describe("session.message-v2.latest", () => {
const NEW_COMPACTION_USER = MessageID.make("msg_006")
const tailUser: MessageV2.WithParts = {
info: userInfo(TAIL_USER),
info: userInfo(TAIL_USER, 1),
parts: [{ ...basePart(TAIL_USER, "p1"), type: "text", text: "original prompt" }] as MessageV2.Part[],
}
const overflowAssistant: MessageV2.WithParts = {
info: {
...assistantInfo(OVERFLOW_ASSISTANT, TAIL_USER),
...assistantInfo(OVERFLOW_ASSISTANT, TAIL_USER, undefined, undefined, 2),
finish: "tool-calls",
tokens: { input: 280_000, output: 200, reasoning: 0, cache: { read: 0, write: 0 }, total: 280_200 },
} as MessageV2.Assistant,
@ -1571,7 +1572,7 @@ describe("session.message-v2.latest", () => {
}
const compactionUser: MessageV2.WithParts = {
info: userInfo(COMPACTION_USER),
info: userInfo(COMPACTION_USER, 3),
parts: [
{
...basePart(COMPACTION_USER, "p1"),
@ -1584,7 +1585,7 @@ describe("session.message-v2.latest", () => {
const summaryAssistant: MessageV2.WithParts = {
info: {
...assistantInfo(SUMMARY_ASSISTANT, COMPACTION_USER),
...assistantInfo(SUMMARY_ASSISTANT, COMPACTION_USER, undefined, undefined, 4),
summary: true,
finish: "stop",
tokens: { input: 150_000, output: 1_500, reasoning: 0, cache: { read: 0, write: 0 }, total: 151_500 },
@ -1593,7 +1594,7 @@ describe("session.message-v2.latest", () => {
}
const continueUser: MessageV2.WithParts = {
info: userInfo(CONTINUE_USER),
info: userInfo(CONTINUE_USER, 5),
parts: [
{
...basePart(CONTINUE_USER, "p1"),
@ -1629,7 +1630,7 @@ describe("session.message-v2.latest", () => {
test("a fresh compaction-user newer than the latest summary surfaces in tasks", () => {
const newCompactionUser: MessageV2.WithParts = {
info: userInfo(NEW_COMPACTION_USER),
info: userInfo(NEW_COMPACTION_USER, 6),
parts: [
{
...basePart(NEW_COMPACTION_USER, "p1"),
@ -1653,4 +1654,27 @@ describe("session.message-v2.latest", () => {
expect(state.tasks).toHaveLength(1)
expect(state.tasks[0]).toMatchObject({ type: "compaction", auto: true })
})
test("latest uses created time when message ids are not chronological", () => {
const newerAssistant = MessageID.make("msg_001")
const olderAssistant = MessageID.make("msg_999")
const state = MessageV2.latest([
{
info: {
...assistantInfo(olderAssistant, TAIL_USER, undefined, undefined, 1),
finish: "stop",
},
parts: [],
},
{
info: {
...assistantInfo(newerAssistant, TAIL_USER, undefined, undefined, 2),
finish: "stop",
},
parts: [],
},
])
expect(state.finished?.id).toBe(newerAssistant)
})
})

View file

@ -173,6 +173,35 @@ describe("MessageV2.page", () => {
),
)
it.instance("uses db order when same-timestamp ids are not chronological", () =>
withSession(({ sessionID }) =>
Effect.gen(function* () {
const session = yield* SessionNs.Service
const older = MessageID.make("msg_999")
const newer = MessageID.make("msg_001")
for (const id of [older, newer]) {
yield* session.updateMessage({
id,
sessionID,
role: "user",
time: { created: 1 },
agent: "test",
model: { providerID: "test", modelID: "test" },
tools: {},
mode: "",
} as unknown as MessageV2.Info)
}
const first = yield* MessageV2.page({ sessionID, limit: 1 })
expect(first.items.map((item) => item.info.id)).toEqual([newer])
expect(first.cursor).toBeTruthy()
const second = yield* MessageV2.page({ sessionID, limit: 1, before: first.cursor! })
expect(second.items.map((item) => item.info.id)).toEqual([older])
}),
),
)
it.instance("returns empty items for session with no messages", () =>
withSession(({ sessionID }) =>
Effect.gen(function* () {
@ -972,22 +1001,22 @@ describe("MessageV2.filterCompacted", () => {
describe("MessageV2.cursor", () => {
test("encode/decode roundtrip", () => {
const input = { id: MessageID.ascending(), time: 1234567890 }
const input = { time: 1234567890, sequence: 1 }
const encoded = MessageV2.cursor.encode(input)
const decoded = MessageV2.cursor.decode(encoded)
expect(decoded.id).toBe(input.id)
expect(decoded.time).toBe(input.time)
expect(decoded.sequence).toBe(input.sequence)
})
test("encode/decode with fractional time", () => {
const input = { id: MessageID.ascending(), time: 1234567890.5 }
const input = { time: 1234567890.5, sequence: 1 }
const encoded = MessageV2.cursor.encode(input)
const decoded = MessageV2.cursor.decode(encoded)
expect(decoded.time).toBe(1234567890.5)
})
test("encoded cursor is base64url", () => {
const encoded = MessageV2.cursor.encode({ id: MessageID.ascending(), time: 0 })
const encoded = MessageV2.cursor.encode({ time: 0, sequence: 1 })
expect(encoded).toMatch(/^[A-Za-z0-9_-]+$/)
})
})