kuber/server/workspace-store.ts

import { randomUUID } from "node:crypto";

export const KUBER_API_VERSION = "kuber.astrxl.dev/v2" as const;
export const MAX_WORKSPACE_CONFIG_BYTES = 768 * 1024;
export const MAX_SOURCE_REFERENCE_BYTES = 64 * 1024;
export const MAX_REVISION_PAYLOAD_BYTES = 900 * 1024;

export const WORKSPACE_ID_PATTERN = /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$/;
export const RESERVED_WORKSPACE_IDS: ReadonlySet<string> = new Set([
  "database",
  "garage-system",
  "routing",
]);

export interface ObjectMeta {
  name: string;
  uid: string;
  resourceVersion: string;
  creationTimestamp: string;
  labels?: Record<string, string>;
  annotations?: Record<string, string>;
}

export interface WorkspaceSourceReference {
  uri: string;
  digest: string;
  revision?: string;
}

export interface WorkspaceSpec {
  source: WorkspaceSourceReference;
  config?: unknown;
}

export interface WorkspaceStatus {
  latestRevision: number;
}

export interface Workspace {
  apiVersion: typeof KUBER_API_VERSION;
  kind: "Workspace";
  metadata: ObjectMeta;
  spec: WorkspaceSpec;
  status: WorkspaceStatus;
}

export interface WorkspaceRevision {
  apiVersion: typeof KUBER_API_VERSION;
  kind: "WorkspaceRevision";
  metadata: ObjectMeta & {
    workspaceUid: string;
  };
  spec: Readonly<WorkspaceSpec> & {
    workspaceId: string;
    revision: number;
  };
}

export interface CreateWorkspaceInput {
  id: string;
  source: WorkspaceSourceReference;
  config?: unknown;
  labels?: Record<string, string>;
  annotations?: Record<string, string>;
}

export interface UpdateWorkspaceInput {
  source: WorkspaceSourceReference;
  config?: unknown;
  labels?: Record<string, string>;
  annotations?: Record<string, string>;
}

export class WorkspaceValidationError extends Error {
  readonly code = "WORKSPACE_INVALID";
}

export class WorkspaceConflictError extends Error {
  readonly code = "WORKSPACE_CONFLICT";
}

export class WorkspaceNotFoundError extends Error {
  readonly code = "WORKSPACE_NOT_FOUND";
}

/** Persistence must atomically compare resourceVersion and append the revision. */
export interface WorkspacePersistence {
  get(id: string): Promise<Workspace | undefined>;
  list(): Promise<Workspace[]>;
  create(workspace: Workspace, revision: WorkspaceRevision): Promise<void>;
  replace(
    workspace: Workspace,
    revision: WorkspaceRevision,
    expectedResourceVersion: string,
  ): Promise<void>;
  getRevision(
    workspaceId: string,
    revision: number,
  ): Promise<WorkspaceRevision | undefined>;
  listRevisions(workspaceId: string): Promise<WorkspaceRevision[]>;
}

export interface WorkspaceStore {
  create(input: CreateWorkspaceInput): Promise<Workspace>;
  get(id: string): Promise<Workspace | undefined>;
  list(): Promise<Workspace[]>;
  update(
    id: string,
    input: UpdateWorkspaceInput,
    ifMatch: string,
  ): Promise<Workspace>;
  getRevision(
    id: string,
    revision: number,
  ): Promise<WorkspaceRevision | undefined>;
  listRevisions(id: string): Promise<WorkspaceRevision[]>;
}

export interface WorkspaceStoreOptions {
  now?: () => Date;
  uid?: () => string;
}

export function validateWorkspaceId(id: string): void {
  if (!WORKSPACE_ID_PATTERN.test(id)) {
    throw new WorkspaceValidationError(
      "Workspace ID must be a lowercase DNS label of at most 63 characters",
    );
  }
  if (id.startsWith("kube-") || RESERVED_WORKSPACE_IDS.has(id)) {
    throw new WorkspaceValidationError(`Workspace ID '${id}' is reserved`);
  }
}

export function workspaceEtag(workspace: Pick<Workspace, "metadata">): string {
  return `"${workspace.metadata.resourceVersion}"`;
}

export function resourceVersionFromEtag(etag: string): string {
  const value = etag.trim();
  const match = /^(?:W\/)?"([^"\\]+)"$/.exec(value);
  if (!match?.[1]) {
    throw new WorkspaceValidationError("If-Match must contain a valid ETag");
  }
  return match[1];
}

function encodedSize(value: unknown): number {
  let serialized: string;
  try {
    serialized = JSON.stringify(value);
  } catch {
    throw new WorkspaceValidationError(
      "Workspace payload must be JSON serializable",
    );
  }
  if (serialized === undefined) {
    throw new WorkspaceValidationError(
      "Workspace payload must be JSON serializable",
    );
  }
  return Buffer.byteLength(serialized);
}

function validateSpec(spec: WorkspaceSpec): void {
  if (!spec.source || typeof spec.source !== "object") {
    throw new WorkspaceValidationError("A source reference is required");
  }
  if (!spec.source.uri?.trim() || !spec.source.digest?.trim()) {
    throw new WorkspaceValidationError("Source uri and digest are required");
  }
  const source = spec.source as unknown as Record<string, unknown>;
  for (const field of ["blob", "content", "data", "archive", "files"]) {
    if (field in source) {
      throw new WorkspaceValidationError(
        "Inline source blobs are not supported; provide a source reference",
      );
    }
  }
  if (encodedSize(spec.source) > MAX_SOURCE_REFERENCE_BYTES) {
    throw new WorkspaceValidationError("Source reference is too large");
  }
  if (encodedSize(spec.config ?? null) > MAX_WORKSPACE_CONFIG_BYTES) {
    throw new WorkspaceValidationError("Workspace config is too large");
  }
  if (encodedSize(spec) > MAX_REVISION_PAYLOAD_BYTES) {
    throw new WorkspaceValidationError(
      "Workspace revision payload is too large",
    );
  }
}

function clone<T>(value: T): T {
  return structuredClone(value);
}

function revisionFor(
  workspace: Workspace,
  uid: () => string,
  creationTimestamp: string,
): WorkspaceRevision {
  const revision = workspace.status.latestRevision;
  return {
    apiVersion: KUBER_API_VERSION,
    kind: "WorkspaceRevision",
    metadata: {
      name: `${workspace.metadata.name}-r${revision}`,
      uid: uid(),
      workspaceUid: workspace.metadata.uid,
      resourceVersion: workspace.metadata.resourceVersion,
      creationTimestamp,
    },
    spec: clone({
      ...workspace.spec,
      workspaceId: workspace.metadata.name,
      revision,
    }),
  };
}

export class PersistentWorkspaceStore implements WorkspaceStore {
  private readonly now: () => Date;
  private readonly uid: () => string;

  constructor(
    private readonly persistence: WorkspacePersistence,
    options: WorkspaceStoreOptions = {},
  ) {
    this.now = options.now ?? (() => new Date());
    this.uid = options.uid ?? randomUUID;
  }

  async create(input: CreateWorkspaceInput): Promise<Workspace> {
    validateWorkspaceId(input.id);
    const spec = clone({ source: input.source, config: input.config });
    validateSpec(spec);
    const workspace: Workspace = {
      apiVersion: KUBER_API_VERSION,
      kind: "Workspace",
      metadata: {
        name: input.id,
        uid: this.uid(),
        resourceVersion: "1",
        creationTimestamp: this.now().toISOString(),
        ...(input.labels && { labels: clone(input.labels) }),
        ...(input.annotations && { annotations: clone(input.annotations) }),
      },
      spec,
      status: { latestRevision: 1 },
    };
    await this.persistence.create(
      workspace,
      revisionFor(workspace, this.uid, workspace.metadata.creationTimestamp),
    );
    return clone(workspace);
  }

  async get(id: string): Promise<Workspace | undefined> {
    const workspace = await this.persistence.get(id);
    return workspace && clone(workspace);
  }

  async list(): Promise<Workspace[]> {
    return clone(await this.persistence.list());
  }

  async update(
    id: string,
    input: UpdateWorkspaceInput,
    ifMatch: string,
  ): Promise<Workspace> {
    validateWorkspaceId(id);
    const current = await this.persistence.get(id);
    if (!current)
      throw new WorkspaceNotFoundError(`Workspace '${id}' not found`);
    const expected = resourceVersionFromEtag(ifMatch);
    if (expected !== current.metadata.resourceVersion) {
      throw new WorkspaceConflictError("Workspace ETag does not match");
    }
    const spec = clone({ source: input.source, config: input.config });
    validateSpec(spec);
    const workspace: Workspace = {
      ...clone(current),
      metadata: {
        ...clone(current.metadata),
        resourceVersion: String(Number(current.metadata.resourceVersion) + 1),
        ...(input.labels !== undefined && { labels: clone(input.labels) }),
        ...(input.annotations !== undefined && {
          annotations: clone(input.annotations),
        }),
      },
      spec,
      status: { latestRevision: current.status.latestRevision + 1 },
    };
    await this.persistence.replace(
      workspace,
      revisionFor(workspace, this.uid, this.now().toISOString()),
      expected,
    );
    return clone(workspace);
  }

  async getRevision(id: string, revision: number) {
    const value = await this.persistence.getRevision(id, revision);
    return value && clone(value);
  }

  async listRevisions(id: string) {
    return clone(await this.persistence.listRevisions(id));
  }
}

export class MemoryWorkspacePersistence implements WorkspacePersistence {
  private readonly workspaces = new Map<string, Workspace>();
  private readonly revisions = new Map<
    string,
    Map<number, WorkspaceRevision>
  >();

  async get(id: string) {
    const value = this.workspaces.get(id);
    return value && clone(value);
  }

  async list() {
    return [...this.workspaces.values()]
      .sort((a, b) => a.metadata.name.localeCompare(b.metadata.name))
      .map(clone);
  }

  async create(workspace: Workspace, revision: WorkspaceRevision) {
    if (this.workspaces.has(workspace.metadata.name)) {
      throw new WorkspaceConflictError("Workspace already exists");
    }
    this.workspaces.set(workspace.metadata.name, clone(workspace));
    this.revisions.set(
      workspace.metadata.name,
      new Map([[revision.spec.revision, clone(revision)]]),
    );
  }

  async replace(
    workspace: Workspace,
    revision: WorkspaceRevision,
    expectedResourceVersion: string,
  ) {
    const current = this.workspaces.get(workspace.metadata.name);
    if (!current) throw new WorkspaceNotFoundError("Workspace not found");
    if (current.metadata.resourceVersion !== expectedResourceVersion) {
      throw new WorkspaceConflictError("Workspace was concurrently modified");
    }
    const revisions = this.revisions.get(workspace.metadata.name);
    if (!revisions || revisions.has(revision.spec.revision)) {
      throw new WorkspaceConflictError("Workspace revision already exists");
    }
    this.workspaces.set(workspace.metadata.name, clone(workspace));
    revisions.set(revision.spec.revision, clone(revision));
  }

  async getRevision(id: string, revision: number) {
    const value = this.revisions.get(id)?.get(revision);
    return value && clone(value);
  }

  async listRevisions(id: string) {
    return [...(this.revisions.get(id)?.values() ?? [])]
      .sort((a, b) => a.spec.revision - b.spec.revision)
      .map(clone);
  }
}

export class MemoryWorkspaceStore extends PersistentWorkspaceStore {
  constructor(options: WorkspaceStoreOptions = {}) {
    super(new MemoryWorkspacePersistence(), options);
  }
}