pim-agent 0.9.0 → 0.11.0

package/packages/server/src/WsGateway.ts

@@ -7,16 +7,27 @@ import { Attachments } from "#core/attachments/Attachments";
 import type { PickerItem } from "#core/picker/PickerItem";
 import { Directories } from "#core/shared/Directories";
 import type { DirectoryListing } from "#core/shared/Directories";
+import { Git, type GitBranch } from "#core/shared/Git";
+import { GitMonitor, type GitRun } from "#core/shared/GitMonitor";
+import {
+  PiExtensions,
+  type ExtensionEntry,
+  type ExtensionScope,
+} from "#core/shared/PiExtensions";
+import { RepoDiff } from "#core/shared/RepoDiff";
 import { ReadCursors } from "#core/session/ReadCursors";
 import type { SessionHost } from "#core/session/SessionHost";
+import { SessionMeta } from "#core/session/SessionMeta";
 import type { SessionRegistry } from "#core/session/SessionRegistry";
 import { PimVersion } from "#core/shared/PimVersion";
 import { SubagentLogs } from "#core/shared/SubagentLogs";
 import type { UpdateOutcome } from "#core/shared/Updater";
 import type { Command } from "#protocol/Command";
-import { CLOSE_PROTOCOL_MISMATCH, PROTOCOL_VERSION } from "#protocol/Protocol";
+import type { ChangeList, FileDiff, FileLines } from "#protocol/Diff";
 import type {
   ModelView,
+  ProjectView,
+  SearchHitView,
   ServerEvent,
   SessionSummaryView,
 } from "#protocol/ServerEvent";
@@ -41,6 +52,14 @@ export type WsGatewayDeps = {
   readonly imagesRoot?: string;
   /** Where the read cursors live; defaults to `~/.pim/read.json`. */
   readonly readCursorsPath?: string;
+  /** Where the per-session and per-project overrides live; defaults to `~/.pim/sessions.json`. */
+  readonly sessionMetaPath?: string;
+  /** How often a session file is polled where `fs.watch` says nothing; the default is a second. */
+  readonly pollMs?: number;
+  /** Longest a dialog may hold an extension waiting; the default is `SessionStream`'s three minutes. */
+  readonly requestCeilingMs?: number;
+  /** How long a pending dialog outlives its last reader; the default is `SessionStream`'s fifteen seconds. */
+  readonly detachGraceMs?: number;
   /** The built web client; defaults to the bundle shipped beside this package. */
   readonly clientDir?: string;
   /** Runs the update a `reload` asks for, reporting each step as it starts. */
@@ -53,10 +72,22 @@ type Outcome = {
   readonly error?: string;
   readonly items?: readonly PickerItem[];
   readonly sessions?: readonly SessionSummaryView[];
+  readonly projects?: readonly ProjectView[];
+  readonly hits?: readonly SearchHitView[];
+  readonly dropped?: readonly string[];
+  readonly scanned?: number;
   readonly models?: readonly ModelView[];
   readonly thinkingLevels?: readonly string[];
+  readonly extensions?: readonly ExtensionEntry[];
   readonly directory?: DirectoryListing;
+  readonly branches?: readonly GitBranch[];
+  readonly commit?: { readonly sha: string };
+  readonly changes?: ChangeList;
+  readonly fileDiff?: FileDiff;
+  readonly fileLines?: FileLines;
   readonly restored?: readonly string[];
+  /** The message was dispatched as an extension command: no turn started, no entry written. */
+  readonly dispatched?: boolean;
   readonly after?: () => void;
 };
 
@@ -73,6 +104,15 @@ export const DEFAULT_HOSTNAME = "127.0.0.1";
 // Bun 1.3.14 never frees the slot of a server-closed socket, so an unbounded `Server.stop(true)` hangs.
 const STOP_GRACE_MS = 250;
 
+/**
+ * Pi's own dispatch parse, three lines of `_tryExecuteExtensionCommand`
+ * (`agent-session.js:954`); a pi upgrade re-verifies it.
+ */
+function commandNameOf(text: string): string {
+  const spaceIndex = text.indexOf(" ");
+  return spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex);
+}
+
 /** The transport half of pim-server: a WebSocket endpoint over the session runtime in `core`. */
 export class WsGateway {
   private readonly registry: SessionRegistry;
@@ -81,6 +121,7 @@ export class WsGateway {
   private readonly uploads: AttachmentEndpoint;
   private readonly images: ImageEndpoint;
   private readonly catalogue: SessionCatalogue;
+  private readonly meta: SessionMeta;
   private readonly client: StaticClient;
   private readonly streams = new Map<string, SessionStream>();
   private readonly opening = new Map<string, Promise<SessionStream>>();
@@ -89,22 +130,39 @@ export class WsGateway {
     ClientConnection
   >();
   private readonly reloader: Reloader;
+  private readonly git = new GitMonitor();
+  private readonly pollMs: number | undefined;
+  private readonly requestCeilingMs: number | undefined;
+  private readonly detachGraceMs: number | undefined;
+  private readonly busyRepos = new Set<string>();
   private versionsRead: Promise<readonly [string, string]> | undefined;
   private server: Server<undefined> | undefined;
 
   public constructor(deps: WsGatewayDeps) {
     this.registry = deps.registry;
+    // `SessionStream` is the sink, so the whole routing is a lookup: a host
+    // outlives the stream that speaks for it, and pi asks per call.
+    this.registry.setUi((host) =>
+      host.sessionId === undefined
+        ? undefined
+        : this.streams.get(host.sessionId)
+    );
     this.hostname = deps.hostname ?? DEFAULT_HOSTNAME;
     this.requestedPort = deps.port ?? DEFAULT_PORT;
+    this.pollMs = deps.pollMs;
+    this.requestCeilingMs = deps.requestCeilingMs;
+    this.detachGraceMs = deps.detachGraceMs;
     this.uploads = new AttachmentEndpoint(
       deps.attachmentsRoot === undefined ? {} : { root: deps.attachmentsRoot }
     );
     this.images = new ImageEndpoint(
       deps.imagesRoot === undefined ? {} : { root: deps.imagesRoot }
     );
+    this.meta = new SessionMeta(deps.sessionMetaPath);
     this.catalogue = new SessionCatalogue({
       registry: deps.registry,
       cursors: new ReadCursors(deps.readCursorsPath),
+      meta: this.meta,
       liveStatus: (sessionId) => this.streams.get(sessionId)?.host.status,
       liveSessionIds: () => this.streams.keys(),
       isBeingRead: (sessionId) => this.isBeingRead(sessionId),
@@ -145,7 +203,9 @@ export class WsGateway {
       fetch: (req, server) => {
         const { pathname } = new URL(req.url);
         if (pathname === "/health") {
-          return Response.json({ ok: true, protocolVersion: PROTOCOL_VERSION });
+          return this.versions().then(([pimVersion]) =>
+            Response.json({ ok: true, pimVersion })
+          );
         }
         if (AttachmentEndpoint.owns(pathname)) {
           return this.uploads.handle(req);
@@ -159,6 +219,8 @@ export class WsGateway {
       websocket: {
         backpressureLimit: 8 << 20,
         closeOnBackpressureLimit: false,
+        // Dedicated, so the window survives between frames: the repeated event envelope is most of a delta.
+        perMessageDeflate: { compress: "dedicated", decompress: "dedicated" },
         open: (ws) => {
           this.connections.set(ws, new ClientConnection(ws));
           this.syncWatches();
@@ -184,6 +246,8 @@ export class WsGateway {
     }
     this.streams.clear();
     this.opening.clear();
+    this.git.dispose();
+    this.busyRepos.clear();
     this.catalogue.watch(false);
     this.catalogue.clear();
     await this.catalogue.flush();
@@ -217,19 +281,6 @@ export class WsGateway {
       connection.send({ type: "error", message: "frame is not a command" });
       return;
     }
-    if (
-      command.type === "attach" &&
-      command.protocolVersion !== PROTOCOL_VERSION
-    ) {
-      connection.send({
-        type: "response",
-        id: command.id,
-        success: false,
-        error: `unsupported protocol version ${String(command.protocolVersion)}; this server speaks ${PROTOCOL_VERSION}`,
-      });
-      ws.close(CLOSE_PROTOCOL_MISMATCH, "protocol version mismatch");
-      return;
-    }
     try {
       const { error, after, ...answer } = await this.dispatch(
         connection,
@@ -260,24 +311,180 @@ export class WsGateway {
     switch (command.type) {
       case "attach":
         return await this.attach(connection, command);
+      case "attention": {
+        const regained = command.value && !connection.attentive;
+        connection.setAttentive(command.value);
+        const sessionId = connection.sessionId;
+        if (regained && sessionId !== undefined) {
+          await this.catalogue.markRead(sessionId);
+        }
+        return {};
+      }
       case "list_sessions":
-        return { sessions: await this.catalogue.list(command) };
+        return await this.catalogue.list(command);
+      case "search_sessions":
+        return await this.catalogue.search(command);
+      // The four below take a session this server may never have opened: a row
+      // is archived or renamed from the sidebar without being attached to, so
+      // none of them may reach for a stream. A sidecar write moves no session
+      // file and produces no `sessions_changed`, so the broadcast each ends
+      // with is the only word a listing already in a client's hands gets.
+      case "set_session_name": {
+        const name = await this.registry.setName(
+          command.sessionId,
+          command.value
+        );
+        this.broadcast({
+          type: "session_meta",
+          sessionId: command.sessionId,
+          name: name ?? null,
+        });
+        return {};
+      }
+      case "set_session_archived":
+        await this.meta.setArchived(command.sessionId, command.value);
+        this.broadcast({
+          type: "session_meta",
+          sessionId: command.sessionId,
+          archived: command.value,
+        });
+        return {};
+      case "set_session_unread":
+        await this.meta.setUnread(command.sessionId, command.value);
+        this.broadcast({
+          type: "session_meta",
+          sessionId: command.sessionId,
+          unread: command.value,
+        });
+        return {};
+      case "set_project_pinned":
+        await this.meta.setPinned(command.cwd, command.value);
+        this.broadcast({
+          type: "project_meta",
+          cwd: command.cwd,
+          pinned: command.value,
+        });
+        // The flag and the place it takes are one fact to a sidebar; a pin
+        // that arrived without a rank would sort by nothing until the next listing.
+        this.broadcast({ type: "pins_changed", order: await this.meta.pins() });
+        return {};
+      case "set_project_expanded":
+        await this.meta.setExpanded(command.cwd, command.value);
+        this.broadcast({
+          type: "project_meta",
+          cwd: command.cwd,
+          expanded: command.value,
+        });
+        return {};
+      case "set_project_label":
+        await this.meta.setLabel(command.cwd, command.value);
+        this.broadcast({
+          type: "project_meta",
+          cwd: command.cwd,
+          label: command.value,
+        });
+        return {};
+      case "set_pin_order":
+        await this.meta.setPinOrder(command.order);
+        this.broadcast({ type: "pins_changed", order: await this.meta.pins() });
+        return {};
       case "list_models":
         return {
           models: this.registry.models(),
           thinkingLevels:
             this.streamFor(connection)?.host.supportedThinkingLevels ?? [],
         };
+      case "list_extensions":
+        return {
+          extensions: await PiExtensions.list(this.scopeFor(connection)),
+        };
+      case "set_extension":
+        await PiExtensions.setEnabled(
+          command.extensionId,
+          command.value,
+          this.scopeFor(connection)
+        );
+        return {
+          after: () => {
+            this.rebuildAgents();
+          },
+        };
       case "list_dirs":
         return { directory: await Directories.list(command.path) };
+      case "create_dir":
+        await Directories.create(command.path);
+        return {};
+      case "refresh_git": {
+        const stream = this.requireStream(connection);
+        await stream.refreshGit(command.fetch === true);
+        return {};
+      }
+      case "list_branches":
+        return {
+          branches: await Git.listBranches(
+            this.requireStream(connection).host.cwd
+          ),
+        };
+      // Reading the repository never moves it, so neither of these takes the `repoBusy` refusal.
+      case "list_changes":
+        return {
+          changes: await RepoDiff.listChanges(
+            this.requireStream(connection).host.cwd,
+            command.base,
+            this.git
+          ),
+        };
+      case "file_diff":
+        return {
+          fileDiff: await RepoDiff.fileDiff(
+            this.requireStream(connection).host.cwd,
+            command.base,
+            command.path,
+            this.git,
+            command.context
+          ),
+        };
+      case "read_lines":
+        return {
+          fileLines: await RepoDiff.readLines(
+            this.requireStream(connection).host.cwd,
+            command.base,
+            command.path,
+            command.spans,
+            this.git
+          ),
+        };
+      case "checkout": {
+        const stream = this.requireStream(connection);
+        return await this.runGit(stream, true, (cwd) =>
+          Git.checkout(cwd, command.branch)
+        );
+      }
+      case "commit":
+        return await this.commit(this.requireStream(connection), command);
+      case "pull":
+      case "push":
+        return await this.runGit(
+          this.requireStream(connection),
+          command.type === "pull",
+          command.type === "pull" ? Git.pull : Git.push
+        );
       case "unwatch_subagent":
         connection.unwatchSubagent(command.callId);
         return {};
       case "reload":
         return this.reload(command.force === true);
       case "user_message":
-        this.promptWithAttachments(this.requireStream(connection), command);
-        return {};
+        return await this.promptWithAttachments(
+          this.requireStream(connection),
+          command
+        );
+      case "ui_response": {
+        const stream = this.requireStream(connection);
+        return stream.answer(command.requestId, command)
+          ? {}
+          : { error: `no request ${command.requestId} is waiting` };
+      }
       case "cancel": {
         const { cancelled, restored } =
           await this.requireStream(connection).host.cancel();
@@ -288,6 +495,7 @@ export class WsGateway {
       case "set_cwd": {
         const stream = this.requireStream(connection);
         const result = await stream.host.setCwd(command.value);
+        stream.syncGit();
         stream.push(stream.sessionState());
         if (!result.ok) {
           return { error: result.error };
@@ -342,6 +550,27 @@ export class WsGateway {
     return sessionId ? this.streams.get(sessionId) : undefined;
   }
 
+  /** A session's directory decides which project extensions are listed; a connection without one reads this process's. */
+  private scopeFor(connection: ClientConnection): ExtensionScope {
+    return {
+      cwd: this.streamFor(connection)?.host.cwd ?? process.cwd(),
+      agentDir: this.registry.agentDir,
+    };
+  }
+
+  private rebuildAgents(): void {
+    for (const stream of this.streams.values()) {
+      const { host } = stream;
+      // Through the turn queue: detaching the agent under a running turn kills it.
+      void host
+        .serialize(() => host.invalidate())
+        .catch((err: unknown) => {
+          this.failed(stream, err, "extension reload");
+        });
+    }
+    this.broadcast({ type: "extensions_changed" });
+  }
+
   private requireStream(connection: ClientConnection): SessionStream {
     const stream = this.streamFor(connection);
     if (!stream) {
@@ -365,7 +594,8 @@ export class WsGateway {
     if (path === null) {
       return { error: `malformed call id: ${command.callId}` };
     }
-    if (!(await Bun.file(path).exists())) {
+    // A running call may not have written its log yet; the watch waits for it.
+    if (!stream.isRunning(command.callId) && !(await Bun.file(path).exists())) {
       return { error: `no subagent log for call ${command.callId}` };
     }
     await connection.watchSubagent(
@@ -380,6 +610,7 @@ export class WsGateway {
     connection: ClientConnection,
     command: Command & { readonly type: "attach" }
   ): Promise<Outcome> {
+    const attentive = command.attentive !== false;
     const like =
       command.like === undefined
         ? undefined
@@ -396,27 +627,98 @@ export class WsGateway {
     ]);
     connection.send({
       type: "attached",
-      protocolVersion: PROTOCOL_VERSION,
       sessionId: stream.sessionId,
       cwd: stream.host.cwd,
       head: await stream.refresh(),
       pimVersion,
       piVersion,
     });
-    await connection.attach(stream, command.fromSeq);
+    await connection.attach(stream, command.fromSeq, attentive);
     this.syncWatches();
-    await this.catalogue.markRead(stream.sessionId);
+    if (attentive) {
+      await this.catalogue.markRead(stream.sessionId);
+    }
     return {};
   }
 
-  /** A watch costs a poll, so only sessions someone is reading get one — and the catalogue only while someone is connected. */
+  /** A watch costs a poll, so only attached sessions get one — and the catalogue only while someone is connected. */
   private syncWatches(): void {
     for (const [sessionId, stream] of this.streams) {
-      stream.watchFiles(this.isBeingRead(sessionId));
+      // Attachment, not attention: a hidden tab that lost its watches would stop hearing about the session.
+      stream.watchFiles(this.isAttached(sessionId));
     }
     this.catalogue.watch(this.connections.size > 0);
   }
 
+  /**
+   * An operation that moves the working tree cannot run under an agent that
+   * may be halfway through an edit; `push` leaves the tree alone, so it can.
+   */
+  private async runGit<T = never>(
+    stream: SessionStream,
+    movesTree: boolean,
+    operation: (cwd: string) => Promise<GitRun<T>>
+  ): Promise<Outcome & { readonly value?: T }> {
+    const cwd = stream.host.cwd;
+    if (movesTree && this.repoBusy(cwd)) {
+      return {
+        error: `an agent is working in ${cwd}; wait for its turn to end`,
+      };
+    }
+    const result = await this.git.run<T>(cwd, () => operation(cwd));
+    if (!result.ok) {
+      return { error: result.error };
+    }
+    return result.value === undefined ? {} : { value: result.value };
+  }
+
+  private async commit(
+    stream: SessionStream,
+    command: Command & { readonly type: "commit" }
+  ): Promise<Outcome> {
+    const { value, ...outcome } = await this.runGit<string>(
+      stream,
+      true,
+      async (cwd) => {
+        const result = await Git.commit(cwd, {
+          message: command.message,
+          paths: command.paths,
+        });
+        return result.ok
+          ? { ok: true, value: result.sha }
+          : { ok: false, error: result.error };
+      }
+    );
+    return value === undefined ? outcome : { commit: { sha: value } };
+  }
+
+  private repoBusy(cwd: string): boolean {
+    for (const stream of this.streams.values()) {
+      if (stream.host.cwd === cwd && stream.host.status !== "idle") {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /** A turn starting anywhere freezes the branch menu of every session beside it, which only their own state can say. */
+  private syncRepoBusy(cwd: string, source: string): void {
+    const busy = this.repoBusy(cwd);
+    if (this.busyRepos.has(cwd) === busy) {
+      return;
+    }
+    if (busy) {
+      this.busyRepos.add(cwd);
+    } else {
+      this.busyRepos.delete(cwd);
+    }
+    for (const [sessionId, stream] of this.streams) {
+      if (sessionId !== source && stream.host.cwd === cwd) {
+        stream.push(stream.sessionState());
+      }
+    }
+  }
+
   private reload(force: boolean): Outcome {
     const busy = [...this.streams.values()]
       .filter((stream) => stream.host.status !== "idle")
@@ -439,13 +741,20 @@ export class WsGateway {
     return this.versionsRead;
   }
 
+  private isAttached(sessionId: string): boolean {
+    return this.connections
+      .values()
+      .some((connection) => connection.sessionId === sessionId);
+  }
+
+  // The TUI never attaches, so reading a session there clears no dot here and a turn it runs trips no mark.
   private isBeingRead(sessionId: string): boolean {
-    for (const connection of this.connections.values()) {
-      if (connection.sessionId === sessionId) {
-        return true;
-      }
-    }
-    return false;
+    return this.connections
+      .values()
+      .some(
+        (connection) =>
+          connection.sessionId === sessionId && connection.attentive
+      );
   }
 
   private broadcast(event: ServerEvent): void {
@@ -498,29 +807,91 @@ export class WsGateway {
     if (existing) {
       return existing;
     }
-    const stream = new SessionStream(id, host, path);
+    const stream = new SessionStream(id, host, path, {
+      git: this.git,
+      repoBusy: () => this.repoBusy(host.cwd),
+      pollMs: this.pollMs,
+      requestCeilingMs: this.requestCeilingMs,
+      detachGraceMs: this.detachGraceMs,
+    });
     stream.start();
     this.catalogue.track(id, host.status);
-    stream.subscribe((event) => {
+    // Observed rather than subscribed: this never leaves, and a reader that
+    // never leaves is a session that is never detached.
+    stream.observe((event) => {
       if (event.type === "session_state") {
         this.catalogue.onStatus(id, event.status);
+        this.syncRepoBusy(stream.host.cwd, id);
       }
     });
     this.streams.set(id, stream);
     return stream;
   }
 
-  private promptWithAttachments(
+  private async promptWithAttachments(
     stream: SessionStream,
     command: Command & { readonly type: "user_message" }
-  ): void {
+  ): Promise<Outcome> {
+    // Trimmed once, here: pi dispatches on the text it is handed after the
+    // fallthrough trims it, so a client that does not trim would otherwise
+    // slip a command past this check and into pi's own.
+    const text = command.text.trim();
+    // Read before anything is rendered: what a command is, is its own text.
+    const dispatched = await this.dispatchCommand(stream, text);
     const taken = this.uploads.take(
       stream.sessionId,
       (command.attachments ?? []).map((ref) => ref.id)
     );
+    // A slash command carries no pictures; the uploads are still taken, so the
+    // endpoint stops holding their bytes for the rest of the session.
+    if (dispatched) {
+      return { dispatched: true };
+    }
     const { lines, images } = Attachments.render(taken);
-    const text = [command.text, ...lines].filter(Boolean).join("\n\n").trim();
-    this.prompt(stream, text, images);
+    this.prompt(stream, [text, ...lines].filter(Boolean).join("\n\n"), images);
+    return {};
+  }
+
+  /**
+   * Whether `text` names an extension command, having started it if it does.
+   * Only the decision is awaited: the handler runs on beyond the ack, exactly
+   * as a turn does.
+   */
+  private async dispatchCommand(
+    stream: SessionStream,
+    text: string
+  ): Promise<boolean> {
+    const host = stream.host;
+    if (!text.startsWith("/")) {
+      return false;
+    }
+    // The picker offers extension commands only once an agent exists, so this
+    // is someone typing from memory; cold is idle, so nothing is queued behind it.
+    const agent = host.agentSession ?? (await host.ensureAgent());
+    const name = commandNameOf(text);
+    if (agent.extensionRunner.getCommand(name) === undefined) {
+      return false;
+    }
+    void stream
+      .dispatch(`/${name}`, async () => {
+        // Prompted whole, never merged into the queue: a merged string starts
+        // with the earlier message, and pi reads the leading `/` before it
+        // looks at `isStreaming` — so the turn we are holding dispatches it.
+        if (host.isStreaming) {
+          await agent.prompt(text, {
+            streamingBehavior: "steer",
+            source: "rpc",
+          });
+          return;
+        }
+        await host.run(async (session) => {
+          await session.prompt(text, { source: "rpc" });
+        });
+      })
+      .catch((err: unknown) => {
+        this.failed(stream, err, "command");
+      });
+    return true;
   }
 
   // Merge into the queued message: pi holds at most one, so a second queued send would be unreachable.
@@ -532,13 +903,6 @@ export class WsGateway {
     const attached = images.length === 0 ? {} : { images: [...images] };
     const host = stream.host;
     const agent = host.agentSession;
-    const failed = (err: unknown, what: string): void => {
-      console.error(`[gateway] ${what} failed:`, err);
-      stream.push({
-        type: "error",
-        message: (err as Error).message || String(err),
-      });
-    };
     if (agent && host.isStreaming) {
       const queued = [...host.takeBack(), text].join("\n\n");
       void agent
@@ -548,7 +912,7 @@ export class WsGateway {
           ...attached,
         })
         .catch((err: unknown) => {
-          failed(err, "steer");
+          this.failed(stream, err, "steer");
         });
       return;
     }
@@ -557,7 +921,15 @@ export class WsGateway {
         await session.prompt(text, { source: "rpc", ...attached });
       })
       .catch((err: unknown) => {
-        failed(err, "turn");
+        this.failed(stream, err, "turn");
       });
   }
+
+  private failed(stream: SessionStream, err: unknown, what: string): void {
+    console.error(`[gateway] ${what} failed:`, err);
+    stream.push({
+      type: "error",
+      message: (err as Error).message || String(err),
+    });
+  }
 }