feat(id): brand PartID through Drizzle and Zod schemas (#16966)

This commit is contained in:
Kit Langton
2026-03-11 19:40:50 -04:00
committed by GitHub
parent d26c6f80e1
commit 090f636354
21 changed files with 102 additions and 97 deletions

View File

@@ -1,8 +1,7 @@
import { BusEvent } from "@/bus/bus-event"
import { Bus } from "@/bus"
import { Session } from "."
import { Identifier } from "../id/id"
import { SessionID, MessageID } from "./schema"
import { SessionID, MessageID, PartID } from "./schema"
import { Instance } from "../project/instance"
import { Provider } from "../provider/provider"
import { MessageV2 } from "./message-v2"
@@ -256,7 +255,7 @@ When constructing the summary, try to stick to this template:
: part
await Session.updatePart({
...replayPart,
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: replayMsg.id,
sessionID: input.sessionID,
})
@@ -276,7 +275,7 @@ When constructing the summary, try to stick to this template:
: "") +
"Continue if you have next steps, or stop and ask for clarification if you are unsure how to proceed."
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: continueMsg.id,
sessionID: input.sessionID,
type: "text",
@@ -317,7 +316,7 @@ When constructing the summary, try to stick to this template:
},
})
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: msg.id,
sessionID: msg.sessionID,
type: "compaction",

View File

@@ -7,7 +7,6 @@ import z from "zod"
import { type ProviderMetadata } from "ai"
import { Config } from "../config/config"
import { Flag } from "../flag/flag"
import { Identifier } from "../id/id"
import { Installation } from "../installation"
import { Database, NotFoundError, eq, and, or, gte, isNull, desc, like, inArray, lt } from "../storage/db"
@@ -25,7 +24,7 @@ import { Snapshot } from "@/snapshot"
import { WorkspaceContext } from "../control-plane/workspace-context"
import { ProjectID } from "../project/schema"
import { WorkspaceID } from "../control-plane/schema"
import { SessionID, MessageID } from "./schema"
import { SessionID, MessageID, PartID } from "./schema"
import type { Provider } from "@/provider/provider"
import { PermissionNext } from "@/permission/next"
@@ -152,7 +151,7 @@ export namespace Session {
revert: z
.object({
messageID: MessageID.zod,
partID: z.string().optional(),
partID: PartID.zod.optional(),
snapshot: z.string().optional(),
diff: z.string().optional(),
})
@@ -269,7 +268,7 @@ export namespace Session {
for (const part of msg.parts) {
await updatePart({
...part,
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: cloned.id,
sessionID: session.id,
})
@@ -731,7 +730,7 @@ export namespace Session {
z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
partID: Identifier.schema("part"),
partID: PartID.zod,
}),
async (input) => {
Database.use((db) => {
@@ -779,7 +778,7 @@ export namespace Session {
z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
partID: z.string(),
partID: PartID.zod,
field: z.string(),
delta: z.string(),
}),

View File

@@ -1,5 +1,5 @@
import { BusEvent } from "@/bus/bus-event"
import { SessionID, MessageID } from "./schema"
import { SessionID, MessageID, PartID } from "./schema"
import z from "zod"
import { NamedError } from "@opencode-ai/util/error"
import { APICallError, convertToModelMessages, LoadAPIKeyError, type ModelMessage, type UIMessage } from "ai"
@@ -78,7 +78,7 @@ export namespace MessageV2 {
export type OutputFormat = z.infer<typeof Format>
const PartBase = z.object({
id: z.string(),
id: PartID.zod,
sessionID: SessionID.zod,
messageID: MessageID.zod,
})
@@ -472,7 +472,7 @@ export namespace MessageV2 {
z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
partID: z.string(),
partID: PartID.zod,
field: z.string(),
delta: z.string(),
}),
@@ -482,7 +482,7 @@ export namespace MessageV2 {
z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
partID: z.string(),
partID: PartID.zod,
}),
),
}

View File

@@ -1,6 +1,5 @@
import { MessageV2 } from "./message-v2"
import { Log } from "@/util/log"
import { Identifier } from "@/id/id"
import { Session } from "."
import { Agent } from "@/agent/agent"
import { Snapshot } from "@/snapshot"
@@ -15,6 +14,7 @@ import { Config } from "@/config/config"
import { SessionCompaction } from "./compaction"
import { PermissionNext } from "@/permission/next"
import { Question } from "@/question"
import { PartID } from "./schema"
import type { SessionID, MessageID } from "./schema"
export namespace SessionProcessor {
@@ -65,7 +65,7 @@ export namespace SessionProcessor {
continue
}
const reasoningPart = {
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.assistantMessage.sessionID,
type: "reasoning" as const,
@@ -111,7 +111,7 @@ export namespace SessionProcessor {
case "tool-input-start":
const part = await Session.updatePart({
id: toolcalls[value.id]?.id ?? Identifier.ascending("part"),
id: toolcalls[value.id]?.id ?? PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.assistantMessage.sessionID,
type: "tool",
@@ -234,7 +234,7 @@ export namespace SessionProcessor {
case "start-step":
snapshot = await Snapshot.track()
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.sessionID,
snapshot,
@@ -252,7 +252,7 @@ export namespace SessionProcessor {
input.assistantMessage.cost += usage.cost
input.assistantMessage.tokens = usage.tokens
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
reason: value.finishReason,
snapshot: await Snapshot.track(),
messageID: input.assistantMessage.id,
@@ -266,7 +266,7 @@ export namespace SessionProcessor {
const patch = await Snapshot.patch(snapshot)
if (patch.files.length) {
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.sessionID,
type: "patch",
@@ -290,7 +290,7 @@ export namespace SessionProcessor {
case "text-start":
currentText = {
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.assistantMessage.sessionID,
type: "text",
@@ -389,7 +389,7 @@ export namespace SessionProcessor {
const patch = await Snapshot.patch(snapshot)
if (patch.files.length) {
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.sessionID,
type: "patch",

View File

@@ -3,8 +3,7 @@ import os from "os"
import fs from "fs/promises"
import z from "zod"
import { Filesystem } from "../util/filesystem"
import { Identifier } from "../id/id"
import { SessionID, MessageID } from "./schema"
import { SessionID, MessageID, PartID } from "./schema"
import { MessageV2 } from "./message-v2"
import { Log } from "../util/log"
import { SessionRevert } from "./revert"
@@ -380,7 +379,7 @@ export namespace SessionPrompt {
},
})) as MessageV2.Assistant
let part = (await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: assistantMessage.id,
sessionID: assistantMessage.sessionID,
type: "tool",
@@ -449,7 +448,7 @@ export namespace SessionPrompt {
})
const attachments = result?.attachments?.map((attachment) => ({
...attachment,
id: Identifier.ascending("part"),
id: PartID.ascending(),
sessionID,
messageID: assistantMessage.id,
}))
@@ -515,7 +514,7 @@ export namespace SessionPrompt {
}
await Session.updateMessage(summaryUserMsg)
await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: summaryUserMsg.id,
sessionID,
type: "text",
@@ -814,7 +813,7 @@ export namespace SessionPrompt {
...result,
attachments: result.attachments?.map((attachment) => ({
...attachment,
id: Identifier.ascending("part"),
id: PartID.ascending(),
sessionID: ctx.sessionID,
messageID: input.processor.message.id,
})),
@@ -917,7 +916,7 @@ export namespace SessionPrompt {
output: truncated.content,
attachments: attachments.map((attachment) => ({
...attachment,
id: Identifier.ascending("part"),
id: PartID.ascending(),
sessionID: ctx.sessionID,
messageID: input.processor.message.id,
})),
@@ -989,7 +988,7 @@ export namespace SessionPrompt {
type Draft<T> = T extends MessageV2.Part ? Omit<T, "id"> & { id?: string } : never
const assign = (part: Draft<MessageV2.Part>): MessageV2.Part => ({
...part,
id: part.id ?? Identifier.ascending("part"),
id: part.id ? PartID.make(part.id) : PartID.ascending(),
})
const parts = await Promise.all(
@@ -1335,7 +1334,7 @@ export namespace SessionPrompt {
if (!Flag.OPENCODE_EXPERIMENTAL_PLAN_MODE) {
if (input.agent.name === "plan") {
userMessage.parts.push({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: userMessage.info.id,
sessionID: userMessage.info.sessionID,
type: "text",
@@ -1346,7 +1345,7 @@ export namespace SessionPrompt {
const wasPlan = input.messages.some((msg) => msg.info.role === "assistant" && msg.info.agent === "plan")
if (wasPlan && input.agent.name === "build") {
userMessage.parts.push({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: userMessage.info.id,
sessionID: userMessage.info.sessionID,
type: "text",
@@ -1366,7 +1365,7 @@ export namespace SessionPrompt {
const exists = await Filesystem.exists(plan)
if (exists) {
const part = await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: userMessage.info.id,
sessionID: userMessage.info.sessionID,
type: "text",
@@ -1385,7 +1384,7 @@ export namespace SessionPrompt {
const exists = await Filesystem.exists(plan)
if (!exists) await fs.mkdir(path.dirname(plan), { recursive: true })
const part = await Session.updatePart({
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: userMessage.info.id,
sessionID: userMessage.info.sessionID,
type: "text",
@@ -1520,7 +1519,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
await Session.updateMessage(userMsg)
const userPart: MessageV2.Part = {
type: "text",
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: userMsg.id,
sessionID: input.sessionID,
text: "The following tool was executed by the user",
@@ -1555,7 +1554,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
await Session.updateMessage(msg)
const part: MessageV2.Part = {
type: "tool",
id: Identifier.ascending("part"),
id: PartID.ascending(),
messageID: msg.id,
sessionID: input.sessionID,
tool: "bash",

View File

@@ -1,6 +1,5 @@
import z from "zod"
import { Identifier } from "../id/id"
import { SessionID, MessageID } from "./schema"
import { SessionID, MessageID, PartID } from "./schema"
import { Snapshot } from "../snapshot"
import { MessageV2 } from "./message-v2"
import { Session } from "."
@@ -18,7 +17,7 @@ export namespace SessionRevert {
export const RevertInput = z.object({
sessionID: SessionID.zod,
messageID: MessageID.zod,
partID: Identifier.schema("part").optional(),
partID: PartID.zod.optional(),
})
export type RevertInput = z.infer<typeof RevertInput>

View File

@@ -27,3 +27,15 @@ export const MessageID = messageIdSchema.pipe(
zod: z.string().startsWith("msg").pipe(z.custom<MessageID>()),
})),
)
const partIdSchema = Schema.String.pipe(Schema.brand("PartId"))
export type PartID = typeof partIdSchema.Type
export const PartID = partIdSchema.pipe(
withStatics((schema: typeof partIdSchema) => ({
make: (id: string) => schema.makeUnsafe(id),
ascending: (id?: string) => schema.makeUnsafe(Identifier.ascending("part", id)),
zod: z.string().startsWith("prt").pipe(z.custom<PartID>()),
})),
)

View File

@@ -4,7 +4,7 @@ import type { MessageV2 } from "./message-v2"
import type { Snapshot } from "../snapshot"
import type { PermissionNext } from "../permission/next"
import type { ProjectID } from "../project/schema"
import type { SessionID, MessageID } from "./schema"
import type { SessionID, MessageID, PartID } from "./schema"
import type { WorkspaceID } from "../control-plane/schema"
import { Timestamps } from "../storage/schema.sql"
@@ -30,7 +30,7 @@ export const SessionTable = sqliteTable(
summary_deletions: integer(),
summary_files: integer(),
summary_diffs: text({ mode: "json" }).$type<Snapshot.FileDiff[]>(),
revert: text({ mode: "json" }).$type<{ messageID: MessageID; partID?: string; snapshot?: string; diff?: string }>(),
revert: text({ mode: "json" }).$type<{ messageID: MessageID; partID?: PartID; snapshot?: string; diff?: string }>(),
permission: text({ mode: "json" }).$type<PermissionNext.Ruleset>(),
...Timestamps,
time_compacting: integer(),
@@ -60,7 +60,7 @@ export const MessageTable = sqliteTable(
export const PartTable = sqliteTable(
"part",
{
id: text().primaryKey(),
id: text().$type<PartID>().primaryKey(),
message_id: text()
.$type<MessageID>()
.notNull()