wip: gateway

16 files changed, 617 insertions, 0 deletions

diff --git a/cloud/core/src/account.ts b/cloud/core/src/account.ts
new file mode 100644
index 000000000..cb123e048
--- /dev/null
+++ b/cloud/core/src/account.ts
@@ -0,0 +1,67 @@
+import { z } from "zod"
+import { and, eq, getTableColumns, isNull } from "drizzle-orm"
+import { fn } from "./util/fn"
+import { Database } from "./drizzle"
+import { Identifier } from "./identifier"
+import { AccountTable } from "./schema/account.sql"
+import { Actor } from "./actor"
+import { WorkspaceTable } from "./schema/workspace.sql"
+import { UserTable } from "./schema/user.sql"
+
+export namespace Account {
+  export const create = fn(
+    z.object({
+      email: z.string().email(),
+      id: z.string().optional(),
+    }),
+    async (input) =>
+      Database.transaction(async (tx) => {
+        const id = input.id ?? Identifier.create("account")
+        await tx.insert(AccountTable).values({
+          id,
+          email: input.email,
+        })
+        return id
+      }),
+  )
+
+  export const fromID = fn(z.string(), async (id) =>
+    Database.transaction(async (tx) => {
+      return tx
+        .select()
+        .from(AccountTable)
+        .where(eq(AccountTable.id, id))
+        .execute()
+        .then((rows) => rows[0])
+    }),
+  )
+
+  export const fromEmail = fn(z.string().email(), async (email) =>
+    Database.transaction(async (tx) => {
+      return tx
+        .select()
+        .from(AccountTable)
+        .where(eq(AccountTable.email, email))
+        .execute()
+        .then((rows) => rows[0])
+    }),
+  )
+
+  export const workspaces = async () => {
+    const actor = Actor.assert("account")
+    return Database.transaction(async (tx) =>
+      tx
+        .select(getTableColumns(WorkspaceTable))
+        .from(WorkspaceTable)
+        .innerJoin(UserTable, eq(UserTable.workspaceID, WorkspaceTable.id))
+        .where(
+          and(
+            eq(UserTable.email, actor.properties.email),
+            isNull(UserTable.timeDeleted),
+            isNull(WorkspaceTable.timeDeleted),
+          ),
+        )
+        .execute(),
+    )
+  }
+}
diff --git a/cloud/core/src/actor.ts b/cloud/core/src/actor.ts
new file mode 100644
index 000000000..beb292bb8
--- /dev/null
+++ b/cloud/core/src/actor.ts
@@ -0,0 +1,75 @@
+import { Context } from "./context"
+import { Log } from "./util/log"
+
+export namespace Actor {
+  interface Account {
+    type: "account"
+    properties: {
+      accountID: string
+      email: string
+    }
+  }
+
+  interface Public {
+    type: "public"
+    properties: {}
+  }
+
+  interface User {
+    type: "user"
+    properties: {
+      userID: string
+      workspaceID: string
+      email: string
+    }
+  }
+
+  interface System {
+    type: "system"
+    properties: {
+      workspaceID: string
+    }
+  }
+
+  export type Info = Account | Public | User | System
+
+  const ctx = Context.create<Info>()
+  export const use = ctx.use
+
+  const log = Log.create().tag("namespace", "actor")
+
+  export function provide<R, T extends Info["type"]>(
+    type: T,
+    properties: Extract<Info, { type: T }>["properties"],
+    cb: () => R,
+  ) {
+    return ctx.provide(
+      {
+        type,
+        properties,
+      } as any,
+      () => {
+        return Log.provide({ ...properties }, () => {
+          log.info("provided")
+          return cb()
+        })
+      },
+    )
+  }
+
+  export function assert<T extends Info["type"]>(type: T) {
+    const actor = use()
+    if (actor.type !== type) {
+      throw new Error(`Expected actor type ${type}, got ${actor.type}`)
+    }
+    return actor as Extract<Info, { type: T }>
+  }
+
+  export function workspace() {
+    const actor = use()
+    if ("workspaceID" in actor.properties) {
+      return actor.properties.workspaceID
+    }
+    throw new Error(`actor of type "${actor.type}" is not associated with a workspace`)
+  }
+}
diff --git a/cloud/core/src/billing.ts b/cloud/core/src/billing.ts
new file mode 100644
index 000000000..1a7bb2946
--- /dev/null
+++ b/cloud/core/src/billing.ts
@@ -0,0 +1,71 @@
+import { Resource } from "sst"
+import { Stripe } from "stripe"
+import { Database, eq, sql } from "./drizzle"
+import { BillingTable, UsageTable } from "./schema/billing.sql"
+import { Actor } from "./actor"
+import { fn } from "./util/fn"
+import { z } from "zod"
+import { Identifier } from "./identifier"
+import { centsToMicroCents } from "./util/price"
+
+export namespace Billing {
+  export const stripe = () =>
+    new Stripe(Resource.STRIPE_SECRET_KEY.value, {
+      apiVersion: "2025-03-31.basil",
+    })
+
+  export const get = async () => {
+    return Database.use(async (tx) =>
+      tx
+        .select({
+          customerID: BillingTable.customerID,
+          paymentMethodID: BillingTable.paymentMethodID,
+          balance: BillingTable.balance,
+          reload: BillingTable.reload,
+        })
+        .from(BillingTable)
+        .where(eq(BillingTable.workspaceID, Actor.workspace()))
+        .then((r) => r[0]),
+    )
+  }
+
+  export const consume = fn(
+    z.object({
+      requestID: z.string().optional(),
+      model: z.string(),
+      inputTokens: z.number(),
+      outputTokens: z.number(),
+      reasoningTokens: z.number().optional(),
+      cacheReadTokens: z.number().optional(),
+      cacheWriteTokens: z.number().optional(),
+      costInCents: z.number(),
+    }),
+    async (input) => {
+      const workspaceID = Actor.workspace()
+      const cost = centsToMicroCents(input.costInCents)
+
+      return await Database.transaction(async (tx) => {
+        await tx.insert(UsageTable).values({
+          workspaceID,
+          id: Identifier.create("usage"),
+          requestID: input.requestID,
+          model: input.model,
+          inputTokens: input.inputTokens,
+          outputTokens: input.outputTokens,
+          reasoningTokens: input.reasoningTokens,
+          cacheReadTokens: input.cacheReadTokens,
+          cacheWriteTokens: input.cacheWriteTokens,
+          cost,
+        })
+        const [updated] = await tx
+          .update(BillingTable)
+          .set({
+            balance: sql`${BillingTable.balance} - ${cost}`,
+          })
+          .where(eq(BillingTable.workspaceID, workspaceID))
+          .returning()
+        return updated.balance
+      })
+    },
+  )
+}
diff --git a/cloud/core/src/context.ts b/cloud/core/src/context.ts
new file mode 100644
index 000000000..c2ca6a313
--- /dev/null
+++ b/cloud/core/src/context.ts
@@ -0,0 +1,21 @@
+import { AsyncLocalStorage } from "node:async_hooks"
+
+export namespace Context {
+  export class NotFound extends Error {}
+
+  export function create<T>() {
+    const storage = new AsyncLocalStorage<T>()
+    return {
+      use() {
+        const result = storage.getStore()
+        if (!result) {
+          throw new NotFound()
+        }
+        return result
+      },
+      provide<R>(value: T, fn: () => R) {
+        return storage.run<R>(value, fn)
+      },
+    }
+  }
+}
diff --git a/cloud/core/src/drizzle/index.ts b/cloud/core/src/drizzle/index.ts
new file mode 100644
index 000000000..76220f2a2
--- /dev/null
+++ b/cloud/core/src/drizzle/index.ts
@@ -0,0 +1,94 @@
+import { drizzle } from "drizzle-orm/postgres-js"
+import { Resource } from "sst"
+export * from "drizzle-orm"
+import postgres from "postgres"
+
+function createClient() {
+  const client = postgres({
+    idle_timeout: 30000,
+    connect_timeout: 30000,
+    host: Resource.Database.host,
+    database: Resource.Database.database,
+    user: Resource.Database.username,
+    password: Resource.Database.password,
+    port: Resource.Database.port,
+    ssl: {
+      rejectUnauthorized: false,
+    },
+    max: 1,
+  })
+
+  return drizzle(client, {})
+}
+
+import { PgTransaction, type PgTransactionConfig } from "drizzle-orm/pg-core"
+import type { ExtractTablesWithRelations } from "drizzle-orm"
+import type { PostgresJsQueryResultHKT } from "drizzle-orm/postgres-js"
+import { Context } from "../context"
+
+export namespace Database {
+  export type Transaction = PgTransaction<
+    PostgresJsQueryResultHKT,
+    Record<string, unknown>,
+    ExtractTablesWithRelations<Record<string, unknown>>
+  >
+
+  export type TxOrDb = Transaction | ReturnType<typeof createClient>
+
+  const TransactionContext = Context.create<{
+    tx: TxOrDb
+    effects: (() => void | Promise<void>)[]
+  }>()
+
+  export async function use<T>(callback: (trx: TxOrDb) => Promise<T>) {
+    try {
+      const { tx } = TransactionContext.use()
+      return tx.transaction(callback)
+    } catch (err) {
+      if (err instanceof Context.NotFound) {
+        const client = createClient()
+        const effects: (() => void | Promise<void>)[] = []
+        const result = await TransactionContext.provide(
+          {
+            effects,
+            tx: client,
+          },
+          () => callback(client),
+        )
+        await Promise.all(effects.map((x) => x()))
+        return result
+      }
+      throw err
+    }
+  }
+  export async function fn<Input, T>(callback: (input: Input, trx: TxOrDb) => Promise<T>) {
+    return (input: Input) => use(async (tx) => callback(input, tx))
+  }
+
+  export async function effect(effect: () => any | Promise<any>) {
+    try {
+      const { effects } = TransactionContext.use()
+      effects.push(effect)
+    } catch {
+      await effect()
+    }
+  }
+
+  export async function transaction<T>(callback: (tx: TxOrDb) => Promise<T>, config?: PgTransactionConfig) {
+    try {
+      const { tx } = TransactionContext.use()
+      return callback(tx)
+    } catch (err) {
+      if (err instanceof Context.NotFound) {
+        const client = createClient()
+        const effects: (() => void | Promise<void>)[] = []
+        const result = await client.transaction(async (tx) => {
+          return TransactionContext.provide({ tx, effects }, () => callback(tx))
+        }, config)
+        await Promise.all(effects.map((x) => x()))
+        return result
+      }
+      throw err
+    }
+  }
+}
diff --git a/cloud/core/src/drizzle/types.ts b/cloud/core/src/drizzle/types.ts
new file mode 100644
index 000000000..5ae95d011
--- /dev/null
+++ b/cloud/core/src/drizzle/types.ts
@@ -0,0 +1,29 @@
+import { bigint, timestamp, varchar } from "drizzle-orm/pg-core"
+
+export const ulid = (name: string) => varchar(name, { length: 30 })
+
+export const workspaceColumns = {
+  get id() {
+    return ulid("id").notNull()
+  },
+  get workspaceID() {
+    return ulid("workspace_id").notNull()
+  },
+}
+
+export const id = () => ulid("id").notNull()
+
+export const utc = (name: string) =>
+  timestamp(name, {
+    withTimezone: true,
+  })
+
+export const currency = (name: string) =>
+  bigint(name, {
+    mode: "number",
+  })
+
+export const timestamps = {
+  timeCreated: utc("time_created").notNull().defaultNow(),
+  timeDeleted: utc("time_deleted"),
+}
diff --git a/cloud/core/src/identifier.ts b/cloud/core/src/identifier.ts
new file mode 100644
index 000000000..f8e73852e
--- /dev/null
+++ b/cloud/core/src/identifier.ts
@@ -0,0 +1,26 @@
+import { ulid } from "ulid"
+import { z } from "zod"
+
+export namespace Identifier {
+  const prefixes = {
+    account: "acc",
+    billing: "bil",
+    key: "key",
+    payment: "pay",
+    usage: "usg",
+    user: "usr",
+    workspace: "wrk",
+  } as const
+
+  export function create(prefix: keyof typeof prefixes, given?: string): string {
+    if (given) {
+      if (given.startsWith(prefixes[prefix])) return given
+      throw new Error(`ID ${given} does not start with ${prefixes[prefix]}`)
+    }
+    return [prefixes[prefix], ulid()].join("_")
+  }
+
+  export function schema(prefix: keyof typeof prefixes) {
+    return z.string().startsWith(prefixes[prefix])
+  }
+}
diff --git a/cloud/core/src/schema/account.sql.ts b/cloud/core/src/schema/account.sql.ts
new file mode 100644
index 000000000..1733f0a15
--- /dev/null
+++ b/cloud/core/src/schema/account.sql.ts
@@ -0,0 +1,12 @@
+import { pgTable, uniqueIndex, varchar } from "drizzle-orm/pg-core"
+import { id, timestamps } from "../drizzle/types"
+
+export const AccountTable = pgTable(
+  "account",
+  {
+    id: id(),
+    ...timestamps,
+    email: varchar("email", { length: 255 }).notNull(),
+  },
+  (table) => [uniqueIndex("email").on(table.email)],
+)
diff --git a/cloud/core/src/schema/billing.sql.ts b/cloud/core/src/schema/billing.sql.ts
new file mode 100644
index 000000000..96b29f5de
--- /dev/null
+++ b/cloud/core/src/schema/billing.sql.ts
@@ -0,0 +1,45 @@
+import { bigint, boolean, integer, pgTable, varchar } from "drizzle-orm/pg-core"
+import { timestamps, workspaceColumns } from "../drizzle/types"
+import { workspaceIndexes } from "./workspace.sql"
+
+export const BillingTable = pgTable(
+  "billing",
+  {
+    ...workspaceColumns,
+    ...timestamps,
+    customerID: varchar("customer_id", { length: 255 }),
+    paymentMethodID: varchar("payment_method_id", { length: 255 }),
+    paymentMethodLast4: varchar("payment_method_last4", { length: 4 }),
+    balance: bigint("balance", { mode: "number" }).notNull(),
+    reload: boolean("reload"),
+  },
+  (table) => [...workspaceIndexes(table)],
+)
+
+export const PaymentTable = pgTable(
+  "payment",
+  {
+    ...workspaceColumns,
+    ...timestamps,
+    customerID: varchar("customer_id", { length: 255 }),
+    paymentID: varchar("payment_id", { length: 255 }),
+    amount: bigint("amount", { mode: "number" }).notNull(),
+  },
+  (table) => [...workspaceIndexes(table)],
+)
+
+export const UsageTable = pgTable(
+  "usage",
+  {
+    ...workspaceColumns,
+    ...timestamps,
+    model: varchar("model", { length: 255 }).notNull(),
+    inputTokens: integer("input_tokens").notNull(),
+    outputTokens: integer("output_tokens").notNull(),
+    reasoningTokens: integer("reasoning_tokens"),
+    cacheReadTokens: integer("cache_read_tokens"),
+    cacheWriteTokens: integer("cache_write_tokens"),
+    cost: bigint("cost", { mode: "number" }).notNull(),
+  },
+  (table) => [...workspaceIndexes(table)],
+)
diff --git a/cloud/core/src/schema/key.sql.ts b/cloud/core/src/schema/key.sql.ts
new file mode 100644
index 000000000..240736b86
--- /dev/null
+++ b/cloud/core/src/schema/key.sql.ts
@@ -0,0 +1,16 @@
+import { text, pgTable, varchar, uniqueIndex } from "drizzle-orm/pg-core"
+import { timestamps, utc, workspaceColumns } from "../drizzle/types"
+import { workspaceIndexes } from "./workspace.sql"
+
+export const KeyTable = pgTable(
+  "key",
+  {
+    ...workspaceColumns,
+    ...timestamps,
+    userID: text("user_id").notNull(),
+    name: varchar("name", { length: 255 }).notNull(),
+    key: varchar("key", { length: 255 }).notNull(),
+    timeUsed: utc("time_used"),
+  },
+  (table) => [...workspaceIndexes(table), uniqueIndex("global_key").on(table.key)],
+)
diff --git a/cloud/core/src/schema/user.sql.ts b/cloud/core/src/schema/user.sql.ts
new file mode 100644
index 000000000..34cbd6beb
--- /dev/null
+++ b/cloud/core/src/schema/user.sql.ts
@@ -0,0 +1,16 @@
+import { text, pgTable, uniqueIndex, varchar, integer } from "drizzle-orm/pg-core"
+import { timestamps, utc, workspaceColumns } from "../drizzle/types"
+import { workspaceIndexes } from "./workspace.sql"
+
+export const UserTable = pgTable(
+  "user",
+  {
+    ...workspaceColumns,
+    ...timestamps,
+    email: text("email").notNull(),
+    name: varchar("name", { length: 255 }).notNull(),
+    timeSeen: utc("time_seen"),
+    color: integer("color"),
+  },
+  (table) => [...workspaceIndexes(table), uniqueIndex("user_email").on(table.workspaceID, table.email)],
+)
diff --git a/cloud/core/src/schema/workspace.sql.ts b/cloud/core/src/schema/workspace.sql.ts
new file mode 100644
index 000000000..3e9379e1f
--- /dev/null
+++ b/cloud/core/src/schema/workspace.sql.ts
@@ -0,0 +1,25 @@
+import { primaryKey, foreignKey, pgTable, uniqueIndex, varchar } from "drizzle-orm/pg-core"
+import { timestamps, ulid } from "../drizzle/types"
+
+export const WorkspaceTable = pgTable(
+  "workspace",
+  {
+    id: ulid("id").notNull().primaryKey(),
+    slug: varchar("slug", { length: 255 }),
+    name: varchar("name", { length: 255 }),
+    ...timestamps,
+  },
+  (table) => [uniqueIndex("slug").on(table.slug)],
+)
+
+export function workspaceIndexes(table: any) {
+  return [
+    primaryKey({
+      columns: [table.workspaceID, table.id],
+    }),
+    foreignKey({
+      foreignColumns: [WorkspaceTable.id],
+      columns: [table.workspaceID],
+    }),
+  ]
+}
diff --git a/cloud/core/src/util/fn.ts b/cloud/core/src/util/fn.ts
new file mode 100644
index 000000000..038a50719
--- /dev/null
+++ b/cloud/core/src/util/fn.ts
@@ -0,0 +1,14 @@
+import { z } from "zod"
+
+export function fn<T extends z.ZodType, Result>(
+  schema: T,
+  cb: (input: z.output<T>) => Result,
+) {
+  const result = (input: z.input<T>) => {
+    const parsed = schema.parse(input)
+    return cb(parsed)
+  }
+  result.force = (input: z.input<T>) => cb(input)
+  result.schema = schema
+  return result
+}
diff --git a/cloud/core/src/util/log.ts b/cloud/core/src/util/log.ts
new file mode 100644
index 000000000..4f2d25c13
--- /dev/null
+++ b/cloud/core/src/util/log.ts
@@ -0,0 +1,55 @@
+import { Context } from "../context"
+
+export namespace Log {
+  const ctx = Context.create<{
+    tags: Record<string, any>
+  }>()
+
+  export function create(tags?: Record<string, any>) {
+    tags = tags || {}
+
+    const result = {
+      info(message?: any, extra?: Record<string, any>) {
+        const prefix = Object.entries({
+          ...use().tags,
+          ...tags,
+          ...extra,
+        })
+          .map(([key, value]) => `${key}=${value}`)
+          .join(" ")
+        console.log(prefix, message)
+        return result
+      },
+      tag(key: string, value: string) {
+        if (tags) tags[key] = value
+        return result
+      },
+      clone() {
+        return Log.create({ ...tags })
+      },
+    }
+
+    return result
+  }
+
+  export function provide<R>(tags: Record<string, any>, cb: () => R) {
+    const existing = use()
+    return ctx.provide(
+      {
+        tags: {
+          ...existing.tags,
+          ...tags,
+        },
+      },
+      cb,
+    )
+  }
+
+  function use() {
+    try {
+      return ctx.use()
+    } catch (e) {
+      return { tags: {} }
+    }
+  }
+}
diff --git a/cloud/core/src/util/price.ts b/cloud/core/src/util/price.ts
new file mode 100644
index 000000000..abdbca032
--- /dev/null
+++ b/cloud/core/src/util/price.ts
@@ -0,0 +1,3 @@
+export function centsToMicroCents(amount: number) {
+  return Math.round(amount * 1000000)
+}
diff --git a/cloud/core/src/workspace.ts b/cloud/core/src/workspace.ts
new file mode 100644
index 000000000..532b22963
--- /dev/null
+++ b/cloud/core/src/workspace.ts
@@ -0,0 +1,48 @@
+import { z } from "zod"
+import { fn } from "./util/fn"
+import { centsToMicroCents } from "./util/price"
+import { Actor } from "./actor"
+import { Database, eq } from "./drizzle"
+import { Identifier } from "./identifier"
+import { UserTable } from "./schema/user.sql"
+import { BillingTable } from "./schema/billing.sql"
+import { WorkspaceTable } from "./schema/workspace.sql"
+
+export namespace Workspace {
+  export const create = fn(z.void(), async () => {
+    const account = Actor.assert("account")
+    const workspaceID = Identifier.create("workspace")
+    await Database.transaction(async (tx) => {
+      await tx.insert(WorkspaceTable).values({
+        id: workspaceID,
+      })
+      await tx.insert(UserTable).values({
+        workspaceID,
+        id: Identifier.create("user"),
+        email: account.properties.email,
+        name: "",
+      })
+      await tx.insert(BillingTable).values({
+        workspaceID,
+        id: Identifier.create("billing"),
+        balance: centsToMicroCents(100),
+      })
+    })
+    return workspaceID
+  })
+
+  export async function list() {
+    const account = Actor.assert("account")
+    return Database.use(async (tx) => {
+      return tx
+        .select({
+          id: WorkspaceTable.id,
+          slug: WorkspaceTable.slug,
+          name: WorkspaceTable.name,
+        })
+        .from(UserTable)
+        .innerJoin(WorkspaceTable, eq(UserTable.workspaceID, WorkspaceTable.id))
+        .where(eq(UserTable.email, account.properties.email))
+    })
+  }
+}