ompchamber 0.3.0

package/src/server/lib/omp/rpc/process.ts

@@ -0,0 +1,319 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Process + protocol layer for `omp --mode rpc-ui` (NDJSON over stdio).
+ * Protocol v1: commands `{id, type, ...}` on stdin; `{type:"response", id, ...}`
+ * plus interleaved event frames on stdout. omp announces readiness with a
+ * `{type:"ready"}` frame before accepting commands. When readiness advertises
+ * protocol v2, callers negotiate it before sending normal commands; oversized
+ * logical frames are then carried as bounded `rpc_chunk` sequences.
+ *
+ * Faithful port of omp-web/lib/omp/rpc-process.ts.
+ */
+
+import type { Subprocess } from 'bun';
+import { resolveOmpBin } from '@/server/lib/omp/core/cli';
+import { RpcFrameDecoder, encodeRpcFrames, type RpcFrameRecord, type RpcProtocolVersion } from '@/shared/lib/omp/rpc/frame';
+import { killProcessTree } from '@/server/lib/omp/rpc/kill-tree';
+import { readLines } from '@/server/lib/omp/rpc/lines';
+import { RpcCommandError, RpcCommandTimeoutError, STDERR_TAIL_LIMIT, sanitizeProjectCommandEnvironment, type PendingCommand, type RpcFrame, type RpcProcessOptions, type RpcResponseFrame } from '@/server/lib/omp/rpc/process-helpers';
+
+export { RpcCommandError, RpcCommandTimeoutError } from '@/server/lib/omp/rpc/process-helpers';
+export type { RpcFrame, RpcProcessOptions, RpcResponseFrame } from '@/server/lib/omp/rpc/process-helpers';
+
+export class RpcProcess {
+  readonly cwd: string;
+  private child!: Subprocess<"pipe", "pipe", "pipe">;
+  private readonly pending = new Map<string, PendingCommand>();
+  private readonly frameListeners = new Set<(frame: RpcFrame) => void>();
+  private readyPromise: Promise<RpcFrame>;
+  private nextId = 1;
+  private stderrTail = '';
+  private exited = false;
+  private exitInfo: { code: number | null; signal: NodeJS.Signals | null } | null = null;
+  private protocolVersion: RpcProtocolVersion = 1;
+  private nextChunkId = 1;
+  // Serializes physical stdin writes: a v2 logical frame can span multiple
+  // `rpc_chunk` records (>1 MiB payloads), and two frames written concurrently
+  // would interleave their chunk sequences on stdin, which RpcFrameDecoder
+  // rejects. Each logical frame is enqueued whole.
+  private writeQueue: Promise<void> = Promise.resolve();
+
+  constructor(options: RpcProcessOptions) {
+    const resolveBin = options.dependencies?.resolveOmpBin ?? resolveOmpBin;
+    const bin = resolveBin();
+    if (!bin) {
+      throw new Error('omp binary not found. Install oh-my-pi or set OMP_WEB_OMP_BIN.');
+    }
+    this.cwd = options.cwd;
+    if (options.onFrame) this.frameListeners.add(options.onFrame);
+
+    const args = ['--mode', 'rpc-ui', '--cwd', options.cwd, ...(options.extraArgs ?? [])];
+    // Strip named profiles (OMP_PROFILE/PI_PROFILE) so the child resolves the
+    // same default agent dir. An explicit options.env entry still wins.
+    const childEnv = sanitizeProjectCommandEnvironment({ ...Bun.env, ...options.env });
+    if (options.env?.OMP_PROFILE === undefined) delete childEnv.OMP_PROFILE;
+    if (options.env?.PI_PROFILE === undefined) delete childEnv.PI_PROFILE;
+    this.child = Bun.spawn({
+      cmd: [bin, ...args],
+      cwd: options.cwd,
+      env: childEnv,
+      stdin: 'pipe',
+      stdout: 'pipe',
+      stderr: 'pipe',
+      windowsHide: true,
+      // On POSIX, omp launches grandchildren (LSP servers, extension subprocesses). Run the
+      // child in its own process group so dispose() can SIGTERM/SIGKILL the whole
+      // tree — otherwise a crashed omp would orphan its LSP children as zombies.
+      // Windows uses taskkill /t instead, so detaching would only create a console.
+      detached: process.platform !== 'win32',
+    });
+
+    let resolveReady: (frame: RpcFrame) => void;
+    let rejectReady: (error: Error) => void;
+    this.readyPromise = new Promise<RpcFrame>((resolve, reject) => {
+      resolveReady = resolve;
+      rejectReady = reject;
+    });
+    // waitReady() is optional for callers; avoid unhandled-rejection noise when
+    // the process dies before anyone awaited readiness.
+    this.readyPromise.catch(() => {});
+
+    const decoder = new RpcFrameDecoder();
+    void readLines(this.child.stdout, (line) => {
+      const trimmed = line.trim();
+      if (!trimmed) return;
+      let parsed: unknown;
+      try {
+        parsed = JSON.parse(trimmed);
+      } catch {
+        // omp guards stdout in RPC mode, but never let a stray line kill the reader.
+        return;
+      }
+      let frame: RpcFrameRecord;
+      try {
+        if (parsed && typeof parsed === 'object' && (parsed as { type?: unknown }).type === 'rpc_chunk' && this.protocolVersion !== 2) {
+          throw new Error('RPC chunk received before protocol negotiation');
+        }
+        const decoded = decoder.push(parsed);
+        if (!decoded) return;
+        frame = decoded;
+      } catch (error) {
+        this.stderrTail = (this.stderrTail + `\nRPC protocol error: ${error instanceof Error ? error.message : String(error)}`).slice(-STDERR_TAIL_LIMIT);
+        void this.dispose(0);
+        return;
+      }
+      if (frame.type === 'ready') {
+        resolveReady(frame);
+        return;
+      }
+      if (frame.type === 'response') {
+        this.handleResponse(frame as unknown as RpcResponseFrame);
+        return;
+      }
+      for (const listener of this.frameListeners) {
+        try {
+          listener(frame);
+        } catch {
+          // Listener bugs must not break the protocol reader.
+        }
+      }
+    });
+
+    void (async () => {
+      await readLines(this.child.stderr, (line) => {
+        this.stderrTail = (this.stderrTail + line + '\n').slice(-STDERR_TAIL_LIMIT);
+      });
+    })();
+
+    const finalize = (code: number | null, signal: NodeJS.Signals | null) => {
+      if (this.exited) return;
+      this.exited = true;
+      this.exitInfo = { code, signal };
+      const exitError = new Error(
+        `omp exited (code ${code ?? 'null'}, signal ${signal ?? 'none'})${this.stderrTail ? `: ${this.stderrTail.slice(-500)}` : ''}`,
+      );
+      rejectReady(exitError);
+      for (const [, entry] of this.pending) {
+        if (entry.timer) clearTimeout(entry.timer);
+        entry.reject(exitError);
+      }
+      this.pending.clear();
+      options.onExit?.({ code, signal, stderrTail: this.stderrTail });
+    };
+    void this.child.exited.then((code) => {
+      finalize(code, this.child.signalCode);
+    });
+  }
+
+  get isAlive(): boolean {
+    return !this.exited;
+  }
+
+  /** OS pid of the spawned omp process (undefined before spawn settles). Used
+   * by the browser viewer to map this process to its owned browser targets. */
+  get exitDetails(): { code: number | null; signal: NodeJS.Signals | null; stderrTail: string } | null {
+    return this.exitInfo ? { ...this.exitInfo, stderrTail: this.stderrTail } : null;
+  }
+
+  /** OS pid of the spawned omp process, undefined once it is gone. Used by the
+   * browser viewer to map this process to its owned browser targets. */
+  get pid(): number | undefined {
+    return this.child.pid ?? undefined;
+  }
+
+  /** Resolves with the `ready` frame; rejects if the process dies first or the
+   * timeout elapses. omp startup can take a few seconds (extensions, LSP). */
+  waitReady(timeoutMs = 60_000): Promise<RpcFrame> {
+    const timeout = new Promise<never>((_, reject) => {
+      const timer = setTimeout(() => reject(new Error(`omp RPC ready timeout after ${timeoutMs}ms`)), timeoutMs);
+      timer.unref?.();
+      this.readyPromise.finally(() => clearTimeout(timer)).catch(() => {});
+    });
+    return Promise.race([this.readyPromise, timeout]);
+  }
+
+  /** Enables bounded protocol-v2 framing when the ready frame advertises it. */
+  async negotiateProtocol(ready: RpcFrame): Promise<RpcProtocolVersion> {
+    const supported = Array.isArray(ready.supportedProtocolVersions) ? ready.supportedProtocolVersions : [];
+    if (!supported.includes(2)) return this.protocolVersion;
+    const response = await this.sendCommand<{ protocolVersion?: unknown }>({ type: 'negotiate_protocol', protocolVersion: 2 });
+    if (response?.protocolVersion !== 2) throw new Error('OMP rejected RPC protocol v2 negotiation');
+    this.protocolVersion = 2;
+    return this.protocolVersion;
+  }
+
+  onFrame(listener: (frame: RpcFrame) => void): () => void {
+    this.frameListeners.add(listener);
+    return () => this.frameListeners.delete(listener);
+  }
+
+  /** Send a command and await its response `data`. A failed response rejects
+   * with RpcCommandError. No timeout by default — some commands (login,
+   * long prompts via bash) legitimately take minutes, and the session wrapper
+   * reclaims wedged children via idle-kill and dispose(). Callers that want a
+   * cap pass `timeoutMs` (>0); when set, the timer is unref'd so it never
+   * keeps the event loop alive on its own. */
+  sendCommand<T = unknown>(command: { type: string; [key: string]: unknown }, timeoutMs?: number): Promise<T> {
+    if (this.exited) {
+      return Promise.reject(new Error('omp RPC process has exited'));
+    }
+    const id = `w${this.nextId++}`;
+    return new Promise<T>((resolve, reject) => {
+      const entry: PendingCommand = {
+        command: command.type,
+        resolve: resolve as (data: unknown) => void,
+        reject,
+      };
+      if (timeoutMs && timeoutMs > 0) {
+        entry.timer = setTimeout(() => {
+          // Only reject if this exact entry is still pending — a reused id or a
+          // response that landed between the timer firing and this callback must
+          // not spuriously reject a different command.
+          if (this.pending.get(id) === entry) {
+            this.pending.delete(id);
+            reject(new RpcCommandTimeoutError(command.type, timeoutMs));
+          }
+        }, timeoutMs);
+        // A pending command timer must never keep the event loop alive on its own
+        // (it would block graceful shutdown when omp has stopped answering).
+        entry.timer.unref?.();
+      }
+      this.pending.set(id, entry);
+      this.writeFrame({ ...command, id }, (error) => {
+        if (error) {
+          const pending = this.pending.get(id);
+          if (pending) {
+            this.pending.delete(id);
+            if (pending.timer) clearTimeout(pending.timer);
+            reject(error);
+          }
+        }
+      });
+    });
+  }
+
+  /** Fire-and-forget frame write (extension_ui_response, host_tool_result). */
+  sendFrame(frame: { type: string; [key: string]: unknown }): void {
+    if (this.exited) return;
+    this.writeFrame(frame, () => {});
+  }
+
+  private writeFrame(frame: RpcFrame, callback: (error?: Error | null) => void): void {
+    let lines: string[];
+    try {
+      lines = encodeRpcFrames(frame, this.protocolVersion, `web-${this.nextChunkId++}`);
+    } catch (error) {
+      callback(error instanceof Error ? error : new Error(String(error)));
+      return;
+    }
+    // Enqueue the entire encoded logical frame; the next frame's physical
+    // records only start after this frame has been fully handed to the stdin
+    // sink, so two frames can never interleave their chunk sequences.
+    this.writeQueue = this.writeQueue.then(
+      () => new Promise<void>((resolve) => {
+        if (this.exited || !this.child.stdin) {
+          callback(new Error('RPC process is not running'));
+          resolve();
+          return;
+        }
+        try {
+          for (const line of lines) this.child.stdin.write(line);
+          callback(null);
+        } catch (error) {
+          callback(error instanceof Error ? error : new Error(String(error)));
+        } finally {
+          resolve();
+        }
+      }),
+    );
+  }
+
+  private handleResponse(response: RpcResponseFrame): void {
+    const id = response.id;
+    const entry = id ? this.pending.get(id) : undefined;
+    if (!entry || !id) {
+      // Unsolicited response (or a command we already timed out) — surface to
+      // frame listeners so nothing is silently dropped.
+      for (const listener of this.frameListeners) {
+        try {
+          listener(response as unknown as RpcFrame);
+        } catch {}
+      }
+      return;
+    }
+    this.pending.delete(id);
+    if (entry.timer) clearTimeout(entry.timer);
+    if (response.success) {
+      entry.resolve(response.data);
+    } else {
+      entry.reject(new RpcCommandError(response.command, response.error ?? 'RPC command failed', response.code));
+    }
+  }
+
+  /** Graceful shutdown: close stdin (omp exits on EOF), escalate to SIGTERM
+   * then SIGKILL on the whole process group. Resolves once the process has
+   * exited. Safe to call during server teardown — escalation timers are
+   * unref'd so they never keep the event loop alive on their own. */
+  async dispose(gracePeriodMs = 5_000): Promise<void> {
+    if (this.exited) return;
+    const exited = this.child.exited.then(() => undefined);
+    try {
+      this.child.stdin.end();
+    } catch {}
+    const timer = setTimeout(() => {
+      if (!this.exited) killProcessTree(this.child.pid, false, 'SIGTERM');
+    }, gracePeriodMs);
+    const killTimer = setTimeout(() => {
+      if (!this.exited) killProcessTree(this.child.pid, true, 'SIGKILL');
+    }, gracePeriodMs * 2);
+    timer.unref?.();
+    killTimer.unref?.();
+    await exited;
+    clearTimeout(timer);
+    clearTimeout(killTimer);
+  }
+}

package/src/server/lib/omp/rpc/session-commands.ts

@@ -0,0 +1,218 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Command dispatch for a live omp session. Extracted from AgentSessionWrapper
+ * (rpc-manager.ts) so that file stays under the repo's per-file size ceiling.
+ * The host object exposes the wrapper's runtime state; behavior is unchanged.
+ */
+
+import { RpcCommandTimeoutError, type RpcProcess } from '@/server/lib/omp/rpc/process';
+import { AWAITING_AGENT_START_TIMEOUT_MS, GET_STATE_TIMEOUT_MS, IMAGE_BEARING_COMMANDS, PASSTHROUGH_COMMANDS, PROMPT_ACK_TIMEOUT_MS, RESTARTING_MESSAGE, WebRpcError, toImageContents, type AgentEvent, type RpcSessionState, validateAgentImages } from '@/server/lib/omp/rpc/constants';
+import { clearSessionFileCaches } from '@/server/lib/omp/session/files';
+import { notifyRunningChange } from '@/server/lib/omp/rpc/session-registry';
+import { markStreamStatus } from '@/shared/lib/omp/session/stream-state.server';
+import { buildWebState, type WebStateHost } from '@/server/lib/omp/rpc/web-state';
+
+/** Runtime surface AgentSessionWrapper exposes to the command dispatcher. */
+export interface SessionCommandHost extends WebStateHost {
+  restarting: boolean;
+  proc: RpcProcess;
+  isAlive(): boolean;
+  /** Real omp session id (empty before the first get_state). */
+  sessionId: string;
+  emit(event: AgentEvent): void;
+  resetIdleTimer(force?: boolean): void;
+  /** Forget a pending ask/approval dialog once its response is sent. */
+  resolvePendingUiDialog(id: string): void;
+  withFinalRunningNotification<T>(operation: () => Promise<T>): Promise<T>;
+  destroyAndWait(): Promise<void>;
+}
+
+export async function dispatchSessionCommand(host: SessionCommandHost, command: Record<string, unknown>): Promise<unknown> {
+  if (host.restarting) throw new WebRpcError(RESTARTING_MESSAGE, 'session_restarting');
+  if (!host.isAlive()) throw new Error('Session is no longer running');
+  host.resetIdleTimer();
+  const type = command.type as string;
+
+  if (IMAGE_BEARING_COMMANDS.has(type)) {
+    const imageError = validateAgentImages(command.images);
+    if (imageError) throw new Error(imageError);
+  }
+
+  switch (type) {
+    case 'prompt': {
+      if (host.bashRunning) {
+        throw new Error('Cannot send a prompt while a shell command is running');
+      }
+      const streamingBehavior = command.streamingBehavior as 'steer' | 'followUp' | undefined;
+      if (!streamingBehavior) {
+        host.promptRunning = true;
+        host.promptDispatchPendingCount += 1;
+        host.awaitingAgentStart = false;
+        host.awaitingAgentStartDeadline = 0;
+        host.continuationGraceUntil = 0;
+        notifyRunningChange();
+      }
+      try {
+        const ack = await host.proc.sendCommand<{ agentInvoked?: boolean } | undefined>({
+          type: 'prompt',
+          message: command.message as string,
+          ...(toImageContents(command.images) ? { images: toImageContents(command.images) } : {}),
+          ...(streamingBehavior ? { streamingBehavior } : {}),
+        }, PROMPT_ACK_TIMEOUT_MS);
+        if (ack?.agentInvoked === false && !streamingBehavior) {
+          host.promptRunning = false;
+          host.awaitingAgentStart = false;
+          host.awaitingAgentStartDeadline = 0;
+          host.emit({ type: 'prompt_result', agentInvoked: false });
+          notifyRunningChange();
+        } else if (!streamingBehavior && ack?.agentInvoked !== false) {
+          host.awaitingAgentStart = true;
+          host.awaitingAgentStartDeadline = Date.now() + AWAITING_AGENT_START_TIMEOUT_MS;
+        }
+      } catch (error) {
+        host.promptRunning = false;
+        host.awaitingAgentStart = false;
+        host.awaitingAgentStartDeadline = 0;
+        notifyRunningChange();
+        if (error instanceof RpcCommandTimeoutError) {
+          await host.destroyAndWait();
+          throw new WebRpcError('The OMP session stopped responding and was reset.', 'session_unresponsive');
+        }
+        throw error;
+      } finally {
+        if (!streamingBehavior) {
+          host.promptDispatchPendingCount = Math.max(0, host.promptDispatchPendingCount - 1);
+        }
+      }
+      return null;
+    }
+
+    case 'abort':
+      await host.withFinalRunningNotification(async () => {
+        await host.proc.sendCommand({ type: 'abort' });
+        host.promptRunning = false;
+        host.awaitingAgentStart = false;
+        host.awaitingAgentStartDeadline = 0;
+        host.continuationGraceUntil = 0;
+        if (host.sessionId) void markStreamStatus(host.sessionId, 'abort');
+      });
+      return null;
+
+    case 'get_state': {
+      try {
+        const state = await host.proc.sendCommand<RpcSessionState>({ type: 'get_state' }, GET_STATE_TIMEOUT_MS);
+        return buildWebState(host, state);
+      } catch (error) {
+        if (error instanceof RpcCommandTimeoutError) {
+          await host.destroyAndWait();
+          throw new WebRpcError('The OMP session stopped responding and was reset.', 'session_unresponsive');
+        }
+        throw error;
+      }
+    }
+
+    case 'set_model': {
+      const { provider, modelId } = command as { provider: string; modelId: string };
+      const model = await host.proc.sendCommand<{ id: string; provider: string }>({ type: 'set_model', provider, modelId });
+      return { id: model.id, provider: model.provider };
+    }
+
+    case 'set_fast_mode': {
+      const enabled = command.enabled === true;
+      const result = await host.proc.sendCommand<{ enabled?: boolean; active?: boolean }>({ type: 'set_fast_mode', enabled });
+      host.fastModeEnabled = result?.enabled ?? enabled;
+      return { enabled: host.fastModeEnabled, active: result?.active ?? false };
+    }
+
+    case 'compact': {
+      try {
+        return await host.withFinalRunningNotification(async () => {
+          host.compacting = true;
+          notifyRunningChange();
+          try {
+            const result = await host.proc.sendCommand<{ summary?: string; tokensBefore?: number; estimatedTokensAfter?: number }>({
+              type: 'compact',
+              ...(command.customInstructions ? { customInstructions: command.customInstructions } : {}),
+            });
+            return result;
+          } finally {
+            host.compacting = false;
+          }
+        });
+      } finally {
+        clearSessionFileCaches();
+      }
+    }
+
+    case 'abort_compaction':
+      await host.withFinalRunningNotification(() => host.proc.sendCommand({ type: 'abort' }));
+      return null;
+
+    case 'set_session_name': {
+      const name = (command.name as string | undefined)?.trim();
+      if (!name) throw new Error('Session name cannot be empty');
+      await host.proc.sendCommand({ type: 'set_session_name', name });
+      clearSessionFileCaches();
+      return null;
+    }
+
+    case 'get_commands': {
+      const data = await host.proc.sendCommand<{ commands: unknown[] }>({
+        type: 'get_available_commands',
+      });
+      return data;
+    }
+
+    case 'bash': {
+      if (host.isRunning()) {
+        throw new Error('Cannot run a shell command while the session is busy');
+      }
+      host.bashRunning = true;
+      notifyRunningChange();
+      try {
+        return await host.proc.sendCommand<{ output?: string; exitCode?: number }>({ type: 'bash', command: command.command as string });
+      } finally {
+        host.bashRunning = false;
+        clearSessionFileCaches();
+        notifyRunningChange();
+      }
+    }
+
+    case 'extension_ui_response': {
+      // Fire-and-forget: omp answers ask/approval dialogs without a response
+      // frame, so a request/response round-trip would time out.
+      const { id, ...rest } = command as { id: string; [key: string]: unknown };
+      if (!id) throw new Error('extension_ui_response requires an id');
+      host.proc.sendFrame({ type: 'extension_ui_response', id, ...rest });
+      // Answered — stop offering it to clients that attach later.
+      host.resolvePendingUiDialog(id);
+      return null;
+    }
+
+    case 'steer':
+    case 'follow_up': {
+      if (!host.isRunning()) {
+        throw new WebRpcError('The session is idle — start a prompt first.', 'session_idle');
+      }
+      await host.proc.sendCommand({
+        type,
+        message: command.message as string,
+        ...(toImageContents(command.images) ? { images: toImageContents(command.images) } : {}),
+      });
+      return null;
+    }
+
+    default: {
+      if (PASSTHROUGH_COMMANDS.has(type)) {
+        const result: unknown = await host.proc.sendCommand(command as { type: string });
+        if (type === 'set_thinking_level') clearSessionFileCaches();
+        return result ?? null;
+      }
+      throw new Error(`Unsupported command: ${type}`);
+    }
+  }
+}