anomalyco_opencode/packages/opencode/src/control-plane/workspace.ts

import z from "zod"
import { setTimeout as sleep } from "node:timers/promises"
import { fn } from "@/util/fn"
import { Database, asc, eq } from "@/storage/db"
import { Project } from "@/project/project"
import { BusEvent } from "@/bus/bus-event"
import { GlobalBus } from "@/bus/global"
import { SyncEvent } from "@/sync"
import { EventTable } from "@/sync/event.sql"
import { Flag } from "@/flag/flag"
import { Log } from "@/util/log"
import { Filesystem } from "@/util/filesystem"
import { ProjectID } from "@/project/schema"
import { Slug } from "@opencode-ai/shared/util/slug"
import { WorkspaceTable } from "./workspace.sql"
import { getAdaptor } from "./adaptors"
import { WorkspaceInfo } from "./types"
import { WorkspaceID } from "./schema"
import { parseSSE } from "./sse"
import { Session } from "@/session"
import { SessionTable } from "@/session/session.sql"
import { SessionID } from "@/session/schema"
import { errorData } from "@/util/error"
import { AppRuntime } from "@/effect/app-runtime"

export namespace Workspace {
  export const Info = WorkspaceInfo.meta({
    ref: "Workspace",
  })
  export type Info = z.infer<typeof Info>

  export const ConnectionStatus = z.object({
    workspaceID: WorkspaceID.zod,
    status: z.enum(["connected", "connecting", "disconnected", "error"]),
    error: z.string().optional(),
  })
  export type ConnectionStatus = z.infer<typeof ConnectionStatus>

  const Restore = z.object({
    workspaceID: WorkspaceID.zod,
    sessionID: SessionID.zod,
    total: z.number().int().min(0),
    step: z.number().int().min(0),
  })

  export const Event = {
    Ready: BusEvent.define(
      "workspace.ready",
      z.object({
        name: z.string(),
      }),
    ),
    Failed: BusEvent.define(
      "workspace.failed",
      z.object({
        message: z.string(),
      }),
    ),
    Restore: BusEvent.define("workspace.restore", Restore),
    Status: BusEvent.define("workspace.status", ConnectionStatus),
  }

  function fromRow(row: typeof WorkspaceTable.$inferSelect): Info {
    return {
      id: row.id,
      type: row.type,
      branch: row.branch,
      name: row.name,
      directory: row.directory,
      extra: row.extra,
      projectID: row.project_id,
    }
  }

  const CreateInput = z.object({
    id: WorkspaceID.zod.optional(),
    type: Info.shape.type,
    branch: Info.shape.branch,
    projectID: ProjectID.zod,
    extra: Info.shape.extra,
  })

  export const create = fn(CreateInput, async (input) => {
    const id = WorkspaceID.ascending(input.id)
    const adaptor = await getAdaptor(input.projectID, input.type)

    const config = await adaptor.configure({ ...input, id, name: Slug.create(), directory: null })

    const info: Info = {
      id,
      type: config.type,
      branch: config.branch ?? null,
      name: config.name ?? null,
      directory: config.directory ?? null,
      extra: config.extra ?? null,
      projectID: input.projectID,
    }

    Database.use((db) => {
      db.insert(WorkspaceTable)
        .values({
          id: info.id,
          type: info.type,
          branch: info.branch,
          name: info.name,
          directory: info.directory,
          extra: info.extra,
          project_id: info.projectID,
        })
        .run()
    })

    await adaptor.create(config)

    startSync(info)

    return info
  })

  const SessionRestoreInput = z.object({
    workspaceID: WorkspaceID.zod,
    sessionID: SessionID.zod,
  })

  export const sessionRestore = fn(SessionRestoreInput, async (input) => {
    log.info("session restore requested", {
      workspaceID: input.workspaceID,
      sessionID: input.sessionID,
    })
    try {
      const space = await get(input.workspaceID)
      if (!space) throw new Error(`Workspace not found: ${input.workspaceID}`)

      const adaptor = await getAdaptor(space.projectID, space.type)
      const target = await adaptor.target(space)

      // Need to switch the workspace of the session
      SyncEvent.run(Session.Event.Updated, {
        sessionID: input.sessionID,
        info: {
          workspaceID: input.workspaceID,
        },
      })

      const rows = Database.use((db) =>
        db
          .select({
            id: EventTable.id,
            aggregateID: EventTable.aggregate_id,
            seq: EventTable.seq,
            type: EventTable.type,
            data: EventTable.data,
          })
          .from(EventTable)
          .where(eq(EventTable.aggregate_id, input.sessionID))
          .orderBy(asc(EventTable.seq))
          .all(),
      )
      if (rows.length === 0) throw new Error(`No events found for session: ${input.sessionID}`)

      const all = rows

      const size = 10
      const sets = Array.from({ length: Math.ceil(all.length / size) }, (_, i) => all.slice(i * size, (i + 1) * size))
      const total = sets.length
      log.info("session restore prepared", {
        workspaceID: input.workspaceID,
        sessionID: input.sessionID,
        workspaceType: space.type,
        directory: space.directory,
        target: target.type === "remote" ? String(route(target.url, "/sync/replay")) : target.directory,
        events: all.length,
        batches: total,
        first: all[0]?.seq,
        last: all.at(-1)?.seq,
      })
      GlobalBus.emit("event", {
        directory: "global",
        workspace: input.workspaceID,
        payload: {
          type: Event.Restore.type,
          properties: {
            workspaceID: input.workspaceID,
            sessionID: input.sessionID,
            total,
            step: 0,
          },
        },
      })
      for (const [i, events] of sets.entries()) {
        log.info("session restore batch starting", {
          workspaceID: input.workspaceID,
          sessionID: input.sessionID,
          step: i + 1,
          total,
          events: events.length,
          first: events[0]?.seq,
          last: events.at(-1)?.seq,
          target: target.type === "remote" ? String(route(target.url, "/sync/replay")) : target.directory,
        })
        if (target.type === "local") {
          SyncEvent.replayAll(events)
          log.info("session restore batch replayed locally", {
            workspaceID: input.workspaceID,
            sessionID: input.sessionID,
            step: i + 1,
            total,
            events: events.length,
          })
        } else {
          const url = route(target.url, "/sync/replay")
          const headers = new Headers(target.headers)
          headers.set("content-type", "application/json")
          const res = await fetch(url, {
            method: "POST",
            headers,
            body: JSON.stringify({
              directory: space.directory ?? "",
              events,
            }),
          })
          if (!res.ok) {
            const body = await res.text()
            log.error("session restore batch failed", {
              workspaceID: input.workspaceID,
              sessionID: input.sessionID,
              step: i + 1,
              total,
              status: res.status,
              body,
            })
            throw new Error(
              `Failed to replay session ${input.sessionID} into workspace ${input.workspaceID}: HTTP ${res.status} ${body}`,
            )
          }
          log.info("session restore batch posted", {
            workspaceID: input.workspaceID,
            sessionID: input.sessionID,
            step: i + 1,
            total,
            status: res.status,
          })
        }
        GlobalBus.emit("event", {
          directory: "global",
          workspace: input.workspaceID,
          payload: {
            type: Event.Restore.type,
            properties: {
              workspaceID: input.workspaceID,
              sessionID: input.sessionID,
              total,
              step: i + 1,
            },
          },
        })
      }

      log.info("session restore complete", {
        workspaceID: input.workspaceID,
        sessionID: input.sessionID,
        batches: total,
      })

      return {
        total,
      }
    } catch (err) {
      log.error("session restore failed", {
        workspaceID: input.workspaceID,
        sessionID: input.sessionID,
        error: errorData(err),
      })
      throw err
    }
  })

  export function list(project: Project.Info) {
    const rows = Database.use((db) =>
      db.select().from(WorkspaceTable).where(eq(WorkspaceTable.project_id, project.id)).all(),
    )
    const spaces = rows.map(fromRow).sort((a, b) => a.id.localeCompare(b.id))

    for (const space of spaces) startSync(space)
    return spaces
  }

  export const get = fn(WorkspaceID.zod, async (id) => {
    const row = Database.use((db) => db.select().from(WorkspaceTable).where(eq(WorkspaceTable.id, id)).get())
    if (!row) return
    const space = fromRow(row)
    startSync(space)
    return space
  })

  export const remove = fn(WorkspaceID.zod, async (id) => {
    const sessions = Database.use((db) =>
      db.select({ id: SessionTable.id }).from(SessionTable).where(eq(SessionTable.workspace_id, id)).all(),
    )
    for (const session of sessions) {
      await AppRuntime.runPromise(Session.Service.use((svc) => svc.remove(session.id)))
    }

    const row = Database.use((db) => db.select().from(WorkspaceTable).where(eq(WorkspaceTable.id, id)).get())

    if (row) {
      stopSync(id)

      const info = fromRow(row)
      try {
        const adaptor = await getAdaptor(info.projectID, row.type)
        await adaptor.remove(info)
      } catch (err) {
        log.error("adaptor not available when removing workspace", { type: row.type })
      }
      Database.use((db) => db.delete(WorkspaceTable).where(eq(WorkspaceTable.id, id)).run())
      return info
    }
  })

  const connections = new Map<WorkspaceID, ConnectionStatus>()
  const aborts = new Map<WorkspaceID, AbortController>()

  function setStatus(id: WorkspaceID, status: ConnectionStatus["status"], error?: string) {
    const prev = connections.get(id)
    if (prev?.status === status && prev?.error === error) return
    const next = { workspaceID: id, status, error }
    connections.set(id, next)
    GlobalBus.emit("event", {
      directory: "global",
      workspace: id,
      payload: {
        type: Event.Status.type,
        properties: next,
      },
    })
  }

  export function status(): ConnectionStatus[] {
    return [...connections.values()]
  }

  const log = Log.create({ service: "workspace-sync" })

  function route(url: string | URL, path: string) {
    const next = new URL(url)
    next.pathname = `${next.pathname.replace(/\/$/, "")}${path}`
    next.search = ""
    next.hash = ""
    return next
  }

  async function syncWorkspace(space: Info, signal: AbortSignal) {
    while (!signal.aborted) {
      log.info("connecting to global sync", { workspace: space.name })

      const adaptor = await getAdaptor(space.projectID, space.type)
      const target = await adaptor.target(space)

      if (target.type === "local") return

      const res = await fetch(route(target.url, "/global/event"), {
        method: "GET",
        headers: target.headers,
        signal,
      }).catch((err: unknown) => {
        setStatus(space.id, "error")

        log.info("failed to connect to global sync", {
          workspace: space.name,
          error: err,
        })
        return undefined
      })

      if (!res || !res.ok || !res.body) {
        log.info("failed to connect to global sync", { workspace: space.name })
        setStatus(space.id, "error")
        await sleep(1000)
        continue
      }

      log.info("global sync connected", { workspace: space.name })
      setStatus(space.id, "connected")

      await parseSSE(res.body, signal, (evt: any) => {
        try {
          if (!("payload" in evt)) return

          if (evt.payload.type === "sync") {
            // This name -> type is temporary
            SyncEvent.replay({ ...evt.payload, type: evt.payload.name } as SyncEvent.SerializedEvent)
          }

          GlobalBus.emit("event", {
            directory: evt.directory,
            project: evt.project,
            workspace: space.id,
            payload: evt.payload,
          })
        } catch (err) {
          log.info("failed to replay global event", {
            workspaceID: space.id,
            error: err,
          })
        }
      })

      log.info("disconnected from global sync: " + space.id)
      setStatus(space.id, "disconnected")

      // TODO: Implement exponential backoff
      await sleep(1000)
    }
  }

  function startSync(space: Info) {
    if (!Flag.OPENCODE_EXPERIMENTAL_WORKSPACES) return

    if (space.type === "worktree") {
      void Filesystem.exists(space.directory!).then((exists) => {
        setStatus(space.id, exists ? "connected" : "error", exists ? undefined : "directory does not exist")
      })
      return
    }

    if (aborts.has(space.id)) return
    const abort = new AbortController()
    aborts.set(space.id, abort)
    setStatus(space.id, "disconnected")

    void syncWorkspace(space, abort.signal).catch((error) => {
      setStatus(space.id, "error", String(error))
      log.warn("workspace listener failed", {
        workspaceID: space.id,
        error,
      })
    })
  }

  function stopSync(id: WorkspaceID) {
    aborts.get(id)?.abort()
    aborts.delete(id)
    connections.delete(id)
  }
}