feat(session): shard-aware todo reads and update export

Author: sjawhar

packages/opencode/src/effect/runner.ts

@@ -184,8 +184,8 @@ export namespace Runner {
         case "ShellThenRun":
           return [
             Effect.gen(function* () {
-              yield* Deferred.fail(st.run.done, new Cancelled()).pipe(Effect.asVoid)
               yield* stopShell(st.shell)
+              yield* Deferred.fail(st.run.done, new Cancelled()).pipe(Effect.asVoid)
               yield* idleIfCurrent()
             }),
             { _tag: "Idle" } as const,

packages/opencode/src/project/project.ts

@@ -133,6 +133,12 @@ export namespace Project {
       const db = <T>(fn: (d: Parameters<typeof Database.use>[0] extends (trx: infer D) => any ? D : never) => T) =>
         Effect.sync(() => Database.use(fn))
 
+      type Sync<T> = T extends Promise<any> ? never : T
+
+      const tx = <T>(
+        fn: (d: Parameters<typeof Database.transaction>[0] extends (trx: infer D) => any ? D : never) => Sync<T>,
+      ) => Effect.sync(() => Database.transaction(fn, { behavior: "immediate" }))
+
       const emitUpdated = (data: Info) =>
         Effect.sync(() =>
           GlobalBus.emit("event", {
@@ -246,7 +252,7 @@ export namespace Project {
         })
 
         // Phase 2: upsert
-        const row = yield* db((d) => d.select().from(ProjectTable).where(eq(ProjectTable.id, data.id)).get())
+        const row = yield* tx((d) => d.select().from(ProjectTable).where(eq(ProjectTable.id, data.id)).get())
         const existing = row
           ? fromRow(row)
           : {
@@ -278,9 +284,8 @@ export namespace Project {
           { concurrency: "unbounded" },
         ).pipe(Effect.map((arr) => arr.filter((x): x is string => x !== undefined)))
 
-        yield* db((d) =>
-          d
-            .insert(ProjectTable)
+        yield* tx((d) => {
+          d.insert(ProjectTable)
             .values({
               id: result.id,
               worktree: result.worktree,
@@ -308,18 +313,15 @@ export namespace Project {
                 commands: result.commands,
               },
             })
-            .run(),
-        )
+            .run()
 
-        if (data.id !== ProjectID.global) {
-          yield* db((d) =>
-            d
-              .update(SessionTable)
-              .set({ project_id: data.id })
-              .where(and(eq(SessionTable.project_id, ProjectID.global), eq(SessionTable.directory, data.worktree)))
-              .run(),
-          )
-        }
+          if (data.id === ProjectID.global) return
+
+          d.update(SessionTable)
+            .set({ project_id: data.id })
+            .where(and(eq(SessionTable.project_id, ProjectID.global), eq(SessionTable.directory, data.worktree)))
+            .run()
+        })
 
         yield* emitUpdated(result)
         return { project: result, sandbox: data.sandbox }

packages/opencode/src/session/index.ts

@@ -23,6 +23,7 @@ import { Snapshot } from "@/snapshot"
 import { ProjectID } from "../project/schema"
 import { WorkspaceID } from "../control-plane/schema"
 import { SessionID, MessageID, PartID } from "./schema"
+import { makeRuntime } from "@/effect/run-service"
 
 import type { Provider } from "@/provider/provider"
 import { Permission } from "@/permission"
@@ -424,6 +425,14 @@ export namespace Session {
 
         yield* Effect.sync(() => SyncEvent.run(Event.Created, { sessionID: result.id, info: result }))
 
+        if (!input.parentID) {
+          try {
+            Database.session(result.id)
+          } catch (e) {
+            log.error("failed to create per-tree db", { error: e })
+          }
+        }
+
         if (!Flag.OPENCODE_EXPERIMENTAL_WORKSPACES) {
           // This only exist for backwards compatibility. We should not be
           // manually publishing this event; it is a sync event now
@@ -498,19 +507,17 @@ export namespace Session {
         }).pipe(Effect.withSpan("Session.updatePart"))
 
       const getPart: Interface["getPart"] = Effect.fn("Session.getPart")(function* (input) {
-        const row = Database.use((db) =>
-          db
-            .select()
-            .from(PartTable)
-            .where(
-              and(
-                eq(PartTable.session_id, input.sessionID),
-                eq(PartTable.message_id, input.messageID),
-                eq(PartTable.id, input.partID),
-              ),
-            )
-            .get(),
-        )
+        const row = Database.resolveSession(input.sessionID)
+          .select()
+          .from(PartTable)
+          .where(
+            and(
+              eq(PartTable.session_id, input.sessionID),
+              eq(PartTable.message_id, input.messageID),
+              eq(PartTable.id, input.partID),
+            ),
+          )
+          .get()
         if (!row) return
         return {
           ...row.data,
@@ -704,6 +711,24 @@ export namespace Session {
 
   export const defaultLayer = layer.pipe(Layer.provide(Bus.layer), Layer.provide(Storage.defaultLayer))
 
+  const { runPromise } = makeRuntime(Service, defaultLayer)
+
+  export async function create(input?: CreateInput) {
+    return runPromise((svc) => svc.create(input))
+  }
+
+  export async function remove(sessionID: SessionID) {
+    return runPromise((svc) => svc.remove(sessionID))
+  }
+
+  export async function updateMessage<T extends MessageV2.Info>(msg: T) {
+    return runPromise((svc) => svc.updateMessage(msg))
+  }
+
+  export async function updatePart<T extends MessageV2.Part>(part: T) {
+    return runPromise((svc) => svc.updatePart(part))
+  }
+
   export function* list(input?: {
     directory?: string
     workspaceID?: WorkspaceID

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

@@ -550,18 +550,20 @@ export namespace MessageV2 {
       and(eq(MessageTable.time_created, row.time), lt(MessageTable.id, row.id)),
     )
 
-  function hydrate(rows: (typeof MessageTable.$inferSelect)[]) {
+  function resolve(sid: SessionID) {
+    return Database.resolveSession(sid)
+  }
+
+  function hydrate(db: Database.TxOrDb, rows: (typeof MessageTable.$inferSelect)[]) {
     const ids = rows.map((row) => row.id)
     const partByMessage = new Map<string, MessageV2.Part[]>()
     if (ids.length > 0) {
-      const partRows = Database.use((db) =>
-        db
-          .select()
-          .from(PartTable)
-          .where(inArray(PartTable.message_id, ids))
-          .orderBy(PartTable.message_id, PartTable.id)
-          .all(),
-      )
+      const partRows = db
+        .select()
+        .from(PartTable)
+        .where(inArray(PartTable.message_id, ids))
+        .orderBy(PartTable.message_id, PartTable.id)
+        .all()
       for (const row of partRows) {
         const next = part(row)
         const list = partByMessage.get(row.message_id)
@@ -850,15 +852,14 @@ export namespace MessageV2 {
     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(),
-    )
+    const db = resolve(input.sessionID)
+    const rows = 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(),
@@ -872,7 +873,7 @@ export namespace MessageV2 {
 
     const more = rows.length > input.limit
     const slice = more ? rows.slice(0, input.limit) : rows
-    const items = hydrate(slice)
+    const items = hydrate(db, slice)
     items.reverse()
     const tail = slice.at(-1)
     return {
@@ -896,33 +897,28 @@ export namespace MessageV2 {
     }
   }
 
-  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(),
-    )
-    return rows.map(
-      (row) =>
-        ({
-          ...row.data,
-          id: row.id,
-          sessionID: row.session_id,
-          messageID: row.message_id,
-        }) as MessageV2.Part,
-    )
+  export function parts(mid: MessageID, sid?: SessionID) {
+    const db = sid ? resolve(sid) : Database.Client()
+    return db.select().from(PartTable).where(eq(PartTable.message_id, mid)).orderBy(PartTable.id).all().map(part)
   }
 
   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(),
-    )
+    const db = resolve(input.sessionID)
+    const row = 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),
+      parts: db
+        .select()
+        .from(PartTable)
+        .where(eq(PartTable.message_id, input.messageID))
+        .orderBy(PartTable.id)
+        .all()
+        .map(part),
     }
   }
 
@@ -944,9 +940,7 @@ export namespace MessageV2 {
     return result
   }
 
-  export const filterCompactedEffect = Effect.fnUntraced(function* (sessionID: SessionID) {
-    return filterCompacted(stream(sessionID))
-  })
+  export const filterCompactedEffect = (sessionID: SessionID) => Effect.succeed(filterCompacted(stream(sessionID)))
 
   export function fromError(
     e: unknown,
@@ -998,7 +992,7 @@ export namespace MessageV2 {
           },
           { cause: e },
         ).toObject()
-      case APICallError.isInstance(e):
+      case APICallError.isInstance(e): {
         const parsed = ProviderError.parseAPICallError({
           providerID: ctx.providerID,
           error: e,
@@ -1024,6 +1018,7 @@ export namespace MessageV2 {
           },
           { cause: e },
         ).toObject()
+      }
       case e instanceof Error:
         return new NamedError.Unknown({ message: errorMessage(e) }, { cause: e }).toObject()
       default:

packages/opencode/src/session/todo.ts

@@ -4,6 +4,7 @@ import { SessionID } from "./schema"
 import { Effect, Layer, Context } from "effect"
 import z from "zod"
 import { Database, eq, asc } from "../storage/db"
+import { makeRuntime } from "@/effect/run-service"
 import { TodoTable } from "./session.sql"
 
 export namespace Todo {
@@ -39,11 +40,12 @@ export namespace Todo {
       const bus = yield* Bus.Service
 
       const update = Effect.fn("Todo.update")(function* (input: { sessionID: SessionID; todos: Info[] }) {
-        yield* Effect.sync(() =>
-          Database.transaction((db) => {
-            db.delete(TodoTable).where(eq(TodoTable.session_id, input.sessionID)).run()
+        yield* Effect.sync(() => {
+          const db = Database.resolveSession(input.sessionID)
+          db.transaction((tx) => {
+            tx.delete(TodoTable).where(eq(TodoTable.session_id, input.sessionID)).run()
             if (input.todos.length === 0) return
-            db.insert(TodoTable)
+            tx.insert(TodoTable)
               .values(
                 input.todos.map((todo, position) => ({
                   session_id: input.sessionID,
@@ -54,22 +56,21 @@ export namespace Todo {
                 })),
               )
               .run()
-          }),
-        )
+          })
+        })
         yield* bus.publish(Event.Updated, input)
       })
 
       const get = Effect.fn("Todo.get")(function* (sessionID: SessionID) {
-        const rows = yield* Effect.sync(() =>
-          Database.use((db) =>
-            db
-              .select()
-              .from(TodoTable)
-              .where(eq(TodoTable.session_id, sessionID))
-              .orderBy(asc(TodoTable.position))
-              .all(),
-          ),
-        )
+        const rows = yield* Effect.sync(() => {
+          const db = Database.resolveSession(sessionID)
+          return db
+            .select()
+            .from(TodoTable)
+            .where(eq(TodoTable.session_id, sessionID))
+            .orderBy(asc(TodoTable.position))
+            .all()
+        })
         return rows.map((row) => ({
           content: row.content,
           status: row.status,
@@ -82,4 +83,14 @@ export namespace Todo {
   )
 
   export const defaultLayer = layer.pipe(Layer.provide(Bus.layer))
+
+  const { runPromise } = makeRuntime(Service, defaultLayer)
+
+  export async function get(sessionID: SessionID) {
+    return runPromise((svc) => svc.get(sessionID))
+  }
+
+  export async function update(input: { sessionID: SessionID; todos: Info[] }) {
+    return runPromise((svc) => svc.update(input))
+  }
 }