pim-agent 0.0.1 → 0.10.0

package/packages/server/src/AttachmentEndpoint.ts

@@ -0,0 +1,184 @@
+import { basename, dirname, join } from "node:path";
+
+import {
+  AttachmentStore,
+  type StoredAttachment,
+} from "#core/attachments/AttachmentStore";
+import { Paths } from "#core/shared/Paths";
+import { IMMUTABLE, serveFile } from "./StaticClient";
+
+export type AttachmentEndpointDeps = {
+  /** Defaults to `~/.pim/attachments`. */
+  readonly root?: string;
+  readonly maxBytes?: number;
+};
+
+const DEFAULT_MAX_BYTES = 25 * 1024 * 1024;
+
+const PREFIX = "/attachment/";
+
+/** Where stored bytes live when nothing says otherwise. */
+export function defaultAttachmentsRoot(): string {
+  return join(Paths.pimHomeDir(), "attachments");
+}
+
+// Every origin is allowed: in dev the client is served from vite's own port and its upload is cross-origin.
+function allow(response: Response): Response {
+  response.headers.set("access-control-allow-origin", "*");
+  response.headers.set("access-control-allow-headers", "content-type");
+  return response;
+}
+
+/** Where `serve` answers for a stored file: the last two path segments only, never the server's layout. */
+export function attachmentUrl(path: string): string {
+  const scope = encodeURIComponent(basename(dirname(path)));
+  return `${PREFIX}${scope}/${encodeURIComponent(basename(path))}`;
+}
+
+/** `POST /upload?session=<id>` with a multipart `file` field, and `GET /attachment/<session>/<file>` to read it back. */
+export class AttachmentEndpoint {
+  private readonly store: AttachmentStore;
+  private readonly maxBytes: number;
+  private readonly bySession = new Map<string, Map<string, StoredAttachment>>();
+
+  public constructor(deps: AttachmentEndpointDeps = {}) {
+    this.store = new AttachmentStore(deps.root ?? defaultAttachmentsRoot());
+    this.maxBytes = deps.maxBytes ?? DEFAULT_MAX_BYTES;
+  }
+
+  /** Whether this endpoint owns the path, so the gateway can route to it. */
+  public static owns(pathname: string): boolean {
+    return pathname === "/upload" || pathname.startsWith(PREFIX);
+  }
+
+  public async handle(request: Request): Promise<Response> {
+    if (request.method === "OPTIONS") {
+      return allow(new Response(null, { status: 204 }));
+    }
+    const { pathname } = new URL(request.url);
+    return allow(
+      pathname.startsWith(PREFIX)
+        ? await this.serve(pathname)
+        : await this.upload(request)
+    );
+  }
+
+  // Every segment goes back through the write's sanitising, so a crafted name cannot escape its scope.
+  private async serve(pathname: string): Promise<Response> {
+    const [scope, name, ...rest] = pathname
+      .slice(PREFIX.length)
+      .split("/")
+      .map(decodeURIComponentSafely);
+    if (!scope || !name || rest.length > 0) {
+      return new Response("not found", { status: 404 });
+    }
+    let path: string;
+    try {
+      path = this.store.locate(scope, name);
+    } catch {
+      return new Response("not found", { status: 404 });
+    }
+    return serveFile(path, IMMUTABLE);
+  }
+
+  private async upload(request: Request): Promise<Response> {
+    if (request.method !== "POST") {
+      return new Response("expected POST", { status: 405 });
+    }
+    const sessionId = new URL(request.url).searchParams.get("session");
+    if (!sessionId) {
+      return Response.json({ error: "missing ?session" }, { status: 400 });
+    }
+    const declared = Number(request.headers.get("content-length") ?? 0);
+    if (declared > this.maxBytes) {
+      return Response.json(
+        { error: `upload exceeds ${this.maxBytes} bytes` },
+        { status: 413 }
+      );
+    }
+
+    let file: File | undefined;
+    try {
+      const form = await request.formData();
+      const field = form.get("file");
+      file = field instanceof File ? field : undefined;
+    } catch {
+      return Response.json(
+        { error: "expected multipart/form-data" },
+        { status: 400 }
+      );
+    }
+    if (!file) {
+      return Response.json({ error: "missing `file` field" }, { status: 400 });
+    }
+    if (file.size > this.maxBytes) {
+      return Response.json(
+        { error: `upload exceeds ${this.maxBytes} bytes` },
+        { status: 413 }
+      );
+    }
+
+    let stored: StoredAttachment;
+    try {
+      stored = await this.store.store(sessionId, {
+        bytes: await file.arrayBuffer(),
+        mimeType: file.type || "application/octet-stream",
+        name: file.name,
+        stem: stemOf(file.name),
+      });
+    } catch (err) {
+      return Response.json({ error: (err as Error).message }, { status: 400 });
+    }
+    this.remember(sessionId, stored);
+
+    return Response.json({
+      id: stored.id,
+      path: stored.path,
+      url: attachmentUrl(stored.path),
+      mimeType: stored.mimeType,
+      isImage: stored.imageBase64 !== undefined,
+    });
+  }
+
+  /** Consume the uploads a prompt referenced; unknown ids are dropped, so a replayed command cannot re-attach. */
+  public take(
+    sessionId: string,
+    ids: readonly string[]
+  ): readonly StoredAttachment[] {
+    const pending = this.bySession.get(sessionId);
+    if (!pending) {
+      return [];
+    }
+    const taken: StoredAttachment[] = [];
+    for (const id of ids) {
+      const stored = pending.get(id);
+      if (stored) {
+        pending.delete(id);
+        taken.push(stored);
+      }
+    }
+    if (pending.size === 0) {
+      this.bySession.delete(sessionId);
+    }
+    return taken;
+  }
+
+  private remember(sessionId: string, stored: StoredAttachment): void {
+    const pending = this.bySession.get(sessionId) ?? new Map();
+    pending.set(stored.id, stored);
+    this.bySession.set(sessionId, pending);
+  }
+}
+
+function stemOf(name: string): string | undefined {
+  const stem = basename(name).replace(/\.[^.]*$/, "");
+  return stem === "" ? undefined : stem;
+}
+
+function decodeURIComponentSafely(segment: string): string {
+  try {
+    return decodeURIComponent(segment);
+  } catch {
+    return "";
+  }
+}

package/packages/server/src/ClientConnection.ts

@@ -0,0 +1,277 @@
+import type { ServerWebSocket } from "bun";
+
+import {
+  isDurableEvent,
+  type ServerEvent,
+  type StreamEvent,
+} from "#protocol/ServerEvent";
+import type { SessionProjection } from "./SessionProjection";
+
+/** The half of `SessionStream` a connection needs: subscribe, and ask for everything missed. */
+export type AttachableStream = {
+  readonly subscribe: (listener: (event: ServerEvent) => void) => () => void;
+  readonly replay: (fromSeq: number) => Promise<readonly StreamEvent[]>;
+  readonly sessionId: string;
+};
+
+type SubagentWatch = {
+  readonly callId: string;
+  readonly projection: SessionProjection;
+  cursor: number;
+};
+
+const HIGH_WATER_MARK = 1 << 20;
+
+const FRAMES = new WeakMap<object, string>();
+
+export function frame(event: object): string {
+  const cached = FRAMES.get(event);
+  if (cached !== undefined) {
+    return cached;
+  }
+  const text = JSON.stringify(event);
+  FRAMES.set(event, text);
+  return text;
+}
+
+/** One WebSocket client, attached to at most one session; a lagging client is paused and re-derived on `drain`. */
+export class ClientConnection {
+  private readonly ws: ServerWebSocket<undefined>;
+  private stream: AttachableStream | undefined;
+  private unsubscribe: (() => void) | undefined;
+  private readonly gate: ServerEvent[] = [];
+  private gated = false;
+  private paused = false;
+  private cursor = 0;
+  private closed = false;
+  private reading = true;
+  private watch: SubagentWatch | undefined;
+
+  public constructor(ws: ServerWebSocket<undefined>) {
+    this.ws = ws;
+  }
+
+  public get sessionId(): string | undefined {
+    return this.stream?.sessionId;
+  }
+
+  /** Whether the reader behind this socket is looking at the session right now. */
+  public get attentive(): boolean {
+    return this.reading;
+  }
+
+  /** Highest durable `seq` this client has been handed. */
+  public get seq(): number {
+    return this.cursor;
+  }
+
+  /** The subagent this client is reading, if it is reading one. */
+  public get watchedCallId(): string | undefined {
+    return this.watch?.callId;
+  }
+
+  public async attach(
+    stream: AttachableStream,
+    fromSeq: number,
+    attentive = true
+  ): Promise<void> {
+    this.detach();
+    this.stream = stream;
+    this.cursor = fromSeq;
+    this.reading = attentive;
+    this.unsubscribe = stream.subscribe((event) => {
+      this.onStreamEvent(event);
+    });
+    await this.sync();
+  }
+
+  public setAttentive(attentive: boolean): void {
+    this.reading = attentive;
+  }
+
+  /** Read a subagent's log alongside the attached session; at most one watch at a time. */
+  public async watchSubagent(
+    callId: string,
+    projection: SessionProjection,
+    fromSeq: number
+  ): Promise<void> {
+    this.watch = { callId, projection, cursor: fromSeq };
+    await this.pumpWatch();
+  }
+
+  /** Closes only the named watch, so it cannot close the one that replaced it. */
+  public unwatchSubagent(callId: string): void {
+    if (this.watch?.callId === callId) {
+      this.watch = undefined;
+    }
+  }
+
+  /** Send something that belongs to no stream, e.g. a command response. */
+  public send(event: ServerEvent): void {
+    if (this.closed) {
+      return;
+    }
+    this.ws.send(frame(event), true);
+  }
+
+  public onDrain(): void {
+    if (!this.paused) {
+      return;
+    }
+    this.paused = false;
+    void this.sync();
+    this.flushWatch();
+  }
+
+  public detach(): void {
+    this.unsubscribe?.();
+    this.unsubscribe = undefined;
+    this.stream = undefined;
+    this.watch = undefined;
+  }
+
+  public close(): void {
+    this.closed = true;
+    this.detach();
+  }
+
+  // Gate live events during the read and reconcile by `seq`; ephemeral ones are already in the snapshot.
+  private async sync(): Promise<void> {
+    const stream = this.stream;
+    if (!stream || this.closed) {
+      return;
+    }
+    this.gated = true;
+    const replayed = await stream.replay(this.cursor);
+    if (this.stream !== stream || this.closed) {
+      return;
+    }
+    const queued = this.gate.filter(isDurableEvent);
+    this.gate.length = 0;
+    this.gated = false;
+    this.writeBatch([...replayed, ...queued]);
+  }
+
+  private onStreamEvent(event: ServerEvent): void {
+    if (this.touchesWatch(event)) {
+      void this.pumpWatch();
+    }
+    // Gate a batch unwrapped: the durable filter cannot see into an envelope.
+    if (event.type === "replay") {
+      if (this.gated) {
+        this.gate.push(...event.events);
+        return;
+      }
+      this.writeBatch(event.events);
+      return;
+    }
+    if (this.gated) {
+      this.gate.push(event);
+      return;
+    }
+    this.write(event);
+  }
+
+  // Sent whole or not at all: the cursor moves only once the frame is on the socket.
+  private writeBatch(events: readonly StreamEvent[]): void {
+    if (this.closed || this.paused) {
+      return;
+    }
+    // Filter against a cursor that moves inside the batch: replay and queue can name the same line.
+    let cursor = this.cursor;
+    const fresh: StreamEvent[] = [];
+    for (const event of events) {
+      if (isDurableEvent(event)) {
+        if (event.seq <= cursor) {
+          continue;
+        }
+        cursor = event.seq;
+      }
+      fresh.push(event);
+    }
+    if (fresh.length === 0) {
+      return;
+    }
+    if (this.deliver(JSON.stringify({ type: "replay", events: fresh }))) {
+      this.cursor = cursor;
+    }
+  }
+
+  private touchesWatch(event: ServerEvent): boolean {
+    const callId = this.watch?.callId;
+    if (callId === undefined) {
+      return false;
+    }
+    switch (event.type) {
+      case "tool_update":
+      case "tool_end":
+      case "tool_result":
+        return event.callId === callId;
+      case "replay":
+        return event.events.some((inner) => this.touchesWatch(inner));
+      default:
+        return false;
+    }
+  }
+
+  private async pumpWatch(): Promise<void> {
+    const watch = this.watch;
+    if (!watch || this.closed) {
+      return;
+    }
+    await watch.projection.drain();
+    // The watch can be closed or replaced during the read.
+    if (this.watch !== watch) {
+      return;
+    }
+    this.flushWatch();
+  }
+
+  private flushWatch(): void {
+    const watch = this.watch;
+    if (!watch || this.closed || this.paused) {
+      return;
+    }
+    const events = watch.projection.since(watch.cursor);
+    const last = events.at(-1);
+    if (!last) {
+      return;
+    }
+    const payload = JSON.stringify({
+      type: "subagent_events",
+      callId: watch.callId,
+      events,
+    });
+    if (this.deliver(payload)) {
+      watch.cursor = last.seq;
+    }
+  }
+
+  private deliver(payload: string): boolean {
+    if (this.closed || this.paused) {
+      return false;
+    }
+    const status = this.ws.send(payload, true);
+    if (status === 0) {
+      this.paused = true;
+      return false;
+    }
+    if (this.ws.getBufferedAmount() > HIGH_WATER_MARK) {
+      this.paused = true;
+    }
+    return true;
+  }
+
+  private write(event: ServerEvent): void {
+    if (this.closed || this.paused) {
+      return;
+    }
+    const durable = isDurableEvent(event);
+    if (durable && event.seq <= this.cursor) {
+      return;
+    }
+    if (this.deliver(frame(event)) && durable) {
+      this.cursor = event.seq;
+    }
+  }
+}

package/packages/server/src/ImageEndpoint.ts

@@ -0,0 +1,36 @@
+import { join } from "node:path";
+
+import { Images } from "#core/shared/Images";
+import { SpillCache } from "#core/shared/SpillCache";
+import { ImageRoute } from "#protocol/ImageRoute";
+import { IMMUTABLE, serveFile } from "./StaticClient";
+
+export type ImageEndpointDeps = {
+  /** Defaults to `~/.pim/cache`, where `read` spills the picture it showed the model. */
+  readonly root?: string;
+};
+
+/** `GET /image/<sha256>.<ext>`: the cached copy of a picture a tool view addresses, or 404 once the TTL sweep took it. */
+export class ImageEndpoint {
+  private readonly root: string;
+
+  public constructor(deps: ImageEndpointDeps = {}) {
+    this.root = deps.root ?? SpillCache.dir();
+  }
+
+  /** Whether this endpoint owns the path, so the gateway can route to it. */
+  public static owns(pathname: string): boolean {
+    return ImageRoute.owns(pathname);
+  }
+
+  public async handle(request: Request): Promise<Response> {
+    if (request.method !== "GET" && request.method !== "HEAD") {
+      return new Response("expected GET", { status: 405 });
+    }
+    // The name is a digest and a format we wrote, which admits no separator to escape the cache with.
+    const name = ImageRoute.nameOf(new URL(request.url).pathname);
+    return name === null
+      ? new Response("not found", { status: 404 })
+      : serveFile(join(this.root, `${Images.CACHE_PREFIX}${name}`), IMMUTABLE);
+  }
+}