Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 75144fedf6 | |||
| 46ea66112e |
@@ -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(
|
||||
|
||||
@@ -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_-]+$/)
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user