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 { and } from "drizzle-orm"
import { desc } from "drizzle-orm" import { desc } from "drizzle-orm"
import { eq } from "drizzle-orm" import { eq } from "drizzle-orm"
import { getTableColumns } from "drizzle-orm"
import { inArray } from "drizzle-orm" import { inArray } from "drizzle-orm"
import { lt } from "drizzle-orm" import { lt } from "drizzle-orm"
import { or } from "drizzle-orm" import { or } from "drizzle-orm"
import { sql } from "drizzle-orm"
import { MessageTable, PartTable, SessionTable } from "./session.sql" import { MessageTable, PartTable, SessionTable } from "./session.sql"
import * as ProviderError from "@/provider/error" import * as ProviderError from "@/provider/error"
import { iife } from "@/util/iife" import { iife } from "@/util/iife"
@@ -561,8 +563,8 @@ export type WithParts = {
} }
const Cursor = Schema.Struct({ const Cursor = Schema.Struct({
id: MessageID,
time: Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)), time: Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)),
sequence: Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)),
}) })
type Cursor = typeof Cursor.Type 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, ...row.data,
time: {
...row.data.time,
created: row.time_created,
},
id: row.id, id: row.id,
sessionID: row.session_id, sessionID: row.session_id,
}) as Info }) as Info
@@ -593,9 +604,9 @@ const part = (row: typeof PartTable.$inferSelect) =>
}) as Part }) as Part
const older = (row: Cursor) => 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 ids = rows.map((row) => row.id)
const partByMessage = new Map<string, Part[]>() const partByMessage = new Map<string, Part[]>()
if (ids.length > 0) { if (ids.length > 0) {
@@ -931,10 +942,10 @@ export const page = Effect.fn("MessageV2.page")(function* (input: {
: eq(MessageTable.session_id, input.sessionID) : eq(MessageTable.session_id, input.sessionID)
const rows = Database.use((db) => const rows = Database.use((db) =>
db db
.select() .select({ ...getTableColumns(MessageTable), sequence: messageRowID })
.from(MessageTable) .from(MessageTable)
.where(where) .where(where)
.orderBy(desc(MessageTable.time_created), desc(MessageTable.id)) .orderBy(desc(MessageTable.time_created), desc(messageRowID))
.limit(input.limit + 1) .limit(input.limit + 1)
.all(), .all(),
) )
@@ -957,7 +968,7 @@ export const page = Effect.fn("MessageV2.page")(function* (input: {
return { return {
items, items,
more, 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) completed.add(msg.info.parentID)
} }
result.reverse() result.reverse()
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" &&
@@ -1068,29 +1080,57 @@ export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: Ses
return filterCompacted(stream(sessionID)) 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 // filterCompacted reorders messages for model consumption
// ([compaction-user, summary, ...retained tail..., continue-user]), so array // ([compaction-user, summary, ...retained tail..., continue-user]), so array
// position is not chronological. Derive each binding by max id (MessageID // position is not chronological. Derive each binding by DB-created time; user
// is monotonic via MessageID.ascending) so a pre-compaction overflowing tail // message IDs can be allocated by clients, so lexical ID order is not reliable.
// assistant doesn't get mistaken for the most recent turn. tasks are // Same-millisecond ties use the DB row order captured before compaction reorder.
// compaction/subtask parts attached to user messages newer than the latest // tasks are compaction/subtask parts attached to user messages newer than the
// 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
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 const info = msg.info
if (info.role === "user" && (!user || info.id > user.id)) user = info if (info.role === "user" && (!user || compare(msg, user, index, userIndex) > 0)) {
if (info.role === "assistant" && (!assistant || info.id > assistant.id)) assistant = info user = msg
if (info.role === "assistant" && info.finish && (!finished || info.id > finished.id)) finished = info 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) => const tasks = msgs.flatMap((m, index) =>
finished && m.info.id <= finished.id 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(
+9 -3
View File
@@ -1250,13 +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 { 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
@@ -1268,7 +1273,7 @@ export const layer = Layer.effect(
lastAssistant?.finish && lastAssistant?.finish &&
!["tool-calls"].includes(lastAssistant.finish) && !["tool-calls"].includes(lastAssistant.finish) &&
!hasToolCalls && !hasToolCalls &&
lastUser.id < lastAssistant.id userBeforeAssistant
) { ) {
yield* slog.info("exiting loop") yield* slog.info("exiting loop")
break break
@@ -1398,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" || 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) { 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
@@ -58,12 +58,12 @@ const model: Provider.Model = {
release_date: "2026-01-01", release_date: "2026-01-01",
} }
function userInfo(id: string): MessageV2.User { function userInfo(id: string, created = 0): MessageV2.User {
return { return {
id, id,
sessionID, sessionID,
role: "user", role: "user",
time: { created: 0 }, time: { created },
agent: "user", agent: "user",
model: { providerID, modelID: ModelID.make("test") }, model: { providerID, modelID: ModelID.make("test") },
tools: {}, tools: {},
@@ -76,13 +76,14 @@ function assistantInfo(
parentID: string, parentID: string,
error?: MessageV2.Assistant["error"], error?: MessageV2.Assistant["error"],
meta?: { providerID: string; modelID: string }, meta?: { providerID: string; modelID: string },
created = 0,
): MessageV2.Assistant { ): MessageV2.Assistant {
const infoModel = meta ?? { providerID: model.providerID, modelID: model.api.id } const infoModel = meta ?? { providerID: model.providerID, modelID: model.api.id }
return { return {
id, id,
sessionID, sessionID,
role: "assistant", role: "assistant",
time: { created: 0 }, time: { created },
error, error,
parentID, parentID,
modelID: infoModel.modelID, modelID: infoModel.modelID,
@@ -1557,13 +1558,13 @@ describe("session.message-v2.latest", () => {
const NEW_COMPACTION_USER = MessageID.make("msg_006") const NEW_COMPACTION_USER = MessageID.make("msg_006")
const tailUser: MessageV2.WithParts = { 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[], parts: [{ ...basePart(TAIL_USER, "p1"), type: "text", text: "original prompt" }] as MessageV2.Part[],
} }
const overflowAssistant: MessageV2.WithParts = { const overflowAssistant: MessageV2.WithParts = {
info: { info: {
...assistantInfo(OVERFLOW_ASSISTANT, TAIL_USER), ...assistantInfo(OVERFLOW_ASSISTANT, TAIL_USER, undefined, undefined, 2),
finish: "tool-calls", finish: "tool-calls",
tokens: { input: 280_000, output: 200, reasoning: 0, cache: { read: 0, write: 0 }, total: 280_200 }, tokens: { input: 280_000, output: 200, reasoning: 0, cache: { read: 0, write: 0 }, total: 280_200 },
} as MessageV2.Assistant, } as MessageV2.Assistant,
@@ -1571,7 +1572,7 @@ describe("session.message-v2.latest", () => {
} }
const compactionUser: MessageV2.WithParts = { const compactionUser: MessageV2.WithParts = {
info: userInfo(COMPACTION_USER), info: userInfo(COMPACTION_USER, 3),
parts: [ parts: [
{ {
...basePart(COMPACTION_USER, "p1"), ...basePart(COMPACTION_USER, "p1"),
@@ -1584,7 +1585,7 @@ describe("session.message-v2.latest", () => {
const summaryAssistant: MessageV2.WithParts = { const summaryAssistant: MessageV2.WithParts = {
info: { info: {
...assistantInfo(SUMMARY_ASSISTANT, COMPACTION_USER), ...assistantInfo(SUMMARY_ASSISTANT, COMPACTION_USER, undefined, undefined, 4),
summary: true, summary: true,
finish: "stop", finish: "stop",
tokens: { input: 150_000, output: 1_500, reasoning: 0, cache: { read: 0, write: 0 }, total: 151_500 }, 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 = { const continueUser: MessageV2.WithParts = {
info: userInfo(CONTINUE_USER), info: userInfo(CONTINUE_USER, 5),
parts: [ parts: [
{ {
...basePart(CONTINUE_USER, "p1"), ...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", () => { test("a fresh compaction-user newer than the latest summary surfaces in tasks", () => {
const newCompactionUser: MessageV2.WithParts = { const newCompactionUser: MessageV2.WithParts = {
info: userInfo(NEW_COMPACTION_USER), info: userInfo(NEW_COMPACTION_USER, 6),
parts: [ parts: [
{ {
...basePart(NEW_COMPACTION_USER, "p1"), ...basePart(NEW_COMPACTION_USER, "p1"),
@@ -1653,4 +1654,27 @@ describe("session.message-v2.latest", () => {
expect(state.tasks).toHaveLength(1) expect(state.tasks).toHaveLength(1)
expect(state.tasks[0]).toMatchObject({ type: "compaction", auto: true }) 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", () => it.instance("returns empty items for session with no messages", () =>
withSession(({ sessionID }) => withSession(({ sessionID }) =>
Effect.gen(function* () { Effect.gen(function* () {
@@ -972,22 +1001,22 @@ describe("MessageV2.filterCompacted", () => {
describe("MessageV2.cursor", () => { describe("MessageV2.cursor", () => {
test("encode/decode roundtrip", () => { test("encode/decode roundtrip", () => {
const input = { id: MessageID.ascending(), time: 1234567890 } const input = { time: 1234567890, sequence: 1 }
const encoded = MessageV2.cursor.encode(input) const encoded = MessageV2.cursor.encode(input)
const decoded = MessageV2.cursor.decode(encoded) const decoded = MessageV2.cursor.decode(encoded)
expect(decoded.id).toBe(input.id)
expect(decoded.time).toBe(input.time) expect(decoded.time).toBe(input.time)
expect(decoded.sequence).toBe(input.sequence)
}) })
test("encode/decode with fractional time", () => { 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 encoded = MessageV2.cursor.encode(input)
const decoded = MessageV2.cursor.decode(encoded) const decoded = MessageV2.cursor.decode(encoded)
expect(decoded.time).toBe(1234567890.5) expect(decoded.time).toBe(1234567890.5)
}) })
test("encoded cursor is base64url", () => { 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_-]+$/) expect(encoded).toMatch(/^[A-Za-z0-9_-]+$/)
}) })
}) })