Przeglądaj źródła

feat(core): add workspace lifecycle

Kit Langton 4 tygodni temu
rodzic
commit
650d5a5e92

+ 183 - 14
packages/core/src/workspace.ts

@@ -1,6 +1,6 @@
 export * as WorkspaceV2 from "./workspace"
 
-import { Context, Effect, Equal, Exit, Layer, RcMap, Schema, Scope } from "effect"
+import { Context, Effect, Equal, Exit, Latch, Layer, RcMap, Schema, Scope, Semaphore } from "effect"
 import { Workspace } from "@opencode-ai/schema/workspace"
 import { eq } from "drizzle-orm"
 import { Database } from "./database/database"
@@ -16,6 +16,13 @@ export type ID = typeof ID.Type
 export const Info = Workspace.Info
 export type Info = Workspace.Info
 
+export interface CreateInput {
+  readonly provider: string
+  readonly name: string
+  readonly directory: AbsolutePath
+  readonly projectID: Info["project"]["id"]
+}
+
 export class NotFoundError extends Schema.TaggedErrorClass<NotFoundError>()("Workspace.NotFoundError", {
   id: ID,
 }) {}
@@ -26,7 +33,11 @@ export class InvalidError extends Schema.TaggedErrorClass<InvalidError>()("Works
 }) {}
 
 export interface Interface {
+  readonly create: (input: CreateInput) => Effect.Effect<Info, Sandbox.Error | Sandbox.ProviderNotFoundError>
   readonly get: (id: ID) => Effect.Effect<Info, NotFoundError | InvalidError>
+  readonly connect: (
+    id: ID,
+  ) => Effect.Effect<void, NotFoundError | InvalidError | Sandbox.Error | Sandbox.ProviderNotFoundError>
   readonly borrow: (
     id: ID,
   ) => Effect.Effect<
@@ -34,6 +45,12 @@ export interface Interface {
     NotFoundError | InvalidError | Sandbox.Error | Sandbox.ProviderNotFoundError,
     Scope.Scope
   >
+  readonly suspend: (
+    id: ID,
+  ) => Effect.Effect<void, NotFoundError | InvalidError | Sandbox.Error | Sandbox.ProviderNotFoundError>
+  readonly remove: (
+    id: ID,
+  ) => Effect.Effect<void, NotFoundError | InvalidError | Sandbox.Error | Sandbox.ProviderNotFoundError>
 }
 
 export class Service extends Context.Service<Service, Interface>()("@opencode/Workspace") {}
@@ -43,6 +60,10 @@ const layer = Layer.effect(
   Effect.gen(function* () {
     const db = (yield* Database.Service).db
     const registry = yield* Sandbox.RegistryService
+    const lifecycle = yield* RcMap.make({
+      idleTimeToLive: 0,
+      lookup: () => Effect.succeed(makeLifecycleGate()),
+    })
 
     const row = Effect.fn("Workspace.row")(function* (id: ID) {
       const value = yield* db.select().from(WorkspaceTable).where(eq(WorkspaceTable.id, id)).get().pipe(Effect.orDie)
@@ -62,6 +83,38 @@ const layer = Layer.effect(
       })
     })
 
+    const create = Effect.fn("Workspace.create")(function* (input: CreateInput) {
+      const id = ID.create()
+      const provider = yield* registry.get(input.provider)
+      return yield* Effect.acquireUseRelease(
+        provider.create({ identity: id }),
+        (binding) =>
+          db
+            .insert(WorkspaceTable)
+            .values({
+              id,
+              type: provider.key,
+              name: input.name,
+              directory: input.directory,
+              extra: Sandbox.Placement.make({ kind: "sandbox", version: 1, binding }),
+              project_id: input.projectID,
+            })
+            .run()
+            .pipe(
+              Effect.orDie,
+              Effect.as(
+                Info.make({
+                  id,
+                  name: input.name,
+                  directory: input.directory,
+                  project: { id: input.projectID, directory: input.directory },
+                }),
+              ),
+            ),
+        (binding, exit) => (Exit.isFailure(exit) ? Effect.uninterruptible(provider.delete(binding)) : Effect.void),
+      )
+    })
+
     const persistBinding = Effect.fnUntraced(function* (id: ID, previous: Sandbox.Binding, next: Sandbox.Binding) {
       if (Equal.equals(previous, next)) return
       yield* db
@@ -72,29 +125,91 @@ const layer = Layer.effect(
         .pipe(Effect.orDie)
     })
 
+    const placement = Effect.fnUntraced(function* (id: ID) {
+      const value = yield* row(id)
+      if (!Schema.is(Sandbox.Placement)(value.extra)) {
+        return yield* new InvalidError({ id, message: "Workspace has no sandbox binding" })
+      }
+      const provider = yield* registry.get(value.type)
+      return { provider, binding: yield* provider.decode(value.extra.binding) }
+    })
+
+    const reconcile = Effect.fnUntraced(function* (id: ID, provider: Sandbox.Provider, binding: Sandbox.Binding) {
+      const next = yield* provider.reconcile(binding)
+      yield* persistBinding(id, binding, next)
+      return next
+    })
+
     const connections = yield* RcMap.make({
       idleTimeToLive: "1 minute",
       lookup: Effect.fn("Workspace.connect")(function* (id: ID) {
-        const placement = yield* row(id)
-        if (!Schema.is(Sandbox.Placement)(placement.extra)) {
-          return yield* new InvalidError({ id, message: "Workspace has no sandbox binding" })
-        }
-        const provider = yield* registry.get(placement.type)
-        const binding = yield* provider.decode(placement.extra.binding)
+        const current = yield* placement(id)
+        const provider = current.provider
+        const binding = current.binding
         const connection = yield* provider.connect(binding)
         yield* persistBinding(id, binding, connection.binding)
-        yield* persistBinding(id, connection.binding, yield* provider.reconcile(connection.binding))
-        return connection
+        yield* reconcile(id, provider, connection.binding)
+        return { provider, connection }
       }),
     })
 
+    const lease = (id: ID) => RcMap.get(lifecycle, id).pipe(Effect.flatMap((gate) => gate.read))
+
+    const borrow = (id: ID) =>
+      lease(id).pipe(
+        Effect.andThen(RcMap.get(connections, id)),
+        Effect.onExit((exit) => (Exit.isFailure(exit) ? RcMap.invalidate(connections, id) : Effect.void)),
+        Effect.map((value) => value.connection.environment),
+      )
+
+    const transition = <A, E, R>(id: ID, effect: Effect.Effect<A, E, R>) =>
+      Effect.scoped(RcMap.get(lifecycle, id).pipe(Effect.flatMap((gate) => gate.write(effect))))
+
+    const suspend = Effect.fn("Workspace.suspend")((id: ID) =>
+      transition(
+        id,
+        Effect.scoped(
+          RcMap.get(connections, id).pipe(
+            Effect.flatMap((value) =>
+              Effect.gen(function* () {
+                const binding = yield* Effect.uninterruptible(
+                  value.provider
+                    .suspend(value.connection)
+                    .pipe(Effect.tap((binding) => persistBinding(id, value.connection.binding, binding))),
+                )
+                yield* reconcile(id, value.provider, binding)
+              }),
+            ),
+          ),
+        ).pipe(Effect.ensuring(RcMap.invalidate(connections, id))),
+      ),
+    )
+
+    const remove = Effect.fn("Workspace.remove")((id: ID) =>
+      transition(
+        id,
+        Effect.gen(function* () {
+          yield* RcMap.invalidate(connections, id)
+          const current = yield* placement(id)
+          const binding = yield* reconcile(id, current.provider, current.binding)
+          yield* Effect.uninterruptible(
+            current.provider
+              .delete(binding)
+              .pipe(
+                Effect.andThen(db.delete(WorkspaceTable).where(eq(WorkspaceTable.id, id)).run().pipe(Effect.orDie)),
+              ),
+          )
+        }),
+      ),
+    )
+
     return Service.of({
+      create,
       get,
-      borrow: (id) =>
-        RcMap.get(connections, id).pipe(
-          Effect.onExit((exit) => (Exit.isFailure(exit) ? RcMap.invalidate(connections, id) : Effect.void)),
-          Effect.map((connection) => connection.environment),
-        ),
+      connect: (id) => Effect.scoped(borrow(id)).pipe(Effect.asVoid),
+      borrow,
+      suspend,
+      remove,
     })
   }),
 )
@@ -104,3 +219,57 @@ export const node = makeGlobalNode({
   layer,
   deps: [Database.node, Sandbox.registryNode],
 })
+
+function makeLifecycleGate() {
+  const admission = Semaphore.makeUnsafe(1)
+  const transition = Semaphore.makeUnsafe(1)
+  const open = Latch.makeUnsafe(true)
+  const drained = Latch.makeUnsafe(true)
+  let active = 0
+  let blocked = false
+
+  const read: Effect.Effect<void, never, Scope.Scope> = Effect.suspend(() =>
+    Effect.uninterruptibleMask((restore) =>
+      Effect.gen(function* () {
+        yield* restore(open.await)
+        yield* restore(admission.take(1))
+        if (blocked) {
+          yield* admission.release(1)
+          return yield* restore(read)
+        }
+        active++
+        if (active === 1) drained.closeUnsafe()
+        const scope = yield* Scope.Scope
+        yield* Scope.addFinalizer(
+          scope,
+          Effect.sync(() => {
+            active--
+            if (active === 0) drained.openUnsafe()
+          }),
+        )
+        yield* admission.release(1)
+      }),
+    ),
+  )
+
+  const write = <A, E, R>(effect: Effect.Effect<A, E, R>) =>
+    transition.withPermit(
+      Effect.uninterruptibleMask((restore) =>
+        Effect.gen(function* () {
+          yield* restore(admission.take(1))
+          blocked = true
+          open.closeUnsafe()
+          yield* admission.release(1)
+          return yield* restore(drained.await.pipe(Effect.andThen(effect))).pipe(
+            Effect.ensuring(
+              Effect.sync(() => {
+                blocked = false
+              }).pipe(Effect.andThen(open.open)),
+            ),
+          )
+        }),
+      ),
+    )
+
+  return { read, write }
+}

+ 7 - 0
packages/core/src/workspace/sandbox.ts

@@ -25,11 +25,18 @@ export interface Connection {
   readonly environment: WorkspaceEnvironment.Interface
 }
 
+export interface CreateInput {
+  readonly identity: string
+}
+
 export interface Provider {
   readonly key: string
+  readonly create: (input: CreateInput) => Effect.Effect<Binding, Error>
   readonly decode: (binding: Binding) => Effect.Effect<Binding, Error>
   readonly connect: (binding: Binding) => Effect.Effect<Connection, Error, Scope.Scope>
+  readonly suspend: (connection: Connection) => Effect.Effect<Binding, Error>
   readonly reconcile: (binding: Binding) => Effect.Effect<Binding, Error>
+  readonly delete: (binding: Binding) => Effect.Effect<void, Error>
 }
 
 export class DuplicateProviderError extends Schema.TaggedErrorClass<DuplicateProviderError>()(

+ 4 - 0
packages/core/test/location.test.ts

@@ -92,6 +92,7 @@ describe("Location", () => {
         const workspaceLayer = Layer.succeed(
           WorkspaceV2.Service,
           WorkspaceV2.Service.of({
+            create: () => Effect.die("Unsupported fake Workspace operation"),
             get: () =>
               Effect.succeed(
                 WorkspaceV2.Info.make({
@@ -109,6 +110,9 @@ describe("Location", () => {
                 connections.count++
                 return providerEnvironment
               }),
+            connect: () => Effect.die("Unsupported fake Workspace operation"),
+            suspend: () => Effect.die("Unsupported fake Workspace operation"),
+            remove: () => Effect.die("Unsupported fake Workspace operation"),
           }),
         )
         const hostedRef = { directory, workspaceID }

+ 113 - 9
packages/core/test/workspace.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect } from "bun:test"
-import { Effect, Exit } from "effect"
+import { Deferred, Effect, Exit, Fiber, Scope } from "effect"
 import { adjust } from "effect/testing/TestClock"
 import { eq } from "drizzle-orm"
 import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder"
@@ -29,7 +29,17 @@ describe("WorkspaceV2", () => {
       const id = WorkspaceV2.ID.make("wrk_hosted")
       const projectID = Project.ID.make("hosted-project")
       const directory = AbsolutePath.make("/workspace/repo")
-      const lifecycle = { connected: 0, reconciled: 0, released: 0 }
+      const lifecycle = {
+        created: [] as string[],
+        connected: [] as Sandbox.Binding[],
+        reconciled: 0,
+        released: 0,
+        suspended: 0,
+        deleted: [] as Sandbox.Binding[],
+        failDelete: false,
+      }
+      const suspendStarted = yield* Deferred.make<void>()
+      const finishSuspend = yield* Deferred.make<void>()
       const unsupported = (operation: string) => Effect.fail(new WorkspaceEnvironment.Error({ operation }))
       const environment = WorkspaceEnvironment.Service.of({
         platform: "linux",
@@ -79,20 +89,38 @@ describe("WorkspaceV2", () => {
         .run()
       yield* registry.register({
         key: "fake",
+        create: (input) =>
+          Effect.sync(() => {
+            lifecycle.created.push(input.identity)
+            return { sandbox: input.identity }
+          }),
         decode: Effect.succeed,
-        connect: () =>
+        connect: (binding) =>
           Effect.acquireRelease(
             Effect.sync(() => {
-              lifecycle.connected++
+              lifecycle.connected.push(binding)
               return { binding: { sandbox: "live", retired: "one" }, environment }
             }),
             () => Effect.sync(() => lifecycle.released++),
           ),
-        reconcile: () =>
+        suspend: () =>
+          Deferred.succeed(suspendStarted, undefined).pipe(
+            Effect.andThen(Deferred.await(finishSuspend)),
+            Effect.tap(() => Effect.sync(() => lifecycle.suspended++)),
+            Effect.as({ snapshot: "snap-one", retired: "live" }),
+          ),
+        reconcile: (binding) =>
           Effect.sync(() => {
             lifecycle.reconciled++
-            return { sandbox: "live" }
+            const next: Sandbox.Binding = "snapshot" in binding ? { snapshot: binding.snapshot } : { sandbox: "live" }
+            return next
           }),
+        delete: (binding) => {
+          if (lifecycle.failDelete) {
+            return Effect.fail(new Sandbox.Error({ provider: "fake", operation: "delete" }))
+          }
+          return Effect.sync(() => lifecycle.deleted.push(binding))
+        },
       })
 
       expect(yield* workspace.get(id)).toEqual({
@@ -101,12 +129,12 @@ describe("WorkspaceV2", () => {
         directory,
         project: { id: projectID, directory },
       })
-      expect(lifecycle.connected).toBe(0)
+      expect(lifecycle.connected).toHaveLength(0)
 
       const borrowed = yield* Effect.all([workspace.borrow(id), workspace.borrow(id)]).pipe(Effect.scoped)
       expect(borrowed[0]).toBe(environment)
       expect(borrowed[1]).toBe(environment)
-      expect(lifecycle.connected).toBe(1)
+      expect(lifecycle.connected).toHaveLength(1)
       expect(lifecycle.reconciled).toBe(1)
       expect(lifecycle.released).toBe(0)
       const placement = yield* db.select().from(WorkspaceTable).where(eq(WorkspaceTable.id, id)).get()
@@ -135,7 +163,7 @@ describe("WorkspaceV2", () => {
         .run()
       const invalid = yield* workspace.borrow(invalidID).pipe(Effect.scoped, Effect.flip)
       expect(invalid._tag).toBe("Workspace.InvalidError")
-      expect(lifecycle.connected).toBe(1)
+      expect(lifecycle.connected).toHaveLength(1)
 
       const retryID = WorkspaceV2.ID.make("wrk_retry")
       const retry = { attempts: 0 }
@@ -153,6 +181,7 @@ describe("WorkspaceV2", () => {
         .run()
       yield* registry.register({
         key: "flaky",
+        create: () => Effect.die("Unexpected create"),
         decode: Effect.succeed,
         connect: (binding) =>
           Effect.sync(() => ++retry.attempts).pipe(
@@ -160,12 +189,87 @@ describe("WorkspaceV2", () => {
               attempt === 1 ? Effect.die("Transient provider defect") : Effect.succeed({ binding, environment }),
             ),
           ),
+        suspend: () => Effect.die("Unexpected suspend"),
         reconcile: Effect.succeed,
+        delete: () => Effect.die("Unexpected delete"),
       })
 
       expect(Exit.isFailure(yield* workspace.borrow(retryID).pipe(Effect.scoped, Effect.exit))).toBe(true)
       expect(yield* workspace.borrow(retryID).pipe(Effect.scoped)).toBe(environment)
       expect(retry.attempts).toBe(2)
+
+      const created = yield* workspace.create({
+        provider: "fake",
+        name: "Created",
+        directory,
+        projectID,
+      })
+      expect(created).toEqual({
+        id: created.id,
+        name: "Created",
+        directory,
+        project: { id: projectID, directory },
+      })
+      expect(lifecycle.created.at(-1)).toBe(created.id)
+
+      const failedCreate = yield* workspace
+        .create({
+          provider: "fake",
+          name: "Invalid project",
+          directory,
+          projectID: Project.ID.make("missing-project"),
+        })
+        .pipe(Effect.exit)
+      expect(Exit.isFailure(failedCreate)).toBe(true)
+      const failedIdentity = lifecycle.created.at(-1)
+      if (!failedIdentity) return yield* Effect.die("Missing failed create identity")
+      expect(lifecycle.deleted).toContainEqual({ sandbox: failedIdentity })
+
+      const active = yield* Scope.make()
+      yield* workspace.borrow(id).pipe(Scope.provide(active))
+      const suspendInvoked = yield* Deferred.make<void>()
+      const suspendFiber = yield* Deferred.succeed(suspendInvoked, undefined).pipe(
+        Effect.andThen(workspace.suspend(id)),
+        Effect.forkChild,
+      )
+      yield* Deferred.await(suspendInvoked)
+      yield* Effect.yieldNow
+      expect(yield* Deferred.isDone(suspendStarted)).toBe(false)
+
+      const admitted = yield* Deferred.make<void>()
+      const borrowFiber = yield* workspace.borrow(id).pipe(
+        Effect.tap(() => Deferred.succeed(admitted, undefined)),
+        Effect.scoped,
+        Effect.forkChild,
+      )
+      yield* Effect.yieldNow
+      expect(yield* Deferred.isDone(admitted)).toBe(false)
+
+      yield* Scope.close(active, Exit.void)
+      yield* Deferred.await(suspendStarted)
+      expect(yield* Deferred.isDone(admitted)).toBe(false)
+
+      yield* Deferred.succeed(finishSuspend, undefined)
+      yield* Fiber.join(suspendFiber)
+      yield* Fiber.join(borrowFiber)
+      expect(lifecycle.suspended).toBe(1)
+      expect(lifecycle.connected).toContainEqual({ snapshot: "snap-one" })
+      expect((yield* db.select().from(WorkspaceTable).where(eq(WorkspaceTable.id, id)).get())?.extra).toEqual({
+        kind: "sandbox",
+        version: 1,
+        binding: { sandbox: "live" },
+      })
+
+      lifecycle.failDelete = true
+      const deleteError = yield* workspace.remove(created.id).pipe(Effect.flip)
+      expect(deleteError._tag).toBe("Sandbox.Error")
+      if (deleteError._tag !== "Sandbox.Error") return yield* Effect.die("Unexpected delete error")
+      expect(deleteError.operation).toBe("delete")
+      expect((yield* workspace.get(created.id)).id).toBe(created.id)
+      lifecycle.failDelete = false
+      yield* workspace.remove(created.id)
+      expect((yield* workspace.get(created.id).pipe(Effect.flip))._tag).toBe("Workspace.NotFoundError")
+      expect(lifecycle.deleted).toContainEqual({ sandbox: "live" })
     }),
   )
 })