enterprise (#4617)

Co-authored-by: GitHub Action <[email protected]>
Co-authored-by: Adam <[email protected]>

2 files changed, 273 insertions, 0 deletions

diff --git a/packages/enterprise/src/core/share.ts b/packages/enterprise/src/core/share.ts
new file mode 100644
index 000000000..a7dfbfdcc
--- /dev/null
+++ b/packages/enterprise/src/core/share.ts
@@ -0,0 +1,139 @@
+import { FileDiff, Message, Part, Session, SessionStatus } from "@opencode-ai/sdk"
+import { fn } from "@opencode-ai/util/fn"
+import { iife } from "@opencode-ai/util/iife"
+import z from "zod"
+import { Storage } from "./storage"
+
+export namespace Share {
+  export const Info = z.object({
+    id: z.string(),
+    secret: z.string(),
+  })
+  export type Info = z.infer<typeof Info>
+
+  export const Data = z.discriminatedUnion("type", [
+    z.object({
+      type: z.literal("session"),
+      data: z.custom<Session>(),
+    }),
+    z.object({
+      type: z.literal("message"),
+      data: z.custom<Message>(),
+    }),
+    z.object({
+      type: z.literal("part"),
+      data: z.custom<Part>(),
+    }),
+    z.object({
+      type: z.literal("session_diff"),
+      data: z.custom<FileDiff[]>(),
+    }),
+    z.object({
+      type: z.literal("session_status"),
+      data: z.custom<SessionStatus>(),
+    }),
+  ])
+  export type Data = z.infer<typeof Data>
+
+  export const create = fn(Info.pick({ id: true }), async (body) => {
+    const info: Info = {
+      id: body.id,
+      secret: crypto.randomUUID(),
+    }
+    const exists = await get(info.id)
+    if (exists) throw new Errors.AlreadyExists(info.id)
+    await Storage.write(["share", info.id], info)
+    console.log("created share", info.id)
+    return info
+  })
+
+  async function get(sessionID: string) {
+    return Storage.read<Info>(["share", sessionID])
+  }
+
+  export const remove = fn(Info.pick({ id: true, secret: true }), async (body) => {
+    const share = await get(body.id)
+    if (!share) throw new Errors.NotFound(body.id)
+    if (share.secret !== body.secret) throw new Errors.InvalidSecret(body.id)
+    await Storage.remove(["share", body.id])
+    const list = await Storage.list(["share_data", body.id])
+    for (const item of list) {
+      await Storage.remove(item)
+    }
+  })
+
+  export async function data(sessionID: string) {
+    const list = await Storage.list(["share_data", sessionID])
+    const promises = []
+    for (const item of list) {
+      promises.push(
+        iife(async () => {
+          const [, , type] = item
+          return {
+            type: type as any,
+            data: await Storage.read<any>(item),
+          } as Data
+        }),
+      )
+    }
+    return await Promise.all(promises)
+  }
+
+  export const sync = fn(
+    z.object({
+      share: Info,
+      data: Data.array(),
+    }),
+    async (input) => {
+      const share = await get(input.share.id)
+      if (!share) throw new Errors.NotFound(input.share.id)
+      if (share.secret !== input.share.secret) throw new Errors.InvalidSecret(input.share.id)
+      const promises = []
+      for (const item of input.data) {
+        promises.push(
+          iife(async () => {
+            switch (item.type) {
+              case "session":
+                await Storage.write(["share_data", input.share.id, "session"], item.data)
+                break
+              case "message":
+                await Storage.write(["share_data", input.share.id, "message", item.data.id], item.data)
+                break
+              case "part":
+                await Storage.write(
+                  ["share_data", input.share.id, "part", item.data.messageID, item.data.id],
+                  item.data,
+                )
+                break
+              case "session_diff":
+                await Storage.write(["share_data", input.share.id, "session_diff"], item.data)
+                break
+              case "session_status":
+                await Storage.write(["share_data", input.share.id, "session_status"], item.data)
+                break
+            }
+          }),
+        )
+      }
+      await Promise.all(promises)
+    },
+  )
+
+  export const Errors = {
+    NotFound: class extends Error {
+      constructor(public id: string) {
+        super(`Share not found: ${id}`)
+      }
+    },
+    InvalidSecret: class extends Error {
+      constructor(public id: string) {
+        super(`Share secret invalid: ${id}`)
+      }
+    },
+    AlreadyExists: class extends Error {
+      constructor(public id: string) {
+        super(`Share already exists: ${id}`)
+      }
+    },
+  }
+}
diff --git a/packages/enterprise/src/core/storage.ts b/packages/enterprise/src/core/storage.ts
new file mode 100644
index 000000000..fb33ef389
--- /dev/null
+++ b/packages/enterprise/src/core/storage.ts
@@ -0,0 +1,134 @@
+import {
+  S3Client,
+  PutObjectCommand,
+  GetObjectCommand,
+  DeleteObjectCommand,
+  ListObjectsV2Command,
+} from "@aws-sdk/client-s3"
+import { lazy } from "@opencode-ai/util/lazy"
+
+export namespace Storage {
+  export interface Adapter {
+    read(path: string): Promise<string | undefined>
+    write(path: string, value: string): Promise<void>
+    remove(path: string): Promise<void>
+    list(prefix: string): Promise<string[]>
+  }
+
+  function createAdapter(client: S3Client, bucket: string): Adapter {
+    return {
+      async read(path: string): Promise<string | undefined> {
+        try {
+          console.log("reading", bucket, path)
+          const command = new GetObjectCommand({
+            Bucket: bucket,
+            Key: path,
+          })
+          const response = await client.send(command)
+          if (!response.Body) return undefined
+          return response.Body.transformToString()
+        } catch (e: any) {
+          if (e.name === "NoSuchKey") return undefined
+          throw e
+        }
+      },
+
+      async write(path: string, value: string): Promise<void> {
+        const command = new PutObjectCommand({
+          Bucket: bucket,
+          Key: path,
+          Body: value,
+          ContentType: "application/json",
+        })
+        await client.send(command)
+      },
+
+      async remove(path: string): Promise<void> {
+        const command = new DeleteObjectCommand({
+          Bucket: bucket,
+          Key: path,
+        })
+        await client.send(command)
+      },
+
+      async list(prefix: string): Promise<string[]> {
+        const command = new ListObjectsV2Command({
+          Bucket: bucket,
+          Prefix: prefix,
+        })
+        const response = await client.send(command)
+        return response.Contents?.map((c) => c.Key!) || []
+      },
+    }
+  }
+
+  function s3(): Adapter {
+    const bucket = process.env.OPENCODE_STORAGE_BUCKET!
+    const client = new S3Client({
+      region: process.env.OPENCODE_STORAGE_REGION,
+      credentials: process.env.OPENCODE_STORAGE_ACCESS_KEY_ID
+        ? {
+            accessKeyId: process.env.OPENCODE_STORAGE_ACCESS_KEY_ID!,
+            secretAccessKey: process.env.OPENCODE_STORAGE_SECRET_ACCESS_KEY!,
+          }
+        : undefined,
+    })
+    return createAdapter(client, bucket)
+  }
+
+  function r2() {
+    const accountId = process.env.OPENCODE_STORAGE_ACCOUNT_ID!
+    const accessKeyId = process.env.OPENCODE_STORAGE_ACCESS_KEY_ID!
+    const secretAccessKey = process.env.OPENCODE_STORAGE_SECRET_ACCESS_KEY!
+    const bucket = process.env.OPENCODE_STORAGE_BUCKET!
+
+    const client = new S3Client({
+      region: "auto",
+      endpoint: `https://${accountId}.r2.cloudflarestorage.com`,
+      credentials: {
+        accessKeyId,
+        secretAccessKey,
+      },
+    })
+    return createAdapter(client, bucket)
+  }
+
+  const adapter = lazy(() => {
+    const type = process.env.OPENCODE_STORAGE_ADAPTER
+    if (type === "r2") return r2()
+    if (type === "s3") return s3()
+    throw new Error("No storage adapter configured")
+  })
+
+  function resolve(key: string[]) {
+    return key.join("/") + ".json"
+  }
+
+  export async function read<T>(key: string[]) {
+    const result = await adapter().read(resolve(key))
+    if (!result) return undefined
+    return JSON.parse(result) as T
+  }
+
+  export function write<T>(key: string[], value: T) {
+    return adapter().write(resolve(key), JSON.stringify(value))
+  }
+
+  export function remove(key: string[]) {
+    return adapter().remove(resolve(key))
+  }
+
+  export async function list(prefix: string[]) {
+    const p = prefix.join("/") + (prefix.length ? "/" : "")
+    const result = await adapter().list(p)
+    return result.map((x) => x.replace(/\.json$/, "").split("/"))
+  }
+
+  export async function update<T>(key: string[], fn: (draft: T) => void) {
+    const val = await read<T>(key)
+    if (!val) throw new Error("Not found")
+    fn(val)
+    await write(key, val)
+    return val
+  }
+}