feat(id): brand WorkspaceID through Drizzle and Zod schemas (#16964)

packages/opencode/src/session/compaction.ts

@@ -2,7 +2,7 @@ import { BusEvent } from "@/bus/bus-event"
 import { Bus } from "@/bus"
 import { Session } from "."
 import { Identifier } from "../id/id"
-import { SessionID } from "./schema"
+import { SessionID, MessageID } from "./schema"
 import { Instance } from "../project/instance"
 import { Provider } from "../provider/provider"
 import { MessageV2 } from "./message-v2"
@@ -100,7 +100,7 @@ export namespace SessionCompaction {
   }
 
   export async function process(input: {
-    parentID: string
+    parentID: MessageID
     messages: MessageV2.WithParts[]
     sessionID: SessionID
     abort: AbortSignal
@@ -134,7 +134,7 @@ export namespace SessionCompaction {
       ? await Provider.getModel(agent.model.providerID, agent.model.modelID)
       : await Provider.getModel(userMessage.model.providerID, userMessage.model.modelID)
     const msg = (await Session.updateMessage({
-      id: Identifier.ascending("message"),
+      id: MessageID.ascending(),
       role: "assistant",
       parentID: input.parentID,
       sessionID: input.sessionID,
@@ -237,7 +237,7 @@ When constructing the summary, try to stick to this template:
       if (replay) {
         const original = replay.info as MessageV2.User
         const replayMsg = await Session.updateMessage({
-          id: Identifier.ascending("message"),
+          id: MessageID.ascending(),
           role: "user",
           sessionID: input.sessionID,
           time: { created: Date.now() },
@@ -263,7 +263,7 @@ When constructing the summary, try to stick to this template:
         }
       } else {
         const continueMsg = await Session.updateMessage({
-          id: Identifier.ascending("message"),
+          id: MessageID.ascending(),
           role: "user",
           sessionID: input.sessionID,
           time: { created: Date.now() },
@@ -307,7 +307,7 @@ When constructing the summary, try to stick to this template:
     }),
     async (input) => {
       const msg = await Session.updateMessage({
-        id: Identifier.ascending("message"),
+        id: MessageID.ascending(),
         role: "user",
         model: input.model,
         sessionID: input.sessionID,

packages/opencode/src/session/index.ts

@@ -24,7 +24,8 @@ import { Command } from "../command"
 import { Snapshot } from "@/snapshot"
 import { WorkspaceContext } from "../control-plane/workspace-context"
 import { ProjectID } from "../project/schema"
-import { SessionID } from "./schema"
+import { WorkspaceID } from "../control-plane/schema"
+import { SessionID, MessageID } from "./schema"
 
 import type { Provider } from "@/provider/provider"
 import { PermissionNext } from "@/permission/next"
@@ -123,7 +124,7 @@ export namespace Session {
       id: SessionID.zod,
       slug: z.string(),
       projectID: ProjectID.zod,
-      workspaceID: z.string().optional(),
+      workspaceID: WorkspaceID.zod.optional(),
       directory: z.string(),
       parentID: SessionID.zod.optional(),
       summary: z
@@ -150,7 +151,7 @@ export namespace Session {
       permission: PermissionNext.Ruleset.optional(),
       revert: z
         .object({
-          messageID: z.string(),
+          messageID: MessageID.zod,
           partID: z.string().optional(),
           snapshot: z.string().optional(),
           diff: z.string().optional(),
@@ -221,7 +222,7 @@ export namespace Session {
         parentID: SessionID.zod.optional(),
         title: z.string().optional(),
         permission: Info.shape.permission,
-        workspaceID: Identifier.schema("workspace").optional(),
+        workspaceID: WorkspaceID.zod.optional(),
       })
       .optional(),
     async (input) => {
@@ -238,7 +239,7 @@ export namespace Session {
   export const fork = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: Identifier.schema("message").optional(),
+      messageID: MessageID.zod.optional(),
     }),
     async (input) => {
       const original = await get(input.sessionID)
@@ -250,11 +251,11 @@ export namespace Session {
         title,
       })
       const msgs = await messages({ sessionID: input.sessionID })
-      const idMap = new Map<string, string>()
+      const idMap = new Map<string, MessageID>()
 
       for (const msg of msgs) {
         if (input.messageID && msg.info.id >= input.messageID) break
-        const newID = Identifier.ascending("message")
+        const newID = MessageID.ascending()
         idMap.set(msg.info.id, newID)
 
         const parentID = msg.info.role === "assistant" && msg.info.parentID ? idMap.get(msg.info.parentID) : undefined
@@ -297,7 +298,7 @@ export namespace Session {
     id?: SessionID
     title?: string
     parentID?: SessionID
-    workspaceID?: string
+    workspaceID?: WorkspaceID
     directory: string
     permission?: PermissionNext.Ruleset
   }) {
@@ -538,7 +539,7 @@ export namespace Session {
 
   export function* list(input?: {
     directory?: string
-    workspaceID?: string
+    workspaceID?: WorkspaceID
     roots?: boolean
     start?: number
     search?: string
@@ -707,7 +708,7 @@ export namespace Session {
   export const removeMessage = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: Identifier.schema("message"),
+      messageID: MessageID.zod,
     }),
     async (input) => {
       // CASCADE delete handles parts automatically
@@ -729,7 +730,7 @@ export namespace Session {
   export const removePart = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: Identifier.schema("message"),
+      messageID: MessageID.zod,
       partID: Identifier.schema("part"),
     }),
     async (input) => {
@@ -777,7 +778,7 @@ export namespace Session {
   export const updatePartDelta = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: z.string(),
+      messageID: MessageID.zod,
       partID: z.string(),
       field: z.string(),
       delta: z.string(),
@@ -877,7 +878,7 @@ export namespace Session {
       sessionID: SessionID.zod,
       modelID: z.string(),
       providerID: z.string(),
-      messageID: Identifier.schema("message"),
+      messageID: MessageID.zod,
     }),
     async (input) => {
       await SessionPrompt.command({

packages/opencode/src/session/message-v2.ts

@@ -1,9 +1,8 @@
 import { BusEvent } from "@/bus/bus-event"
-import { SessionID } from "./schema"
+import { SessionID, MessageID } from "./schema"
 import z from "zod"
 import { NamedError } from "@opencode-ai/util/error"
 import { APICallError, convertToModelMessages, LoadAPIKeyError, type ModelMessage, type UIMessage } from "ai"
-import { Identifier } from "../id/id"
 import { LSP } from "../lsp"
 import { Snapshot } from "@/snapshot"
 import { fn } from "@/util/fn"
@@ -81,7 +80,7 @@ export namespace MessageV2 {
   const PartBase = z.object({
     id: z.string(),
     sessionID: SessionID.zod,
-    messageID: z.string(),
+    messageID: MessageID.zod,
   })
 
   export const SnapshotPart = PartBase.extend({
@@ -344,7 +343,7 @@ export namespace MessageV2 {
   export type ToolPart = z.infer<typeof ToolPart>
 
   const Base = z.object({
-    id: z.string(),
+    id: MessageID.zod,
     sessionID: SessionID.zod,
   })
 
@@ -411,7 +410,7 @@ export namespace MessageV2 {
         APIError.Schema,
       ])
       .optional(),
-    parentID: z.string(),
+    parentID: MessageID.zod,
     modelID: z.string(),
     providerID: z.string(),
     /**
@@ -459,7 +458,7 @@ export namespace MessageV2 {
       "message.removed",
       z.object({
         sessionID: SessionID.zod,
-        messageID: z.string(),
+        messageID: MessageID.zod,
       }),
     ),
     PartUpdated: BusEvent.define(
@@ -472,7 +471,7 @@ export namespace MessageV2 {
       "message.part.delta",
       z.object({
         sessionID: SessionID.zod,
-        messageID: z.string(),
+        messageID: MessageID.zod,
         partID: z.string(),
         field: z.string(),
         delta: z.string(),
@@ -482,7 +481,7 @@ export namespace MessageV2 {
       "message.part.removed",
       z.object({
         sessionID: SessionID.zod,
-        messageID: z.string(),
+        messageID: MessageID.zod,
         partID: z.string(),
       }),
     ),
@@ -699,7 +698,7 @@ export namespace MessageV2 {
           // media (images, PDFs) in tool results
           if (media.length > 0) {
             result.push({
-              id: Identifier.ascending("message"),
+              id: MessageID.ascending(),
               role: "user",
               parts: [
                 {
@@ -782,7 +781,7 @@ export namespace MessageV2 {
     }
   })
 
-  export const parts = fn(Identifier.schema("message"), async (message_id) => {
+  export const parts = fn(MessageID.zod, async (message_id) => {
     const rows = Database.use((db) =>
       db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).orderBy(PartTable.id).all(),
     )
@@ -794,7 +793,7 @@ export namespace MessageV2 {
   export const get = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: Identifier.schema("message"),
+      messageID: MessageID.zod,
     }),
     async (input): Promise<WithParts> => {
       const row = Database.use((db) => db.select().from(MessageTable).where(eq(MessageTable.id, input.messageID)).get())

packages/opencode/src/session/processor.ts

@@ -15,7 +15,7 @@ import { Config } from "@/config/config"
 import { SessionCompaction } from "./compaction"
 import { PermissionNext } from "@/permission/next"
 import { Question } from "@/question"
-import type { SessionID } from "./schema"
+import type { SessionID, MessageID } from "./schema"
 
 export namespace SessionProcessor {
   const DOOM_LOOP_THRESHOLD = 3

packages/opencode/src/session/prompt.ts

@@ -4,7 +4,7 @@ import fs from "fs/promises"
 import z from "zod"
 import { Filesystem } from "../util/filesystem"
 import { Identifier } from "../id/id"
-import { SessionID } from "./schema"
+import { SessionID, MessageID } from "./schema"
 import { MessageV2 } from "./message-v2"
 import { Log } from "../util/log"
 import { SessionRevert } from "./revert"
@@ -92,7 +92,7 @@ export namespace SessionPrompt {
 
   export const PromptInput = z.object({
     sessionID: SessionID.zod,
-    messageID: Identifier.schema("message").optional(),
+    messageID: MessageID.zod.optional(),
     model: z
       .object({
         providerID: z.string(),
@@ -355,7 +355,7 @@ export namespace SessionPrompt {
         const taskTool = await TaskTool.init()
         const taskModel = task.model ? await Provider.getModel(task.model.providerID, task.model.modelID) : model
         const assistantMessage = (await Session.updateMessage({
-          id: Identifier.ascending("message"),
+          id: MessageID.ascending(),
           role: "assistant",
           parentID: lastUser.id,
           sessionID,
@@ -504,7 +504,7 @@ export namespace SessionPrompt {
           // If we create assistant messages w/ out user ones following mid loop thinking signatures
           // will be missing and it can cause errors for models like gemini for example
           const summaryUserMsg: MessageV2.User = {
-            id: Identifier.ascending("message"),
+            id: MessageID.ascending(),
             sessionID,
             role: "user",
             time: {
@@ -568,7 +568,7 @@ export namespace SessionPrompt {
 
       const processor = SessionProcessor.create({
         assistantMessage: (await Session.updateMessage({
-          id: Identifier.ascending("message"),
+          id: MessageID.ascending(),
           parentID: lastUser.id,
           role: "assistant",
           mode: agent.name,
@@ -971,7 +971,7 @@ export namespace SessionPrompt {
     const variant = input.variant ?? (agent.variant && full?.variants?.[agent.variant] ? agent.variant : undefined)
 
     const info: MessageV2.Info = {
-      id: input.messageID ?? Identifier.ascending("message"),
+      id: input.messageID ?? MessageID.ascending(),
       role: "user",
       sessionID: input.sessionID,
       time: {
@@ -1505,7 +1505,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
     const agent = await Agent.get(input.agent)
     const model = input.model ?? agent.model ?? (await lastModel(input.sessionID))
     const userMsg: MessageV2.User = {
-      id: Identifier.ascending("message"),
+      id: MessageID.ascending(),
       sessionID: input.sessionID,
       time: {
         created: Date.now(),
@@ -1529,7 +1529,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
     await Session.updatePart(userPart)
 
     const msg: MessageV2.Assistant = {
-      id: Identifier.ascending("message"),
+      id: MessageID.ascending(),
       sessionID: input.sessionID,
       parentID: userMsg.id,
       mode: input.agent,
@@ -1719,7 +1719,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
   }
 
   export const CommandInput = z.object({
-    messageID: Identifier.schema("message").optional(),
+    messageID: MessageID.zod.optional(),
     sessionID: SessionID.zod,
     agent: z.string().optional(),
     model: z.string().optional(),

packages/opencode/src/session/revert.ts

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

packages/opencode/src/session/schema.ts

@@ -15,3 +15,15 @@ export const SessionID = sessionIdSchema.pipe(
     zod: z.string().startsWith("ses").pipe(z.custom<SessionID>()),
   })),
 )
+
+const messageIdSchema = Schema.String.pipe(Schema.brand("MessageId"))
+
+export type MessageID = typeof messageIdSchema.Type
+
+export const MessageID = messageIdSchema.pipe(
+  withStatics((schema: typeof messageIdSchema) => ({
+    make: (id: string) => schema.makeUnsafe(id),
+    ascending: (id?: string) => schema.makeUnsafe(Identifier.ascending("message", id)),
+    zod: z.string().startsWith("msg").pipe(z.custom<MessageID>()),
+  })),
+)

packages/opencode/src/session/session.sql.ts

@@ -4,7 +4,8 @@ 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 } from "./schema"
+import type { SessionID, MessageID } from "./schema"
+import type { WorkspaceID } from "../control-plane/schema"
 import { Timestamps } from "../storage/schema.sql"
 
 type PartData = Omit<MessageV2.Part, "id" | "sessionID" | "messageID">
@@ -18,7 +19,7 @@ export const SessionTable = sqliteTable(
       .$type<ProjectID>()
       .notNull()
       .references(() => ProjectTable.id, { onDelete: "cascade" }),
-    workspace_id: text(),
+    workspace_id: text().$type<WorkspaceID>(),
     parent_id: text().$type<SessionID>(),
     slug: text().notNull(),
     directory: text().notNull(),
@@ -29,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: string; partID?: string; snapshot?: string; diff?: string }>(),
+    revert: text({ mode: "json" }).$type<{ messageID: MessageID; partID?: string; snapshot?: string; diff?: string }>(),
     permission: text({ mode: "json" }).$type<PermissionNext.Ruleset>(),
     ...Timestamps,
     time_compacting: integer(),
@@ -45,7 +46,7 @@ export const SessionTable = sqliteTable(
 export const MessageTable = sqliteTable(
   "message",
   {
-    id: text().primaryKey(),
+    id: text().$type<MessageID>().primaryKey(),
     session_id: text()
       .$type<SessionID>()
       .notNull()
@@ -61,6 +62,7 @@ export const PartTable = sqliteTable(
   {
     id: text().primaryKey(),
     message_id: text()
+      .$type<MessageID>()
       .notNull()
       .references(() => MessageTable.id, { onDelete: "cascade" }),
     session_id: text().$type<SessionID>().notNull(),

packages/opencode/src/session/summary.ts

@@ -4,7 +4,7 @@ import { Session } from "."
 
 import { MessageV2 } from "./message-v2"
 import { Identifier } from "@/id/id"
-import { SessionID } from "./schema"
+import { SessionID, MessageID } from "./schema"
 import { Snapshot } from "@/snapshot"
 
 import { Storage } from "@/storage/storage"
@@ -70,7 +70,7 @@ export namespace SessionSummary {
   export const summarize = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: z.string(),
+      messageID: MessageID.zod,
     }),
     async (input) => {
       const all = await Session.messages({ sessionID: input.sessionID })
@@ -115,7 +115,7 @@ export namespace SessionSummary {
   export const diff = fn(
     z.object({
       sessionID: SessionID.zod,
-      messageID: Identifier.schema("message").optional(),
+      messageID: MessageID.zod.optional(),
     }),
     async (input) => {
       const diffs = await Storage.read<Snapshot.FileDiff[]>(["session_diff", input.sessionID]).catch(() => [])