fix(session): order prompt loop by message creation
This commit is contained in:
parent
3e931152d1
commit
46ea66112e
4 changed files with 118 additions and 33 deletions
|
|
@ -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"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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_-]+$/)
|
||||
})
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue