Compare commits

...

2 Commits

Author SHA1 Message Date
Aiden Cline 75144fedf6 refactor(session): make message ordering explicit 2026-05-23 17:53:32 -05:00
Aiden Cline 46ea66112e fix(session): order prompt loop by message creation 2026-05-23 17:33:22 -05:00
4 changed files with 137 additions and 38 deletions
+62 -22
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,18 @@ export const cursor = {
},
}
const info = (row: typeof MessageTable.$inferSelect) =>
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) =>
({
...row.data,
time: {
...row.data.time,
created: row.time_created,
},
id: row.id,
sessionID: row.session_id,
}) as Info
@@ -593,9 +604,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 +942,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 +968,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 +1046,7 @@ export function filterCompacted(msgs: Iterable<WithParts>) {
completed.add(msg.info.parentID)
}
result.reverse()
result.forEach((msg, index) => ((msg as Chronological)[chronologicalOrder] = index))
const compactionIndex = result.findLastIndex(
(msg) =>
msg.info.role === "user" &&
@@ -1068,29 +1080,57 @@ export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: Ses
return filterCompacted(stream(sessionID))
})
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
}
// 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 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 || 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(msg, user, index, userIndex) > 0)) {
user = msg
userIndex = index
}
if (info.role === "assistant" && (!assistant || compare(msg, assistant, index, assistantIndex) > 0)) {
assistant = msg
assistantIndex = index
}
if (info.role === "assistant" && info.finish && (!finished || compare(msg, finished, index, finishedIndex) > 0)) {
finished = msg
finishedIndex = index
}
}
const tasks = msgs.flatMap((m) =>
finished && m.info.id <= finished.id
const tasks = msgs.flatMap((m, index) =>
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 -3
View File
@@ -1250,13 +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 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
@@ -1268,7 +1273,7 @@ export const layer = Layer.effect(
lastAssistant?.finish &&
!["tool-calls"].includes(lastAssistant.finish) &&
!hasToolCalls &&
lastUser.id < lastAssistant.id
userBeforeAssistant
) {
yield* slog.info("exiting loop")
break
@@ -1398,7 +1403,8 @@ 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
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
@@ -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)
})
})
@@ -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_-]+$/)
})
})