refactor: add Effect-returning versions of MessageV2 functions (#20374)

This commit is contained in:
Kit Langton
2026-04-01 19:48:36 -04:00
committed by GitHub
parent 4214ae205d
commit 26fb6b8788
8 changed files with 878 additions and 114 deletions

View File

@@ -218,7 +218,7 @@ When constructing the summary, try to stick to this template:
const prompt = compacting.prompt ?? [defaultPrompt, ...compacting.context].join("\n\n")
const msgs = structuredClone(messages)
yield* plugin.trigger("experimental.chat.messages.transform", {}, { messages: msgs })
const modelMessages = yield* Effect.promise(() => MessageV2.toModelMessages(msgs, model, { stripMedia: true }))
const modelMessages = yield* MessageV2.toModelMessagesEffect(msgs, model, { stripMedia: true })
const ctx = yield* InstanceState.context
const msg: MessageV2.Assistant = {
id: MessageID.ascending(),

View File

@@ -593,15 +593,10 @@ export namespace Session {
})
const messages = Effect.fn("Session.messages")(function* (input: { sessionID: SessionID; limit?: number }) {
return yield* Effect.promise(async () => {
const result = [] as MessageV2.WithParts[]
for await (const msg of MessageV2.stream(input.sessionID)) {
if (input.limit && result.length >= input.limit) break
result.push(msg)
}
result.reverse()
return result
})
if (input.limit) {
return MessageV2.page({ sessionID: input.sessionID, limit: input.limit }).items
}
return Array.from(MessageV2.stream(input.sessionID)).reverse()
})
const removeMessage = Effect.fn("Session.removeMessage")(function* (input: {

View File

@@ -5,7 +5,6 @@ import { NamedError } from "@opencode-ai/util/error"
import { APICallError, convertToModelMessages, LoadAPIKeyError, type ModelMessage, type UIMessage } from "ai"
import { LSP } from "../lsp"
import { Snapshot } from "@/snapshot"
import { fn } from "@/util/fn"
import { SyncEvent } from "../sync"
import { Database, NotFoundError, and, desc, eq, inArray, lt, or } from "@/storage/db"
import { MessageTable, PartTable, SessionTable } from "./session.sql"
@@ -15,6 +14,7 @@ import { errorMessage } from "@/util/error"
import type { SystemError } from "bun"
import type { Provider } from "@/provider/provider"
import { ModelID, ProviderID } from "@/provider/schema"
import { Effect } from "effect"
/** Error shape thrown by Bun's fetch() when gzip/br decompression fails mid-stream */
interface FetchDecompressionError extends Error {
@@ -547,7 +547,7 @@ export namespace MessageV2 {
and(eq(MessageTable.time_created, row.time), lt(MessageTable.id, row.id)),
)
async function hydrate(rows: (typeof MessageTable.$inferSelect)[]) {
function hydrate(rows: (typeof MessageTable.$inferSelect)[]) {
const ids = rows.map((row) => row.id)
const partByMessage = new Map<string, MessageV2.Part[]>()
if (ids.length > 0) {
@@ -573,11 +573,11 @@ export namespace MessageV2 {
}))
}
export async function toModelMessages(
export const toModelMessagesEffect = Effect.fnUntraced(function* (
input: WithParts[],
model: Provider.Model,
options?: { stripMedia?: boolean },
): Promise<ModelMessage[]> {
) {
const result: UIMessage[] = []
const toolNames = new Set<string>()
// Track media from tool results that need to be injected as user messages
@@ -800,64 +800,67 @@ export namespace MessageV2 {
const tools = Object.fromEntries(Array.from(toolNames).map((toolName) => [toolName, { toModelOutput }]))
return await convertToModelMessages(
result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
{
//@ts-expect-error (convertToModelMessages expects a ToolSet but only actually needs tools[name]?.toModelOutput)
tools,
},
return yield* Effect.promise(() =>
convertToModelMessages(
result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
{
//@ts-expect-error (convertToModelMessages expects a ToolSet but only actually needs tools[name]?.toModelOutput)
tools,
},
),
)
})
export function toModelMessages(
input: WithParts[],
model: Provider.Model,
options?: { stripMedia?: boolean },
): Promise<ModelMessage[]> {
return Effect.runPromise(toModelMessagesEffect(input, model, options))
}
export const page = fn(
z.object({
sessionID: SessionID.zod,
limit: z.number().int().positive(),
before: z.string().optional(),
}),
async (input) => {
const before = input.before ? cursor.decode(input.before) : undefined
const where = before
? and(eq(MessageTable.session_id, input.sessionID), older(before))
: eq(MessageTable.session_id, input.sessionID)
const rows = Database.use((db) =>
db
.select()
.from(MessageTable)
.where(where)
.orderBy(desc(MessageTable.time_created), desc(MessageTable.id))
.limit(input.limit + 1)
.all(),
export function page(input: { sessionID: SessionID; limit: number; before?: string }) {
const before = input.before ? cursor.decode(input.before) : undefined
const where = before
? and(eq(MessageTable.session_id, input.sessionID), older(before))
: eq(MessageTable.session_id, input.sessionID)
const rows = Database.use((db) =>
db
.select()
.from(MessageTable)
.where(where)
.orderBy(desc(MessageTable.time_created), desc(MessageTable.id))
.limit(input.limit + 1)
.all(),
)
if (rows.length === 0) {
const row = Database.use((db) =>
db.select({ id: SessionTable.id }).from(SessionTable).where(eq(SessionTable.id, input.sessionID)).get(),
)
if (rows.length === 0) {
const row = Database.use((db) =>
db.select({ id: SessionTable.id }).from(SessionTable).where(eq(SessionTable.id, input.sessionID)).get(),
)
if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
return {
items: [] as MessageV2.WithParts[],
more: false,
}
}
const more = rows.length > input.limit
const page = more ? rows.slice(0, input.limit) : rows
const items = await hydrate(page)
items.reverse()
const tail = page.at(-1)
if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
return {
items,
more,
cursor: more && tail ? cursor.encode({ id: tail.id, time: tail.time_created }) : undefined,
items: [] as MessageV2.WithParts[],
more: false,
}
},
)
}
export const stream = fn(SessionID.zod, async function* (sessionID) {
const more = rows.length > input.limit
const slice = more ? rows.slice(0, input.limit) : rows
const items = hydrate(slice)
items.reverse()
const tail = slice.at(-1)
return {
items,
more,
cursor: more && tail ? cursor.encode({ id: tail.id, time: tail.time_created }) : undefined,
}
}
export function* stream(sessionID: SessionID) {
const size = 50
let before: string | undefined
while (true) {
const next = await page({ sessionID, limit: size, before })
const next = page({ sessionID, limit: size, before })
if (next.items.length === 0) break
for (let i = next.items.length - 1; i >= 0; i--) {
yield next.items[i]
@@ -865,9 +868,9 @@ export namespace MessageV2 {
if (!next.more || !next.cursor) break
before = next.cursor
}
})
}
export const parts = fn(MessageID.zod, async (message_id) => {
export function parts(message_id: MessageID) {
const rows = Database.use((db) =>
db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).orderBy(PartTable.id).all(),
)
@@ -880,33 +883,28 @@ export namespace MessageV2 {
messageID: row.message_id,
}) as MessageV2.Part,
)
})
}
export const get = fn(
z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
}),
async (input): Promise<WithParts> => {
const row = Database.use((db) =>
db
.select()
.from(MessageTable)
.where(and(eq(MessageTable.id, input.messageID), eq(MessageTable.session_id, input.sessionID)))
.get(),
)
if (!row) throw new NotFoundError({ message: `Message not found: ${input.messageID}` })
return {
info: info(row),
parts: await parts(input.messageID),
}
},
)
export function get(input: { sessionID: SessionID; messageID: MessageID }): WithParts {
const row = Database.use((db) =>
db
.select()
.from(MessageTable)
.where(and(eq(MessageTable.id, input.messageID), eq(MessageTable.session_id, input.sessionID)))
.get(),
)
if (!row) throw new NotFoundError({ message: `Message not found: ${input.messageID}` })
return {
info: info(row),
parts: parts(input.messageID),
}
}
export async function filterCompacted(stream: AsyncIterable<MessageV2.WithParts>) {
export function filterCompacted(msgs: Iterable<MessageV2.WithParts>) {
const result = [] as MessageV2.WithParts[]
const completed = new Set<string>()
for await (const msg of stream) {
for (const msg of msgs) {
result.push(msg)
if (
msg.info.role === "user" &&
@@ -921,6 +919,10 @@ export namespace MessageV2 {
return result
}
export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: SessionID) {
return filterCompacted(stream(sessionID))
})
export function fromError(
e: unknown,
ctx: { providerID: ProviderID; aborted?: boolean },

View File

@@ -184,7 +184,7 @@ export namespace SessionProcessor {
metadata: value.providerMetadata,
} satisfies MessageV2.ToolPart)
const parts = yield* Effect.promise(() => MessageV2.parts(ctx.assistantMessage.id))
const parts = MessageV2.parts(ctx.assistantMessage.id)
const recentParts = parts.slice(-DOOM_LOOP_THRESHOLD)
if (
@@ -396,7 +396,7 @@ export namespace SessionProcessor {
}
ctx.reasoningMap = {}
const parts = yield* Effect.promise(() => MessageV2.parts(ctx.assistantMessage.id))
const parts = MessageV2.parts(ctx.assistantMessage.id)
for (const part of parts) {
if (part.type !== "tool" || part.state.status === "completed" || part.state.status === "error") continue
yield* session.updatePart({

View File

@@ -216,7 +216,7 @@ export namespace SessionPrompt {
(yield* provider.getModel(input.providerID, input.modelID)))
const msgs = onlySubtasks
? [{ role: "user" as const, content: subtasks.map((p) => p.prompt).join("\n") }]
: yield* Effect.promise(() => MessageV2.toModelMessages(context, mdl))
: yield* MessageV2.toModelMessagesEffect(context, mdl)
const text = yield* Effect.promise(async (signal) => {
const result = await LLM.stream({
agent: ag,
@@ -1342,7 +1342,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
yield* status.set(sessionID, { type: "busy" })
log.info("loop", { step, sessionID })
let msgs = yield* Effect.promise(() => MessageV2.filterCompacted(MessageV2.stream(sessionID)))
let msgs = yield* MessageV2.filterCompactedEffect(sessionID)
let lastUser: MessageV2.User | undefined
let lastAssistant: MessageV2.Assistant | undefined