@larose/pi-web 0.3.0

package/src/server/runtime-registry.ts

@@ -0,0 +1,943 @@
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
+import {
+  type AgentSession,
+  type AgentSessionEvent,
+  type AgentSessionRuntime,
+  type ContextUsage,
+  createAgentSessionFromServices,
+  createAgentSessionRuntime,
+  createAgentSessionServices,
+  getAgentDir,
+  ModelRuntime,
+  type SessionManager,
+  type SessionStats,
+  type SlashCommandInfo,
+} from "@earendil-works/pi-coding-agent";
+
+import { AppError } from "./errors.js";
+import { EventBuffer, type BufferedEvent } from "./event-buffer.js";
+import {
+  type ExtensionUIEvent,
+  type ExtensionUIResponse,
+  type ExtensionUIState,
+  WebExtensionUI,
+} from "./extension-ui.js";
+import type { SessionRepository } from "./session-repository.js";
+
+export type RuntimeEvent =
+  | { type: "agent_event"; event: AgentSessionEvent }
+  | ExtensionUIEvent
+  | { type: "session_replaced"; previousId: string; sessionId: string }
+  | { type: "runtime_error"; message: string }
+  | { type: "runtime_disposed" };
+
+export interface RuntimeUsage {
+  tokens: SessionStats["tokens"];
+  context: ContextUsage | null;
+}
+
+export interface RuntimeState {
+  id: string;
+  cwd: string;
+  sessionName?: string;
+  isStreaming: boolean;
+  isCompacting: boolean;
+  isWorking: boolean;
+  pendingMessageCount: number;
+  queue: {
+    steering: string[];
+    followUp: string[];
+  };
+  connectedClients: number;
+  model: { provider: string; id: string; name: string } | null;
+  thinkingLevel: string;
+  autoCompactionEnabled: boolean;
+  usage: RuntimeUsage;
+  messages: AgentMessage[];
+  commands: SlashCommandInfo[];
+  extensionUI: ExtensionUIState;
+  streamingMessage?: AgentMessage;
+  latestEventId: number;
+}
+
+export interface RuntimeSession {
+  readonly sessionId: string;
+  readonly sessionManager: SessionManager;
+  readonly sessionName: string | undefined;
+  readonly cwd?: string;
+  readonly isStreaming: boolean;
+  readonly isCompacting: boolean;
+  readonly pendingMessageCount: number;
+  getSteeringMessages(): readonly string[];
+  getFollowUpMessages(): readonly string[];
+  readonly messages: AgentMessage[];
+  readonly commands: SlashCommandInfo[];
+  readonly model: AgentSession["model"];
+  readonly thinkingLevel: string;
+  readonly autoCompactionEnabled: boolean;
+  readonly agent: { state: { streamingMessage?: AgentMessage } };
+  getSessionStats(): SessionStats;
+  readonly extensionUI?: WebExtensionUI;
+  readonly ready?: Promise<void>;
+  subscribe(listener: (event: AgentSessionEvent) => void): () => void;
+  onSessionReplaced?(listener: (previousId: string, sessionId: string) => void): () => void;
+  prompt(
+    message: string,
+    options: { streamingBehavior?: "steer"; preflightResult(success: boolean): void },
+  ): Promise<void>;
+  steer(message: string): Promise<void>;
+  clearQueue(): { steering: string[]; followUp: string[] };
+  restoreQueue(queue: { steering: string[]; followUp: string[] }): Promise<void>;
+  abort(): Promise<void>;
+  setSessionName(name: string): void;
+  dispose(): void | Promise<void>;
+}
+
+export type RuntimeSessionFactory = (manager: SessionManager) => Promise<RuntimeSession>;
+
+export interface RuntimeRegistryOptions {
+  idleTimeoutMs?: number;
+  eventCapacity?: number;
+  createSession: RuntimeSessionFactory;
+}
+
+function errorMessage(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
+function sessionCommands(session: AgentSession): SlashCommandInfo[] {
+  const commands: SlashCommandInfo[] = [];
+
+  for (const command of session.extensionRunner.getRegisteredCommands()) {
+    commands.push({
+      name: command.invocationName,
+      ...(command.description === undefined ? {} : { description: command.description }),
+      source: "extension",
+      sourceInfo: command.sourceInfo,
+    });
+  }
+
+  for (const template of session.promptTemplates) {
+    commands.push({
+      name: template.name,
+      ...(template.description === undefined ? {} : { description: template.description }),
+      source: "prompt",
+      sourceInfo: template.sourceInfo,
+    });
+  }
+
+  for (const skill of session.resourceLoader.getSkills().skills) {
+    commands.push({
+      name: `skill:${skill.name}`,
+      ...(skill.description === undefined ? {} : { description: skill.description }),
+      source: "skill",
+      sourceInfo: skill.sourceInfo,
+    });
+  }
+
+  return commands;
+}
+
+class SdkRuntimeSession implements RuntimeSession {
+  readonly extensionUI = new WebExtensionUI();
+
+  private readonly runtime: AgentSessionRuntime;
+  private readonly listeners = new Set<(event: AgentSessionEvent) => void>();
+  private readonly replacementListeners = new Set<(previousId: string, sessionId: string) => void>();
+  private unsubscribeSession: (() => void) | undefined;
+  private binding: Promise<void>;
+  private commandCatalogue: SlashCommandInfo[];
+  private currentSessionId: string;
+  private currentSessionManager: SessionManager;
+  private disposed = false;
+
+  constructor(runtime: AgentSessionRuntime) {
+    this.runtime = runtime;
+    this.commandCatalogue = sessionCommands(runtime.session);
+    this.currentSessionId = runtime.session.sessionId;
+    this.currentSessionManager = runtime.session.sessionManager;
+    this.attachSession(runtime.session);
+    this.binding = this.bindSession(runtime.session);
+    void this.binding.catch(() => undefined);
+
+    runtime.setBeforeSessionInvalidate(() => this.extensionUI.reset());
+    runtime.setRebindSession(async (session) => {
+      const previousId = this.currentSessionId;
+      this.commandCatalogue = sessionCommands(session);
+      this.currentSessionId = session.sessionId;
+      this.currentSessionManager = session.sessionManager;
+      this.attachSession(session);
+      this.binding = this.bindSession(session);
+      void this.binding.catch(() => undefined);
+
+      for (const listener of this.replacementListeners) {
+        listener(previousId, session.sessionId);
+      }
+      await this.binding;
+    });
+  }
+
+  get sessionId(): string {
+    return this.currentSessionId;
+  }
+  get sessionManager(): SessionManager {
+    return this.currentSessionManager;
+  }
+  get sessionName(): string | undefined {
+    return this.runtime.session.sessionName;
+  }
+  get cwd(): string {
+    return this.runtime.session.sessionManager.getCwd();
+  }
+  get isStreaming(): boolean {
+    return this.runtime.session.isStreaming;
+  }
+  get isCompacting(): boolean {
+    return this.runtime.session.isCompacting;
+  }
+  get pendingMessageCount(): number {
+    return this.runtime.session.pendingMessageCount;
+  }
+  getSteeringMessages(): readonly string[] {
+    return [...this.runtime.session.getSteeringMessages()];
+  }
+  getFollowUpMessages(): readonly string[] {
+    return [...this.runtime.session.getFollowUpMessages()];
+  }
+  get messages(): AgentMessage[] {
+    return this.runtime.session.messages;
+  }
+  get commands(): SlashCommandInfo[] {
+    return this.commandCatalogue;
+  }
+  get model(): AgentSession["model"] {
+    return this.runtime.session.model;
+  }
+  get thinkingLevel(): string {
+    return this.runtime.session.thinkingLevel;
+  }
+  get autoCompactionEnabled(): boolean {
+    return this.runtime.session.autoCompactionEnabled;
+  }
+  get agent(): RuntimeSession["agent"] {
+    return this.runtime.session.agent;
+  }
+  get ready(): Promise<void> {
+    return this.binding;
+  }
+
+  getSessionStats(): SessionStats {
+    return this.runtime.session.getSessionStats();
+  }
+
+  subscribe(listener: (event: AgentSessionEvent) => void): () => void {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  }
+
+  onSessionReplaced(listener: (previousId: string, sessionId: string) => void): () => void {
+    this.replacementListeners.add(listener);
+    return () => this.replacementListeners.delete(listener);
+  }
+
+  prompt(
+    message: string,
+    options: { streamingBehavior?: "steer"; preflightResult(success: boolean): void },
+  ): Promise<void> {
+    return this.runtime.session.prompt(message, options);
+  }
+
+  steer(message: string): Promise<void> {
+    return this.runtime.session.steer(message);
+  }
+
+  clearQueue(): { steering: string[]; followUp: string[] } {
+    return this.runtime.session.clearQueue();
+  }
+
+  async restoreQueue(queue: { steering: string[]; followUp: string[] }): Promise<void> {
+    for (const message of queue.steering) {
+      await this.runtime.session.sendUserMessage(message, {
+        deliverAs: "steer",
+        expandPromptTemplates: false,
+      });
+    }
+    for (const message of queue.followUp) {
+      await this.runtime.session.sendUserMessage(message, {
+        deliverAs: "followUp",
+        expandPromptTemplates: false,
+      });
+    }
+  }
+
+  abort(): Promise<void> {
+    return this.runtime.session.abort();
+  }
+
+  setSessionName(name: string): void {
+    this.runtime.session.setSessionName(name);
+  }
+
+  async dispose(): Promise<void> {
+    if (this.disposed) {
+      return;
+    }
+    this.disposed = true;
+    this.extensionUI.dispose();
+    await this.runtime.dispose();
+    this.unsubscribeSession?.();
+    this.listeners.clear();
+    this.replacementListeners.clear();
+  }
+
+  private attachSession(session: AgentSession): void {
+    this.unsubscribeSession?.();
+    this.unsubscribeSession = session.subscribe((event) => {
+      for (const listener of this.listeners) {
+        listener(event);
+      }
+    });
+  }
+
+  private bindSession(session: AgentSession): Promise<void> {
+    return session.bindExtensions({
+      mode: "rpc",
+      uiContext: this.extensionUI.context,
+      commandContextActions: {
+        waitForIdle: () => session.waitForIdle(),
+        newSession: (options) => this.runtime.newSession(options),
+        fork: async (entryId, options) => {
+          const result = await this.runtime.fork(entryId, options);
+          return { cancelled: result.cancelled };
+        },
+        navigateTree: async (targetId, options) => {
+          const result = await session.navigateTree(targetId, options);
+          return { cancelled: result.cancelled };
+        },
+        switchSession: (sessionPath, options) => this.runtime.switchSession(sessionPath, options),
+        reload: async () => {
+          await session.reload();
+          this.commandCatalogue = sessionCommands(session);
+        },
+      },
+      shutdownHandler: () =>
+        this.extensionUI.context.notify(
+          "An extension requested shutdown. Stop Pi Web from its server process.",
+          "warning",
+        ),
+      onError: (error) =>
+        this.extensionUI.context.notify(
+          `Extension error in ${error.extensionPath} (${error.event}): ${error.error}`,
+          "error",
+        ),
+    });
+  }
+}
+
+export class RuntimeHandle {
+  readonly events: EventBuffer<RuntimeEvent>;
+
+  private readonly session: RuntimeSession;
+  private readonly initialCwd: string;
+  private readonly idleTimeoutMs: number;
+  private readonly onIdle: (handle: RuntimeHandle) => void;
+  private readonly onSessionIdChanged: (
+    handle: RuntimeHandle,
+    previousId: string,
+    sessionId: string,
+    manager: SessionManager,
+  ) => void;
+  private readonly unsubscribeSession: () => void;
+  private readonly unsubscribeExtensionUI: (() => void) | undefined;
+  private readonly unsubscribeReplacement: (() => void) | undefined;
+  private sessionId: string;
+  private activeTools = new Set<string>();
+  private activePrompt: Promise<void> | undefined;
+  private connectedClients = 0;
+  private disposed = false;
+  private idleTimer: NodeJS.Timeout | undefined;
+  private lastSettledAt = Date.now();
+  private operationTail: Promise<void> = Promise.resolve();
+  private queuedMessages = 0;
+  private retrying = false;
+  private suppressQueueUpdates = false;
+  private suppressedQueueUpdate: Extract<AgentSessionEvent, { type: "queue_update" }> | undefined;
+
+  constructor(
+    session: RuntimeSession,
+    cwd: string,
+    idleTimeoutMs: number,
+    eventCapacity: number,
+    onIdle: (handle: RuntimeHandle) => void,
+    onSessionIdChanged: (
+      handle: RuntimeHandle,
+      previousId: string,
+      sessionId: string,
+      manager: SessionManager,
+    ) => void = () => undefined,
+  ) {
+    this.session = session;
+    this.sessionId = session.sessionId;
+    this.initialCwd = cwd;
+    this.idleTimeoutMs = idleTimeoutMs;
+    this.events = new EventBuffer(eventCapacity);
+    this.onIdle = onIdle;
+    this.onSessionIdChanged = onSessionIdChanged;
+    this.unsubscribeSession = session.subscribe((event) => this.onSessionEvent(event));
+    this.unsubscribeExtensionUI = session.extensionUI?.setPublisher((event) => {
+      this.events.publish(event);
+
+      if (
+        event.type === "extension_ui_request" &&
+        (event.method === "select" ||
+          event.method === "confirm" ||
+          event.method === "input" ||
+          event.method === "editor")
+      ) {
+        this.cancelEviction();
+      }
+
+      if (event.type === "extension_ui_closed") {
+        this.scheduleEviction();
+      }
+    });
+    this.unsubscribeReplacement = session.onSessionReplaced?.((previousId, sessionId) => {
+      this.sessionId = sessionId;
+      this.onSessionIdChanged(this, previousId, sessionId, session.sessionManager);
+      this.events.publish({ type: "session_replaced", previousId, sessionId });
+    });
+    this.scheduleEviction();
+  }
+
+  get id(): string {
+    return this.sessionId;
+  }
+
+  get cwd(): string {
+    return this.session.cwd ?? this.initialCwd;
+  }
+
+  snapshot(): RuntimeState {
+    const streamingMessage = this.session.agent.state.streamingMessage;
+    const stats = this.session.getSessionStats();
+
+    return {
+      id: this.id,
+      cwd: this.cwd,
+      ...(this.session.sessionName ? { sessionName: this.session.sessionName } : {}),
+      isStreaming: this.session.isStreaming,
+      isCompacting: this.session.isCompacting,
+      isWorking: this.isWorking(),
+      pendingMessageCount: this.session.pendingMessageCount,
+      queue: {
+        steering: [...this.session.getSteeringMessages()],
+        followUp: [...this.session.getFollowUpMessages()],
+      },
+      connectedClients: this.connectedClients,
+      model: this.session.model
+        ? {
+            provider: this.session.model.provider,
+            id: this.session.model.id,
+            name: this.session.model.name,
+          }
+        : null,
+      thinkingLevel: this.session.thinkingLevel,
+      autoCompactionEnabled: this.session.autoCompactionEnabled,
+      usage: {
+        tokens: stats.tokens,
+        context: stats.contextUsage ?? null,
+      },
+      messages: this.session.messages,
+      commands: this.session.commands,
+      extensionUI: this.session.extensionUI?.snapshot() ?? { pending: [], statuses: [], widgets: [] },
+      ...(streamingMessage ? { streamingMessage } : {}),
+      latestEventId: this.events.latestId,
+    };
+  }
+
+  attachClient(listener: (event: BufferedEvent<RuntimeEvent>) => void): () => void {
+    this.assertOpen();
+    this.connectedClients += 1;
+    this.cancelEviction();
+    const unsubscribe = this.events.subscribe(listener);
+    let attached = true;
+
+    return () => {
+      if (!attached) {
+        return;
+      }
+      attached = false;
+      unsubscribe();
+      this.connectedClients -= 1;
+      this.scheduleEviction();
+    };
+  }
+
+  send(message: string): Promise<{ queued: boolean }> {
+    const trimmed = message.trim();
+    if (!trimmed) {
+      return Promise.reject(new AppError("invalid_message", "A message is required"));
+    }
+
+    return this.enqueue(async () => {
+      this.assertOpen();
+      this.cancelEviction();
+
+      if (this.session.isStreaming || this.activePrompt) {
+        await this.session.prompt(message, {
+          streamingBehavior: "steer",
+          preflightResult: () => undefined,
+        });
+        return { queued: true };
+      }
+
+      await this.startPrompt(message);
+      return { queued: false };
+    });
+  }
+
+  removePendingSteering(index: number, expectedMessage: string, expectedQueue: string[]): Promise<void> {
+    return this.enqueue(async () => {
+      this.assertOpen();
+      this.cancelEviction();
+
+      const steering = [...this.session.getSteeringMessages()];
+      const unchanged =
+        steering.length === expectedQueue.length &&
+        steering.every((message, queueIndex) => message === expectedQueue[queueIndex]);
+      if (!unchanged || steering[index] !== expectedMessage) {
+        throw new AppError("stale_pending_message", "The pending message queue has changed", 409);
+      }
+
+      this.suppressQueueUpdates = true;
+      try {
+        const queue = this.session.clearQueue();
+        queue.steering.splice(index, 1);
+        await this.session.restoreQueue(queue);
+      } finally {
+        this.suppressQueueUpdates = false;
+        const update = this.suppressedQueueUpdate;
+        this.suppressedQueueUpdate = undefined;
+        if (update) {
+          this.events.publish({ type: "agent_event", event: update });
+        }
+        this.scheduleEviction();
+      }
+    });
+  }
+
+  respondToExtensionUI(response: ExtensionUIResponse): void {
+    this.assertOpen();
+    if (!this.session.extensionUI?.respond(response)) {
+      throw new AppError("unknown_ui_request", `Unknown extension UI request: ${response.id}`, 409);
+    }
+  }
+
+  setSessionName(name: string | undefined): string | undefined {
+    this.assertOpen();
+    this.session.setSessionName(name ?? "");
+    return this.session.sessionName;
+  }
+
+  abort(): Promise<{ restoredMessages: string[] }> {
+    return this.enqueue(async () => {
+      this.assertOpen();
+      this.cancelEviction();
+      const queued = this.session.clearQueue();
+      this.queuedMessages = 0;
+      this.session.extensionUI?.cancelPending();
+      await this.session.abort();
+      this.retrying = false;
+      this.activeTools.clear();
+      this.lastSettledAt = Date.now();
+      this.scheduleEviction();
+
+      return { restoredMessages: [...queued.steering, ...queued.followUp] };
+    });
+  }
+
+  isEvictable(): boolean {
+    return !this.disposed && this.connectedClients === 0 && !this.isWorking();
+  }
+
+  async dispose(): Promise<void> {
+    if (this.disposed) {
+      return;
+    }
+    this.disposed = true;
+    this.cancelEviction();
+    this.session.clearQueue();
+
+    try {
+      await this.session.abort();
+    } finally {
+      this.events.publish({ type: "runtime_disposed" });
+      this.unsubscribeSession();
+      this.unsubscribeExtensionUI?.();
+      this.unsubscribeReplacement?.();
+      await this.session.dispose();
+    }
+  }
+
+  private enqueue<T>(operation: () => Promise<T>): Promise<T> {
+    const result = this.operationTail.then(operation, operation);
+    this.operationTail = result.then(
+      () => undefined,
+      () => undefined,
+    );
+    return result;
+  }
+
+  private async startPrompt(message: string): Promise<void> {
+    let preflightSettled = false;
+    let acceptPrompt!: () => void;
+    let rejectPrompt!: (error: unknown) => void;
+    const accepted = new Promise<void>((resolve, reject) => {
+      acceptPrompt = resolve;
+      rejectPrompt = reject;
+    });
+
+    const completion = (async () => {
+      await this.session.ready;
+      await this.session.prompt(message, {
+        preflightResult: (success) => {
+          if (preflightSettled) {
+            return;
+          }
+          preflightSettled = true;
+
+          if (success) {
+            acceptPrompt();
+          } else {
+            rejectPrompt(new AppError("prompt_rejected", "Pi rejected the message", 409));
+          }
+        },
+      });
+    })();
+    this.activePrompt = completion;
+
+    void completion
+      .then(() => {
+        if (preflightSettled) {
+          return;
+        }
+        preflightSettled = true;
+        acceptPrompt();
+      })
+      .catch((error: unknown) => {
+        if (!preflightSettled) {
+          preflightSettled = true;
+          rejectPrompt(error);
+        } else {
+          this.events.publish({ type: "runtime_error", message: errorMessage(error) });
+        }
+      })
+      .finally(() => {
+        if (this.activePrompt === completion) {
+          this.activePrompt = undefined;
+        }
+        this.scheduleEviction();
+      });
+
+    await accepted;
+  }
+
+  private onSessionEvent(event: AgentSessionEvent): void {
+    if (this.disposed) {
+      return;
+    }
+
+    switch (event.type) {
+      case "agent_start":
+        this.cancelEviction();
+        break;
+      case "tool_execution_start":
+        this.activeTools.add(event.toolCallId);
+        this.cancelEviction();
+        break;
+      case "tool_execution_end":
+        this.activeTools.delete(event.toolCallId);
+        break;
+      case "queue_update":
+        this.queuedMessages = event.steering.length + event.followUp.length;
+        if (this.queuedMessages > 0) {
+          this.cancelEviction();
+        }
+        break;
+      case "compaction_start":
+      case "auto_retry_start":
+      case "summarization_retry_scheduled":
+      case "summarization_retry_attempt_start":
+        this.retrying = true;
+        this.cancelEviction();
+        break;
+      case "compaction_end":
+        this.retrying = event.willRetry;
+        break;
+      case "auto_retry_end":
+      case "summarization_retry_finished":
+        this.retrying = false;
+        break;
+      case "agent_settled":
+        this.retrying = false;
+        this.activeTools.clear();
+        this.lastSettledAt = Date.now();
+        this.scheduleEviction();
+        break;
+    }
+
+    if (event.type === "queue_update" && this.suppressQueueUpdates) {
+      this.suppressedQueueUpdate = event;
+      return;
+    }
+
+    this.events.publish({ type: "agent_event", event });
+  }
+
+  isWorking(): boolean {
+    return (
+      this.session.isStreaming ||
+      this.session.isCompacting ||
+      this.activePrompt !== undefined ||
+      this.activeTools.size > 0 ||
+      this.queuedMessages > 0 ||
+      (this.session.extensionUI?.snapshot().pending.length ?? 0) > 0 ||
+      this.retrying
+    );
+  }
+
+  private scheduleEviction(): void {
+    this.cancelEviction();
+    if (!this.isEvictable()) {
+      return;
+    }
+
+    const elapsed = Date.now() - this.lastSettledAt;
+    const delay = Math.max(0, this.idleTimeoutMs - elapsed);
+    this.idleTimer = setTimeout(() => {
+      this.idleTimer = undefined;
+      if (this.isEvictable()) {
+        this.onIdle(this);
+      }
+    }, delay);
+    this.idleTimer.unref();
+  }
+
+  private cancelEviction(): void {
+    if (!this.idleTimer) {
+      return;
+    }
+    clearTimeout(this.idleTimer);
+    this.idleTimer = undefined;
+  }
+
+  private assertOpen(): void {
+    if (this.disposed) {
+      throw new AppError("runtime_closed", "The session runtime is closed", 409);
+    }
+  }
+}
+
+export class RuntimeRegistry {
+  private readonly repository: SessionRepository;
+  private readonly options: Required<Pick<RuntimeRegistryOptions, "idleTimeoutMs" | "eventCapacity">> &
+    RuntimeRegistryOptions;
+  private readonly handles = new Map<string, RuntimeHandle>();
+  private readonly starts = new Map<string, Promise<RuntimeHandle>>();
+  private readonly deletions = new Map<string, Promise<void>>();
+  private closing = false;
+
+  constructor(repository: SessionRepository, options: RuntimeRegistryOptions) {
+    this.repository = repository;
+    this.options = {
+      idleTimeoutMs: options.idleTimeoutMs ?? 10 * 60_000,
+      eventCapacity: options.eventCapacity ?? 512,
+      createSession: options.createSession,
+    };
+  }
+
+  async get(id: string): Promise<RuntimeHandle> {
+    if (this.closing) {
+      throw new AppError("server_stopping", "The server is stopping", 503);
+    }
+
+    const deleting = this.deletions.get(id);
+    if (deleting) {
+      await deleting;
+    }
+    if (this.closing) {
+      throw new AppError("server_stopping", "The server is stopping", 503);
+    }
+
+    const existing = this.handles.get(id);
+    if (existing) {
+      return existing;
+    }
+
+    const starting = this.starts.get(id);
+    if (starting) {
+      return starting;
+    }
+
+    const start = this.start(id);
+    this.starts.set(id, start);
+
+    try {
+      return await start;
+    } finally {
+      if (this.starts.get(id) === start) {
+        this.starts.delete(id);
+      }
+    }
+  }
+
+  getActive(id: string): RuntimeHandle | undefined {
+    return this.handles.get(id);
+  }
+
+  activeIds(): string[] {
+    return [...this.handles.keys()];
+  }
+
+  async delete(id: string): Promise<void> {
+    if (this.closing) {
+      throw new AppError("server_stopping", "The server is stopping", 503);
+    }
+
+    const existing = this.deletions.get(id);
+    if (existing) {
+      return existing;
+    }
+
+    const deletion = this.deleteSession(id);
+    this.deletions.set(id, deletion);
+
+    try {
+      await deletion;
+    } finally {
+      if (this.deletions.get(id) === deletion) {
+        this.deletions.delete(id);
+      }
+    }
+  }
+
+  async close(): Promise<void> {
+    if (this.closing) {
+      return;
+    }
+    this.closing = true;
+
+    await Promise.allSettled([...this.starts.values()]);
+    await Promise.allSettled([...this.deletions.values()]);
+    const handles = [...this.handles.values()];
+    this.handles.clear();
+    await Promise.allSettled(handles.map((handle) => handle.dispose()));
+  }
+
+  private async deleteSession(id: string): Promise<void> {
+    const starting = this.starts.get(id);
+    if (starting) {
+      await starting.catch(() => undefined);
+    }
+
+    const handle = this.handles.get(id);
+    if (handle?.isWorking()) {
+      throw new AppError("session_busy", "The session is busy and cannot be deleted", 409);
+    }
+
+    if (handle) {
+      for (const [key, candidate] of this.handles) {
+        if (candidate === handle) {
+          this.handles.delete(key);
+        }
+      }
+      await handle.dispose();
+    }
+
+    await this.repository.delete(id);
+  }
+
+  private async start(id: string): Promise<RuntimeHandle> {
+    const manager = await this.repository.openManager(id);
+    const session = await this.options.createSession(manager);
+
+    if (this.closing) {
+      await session.dispose();
+      throw new AppError("server_stopping", "The server is stopping", 503);
+    }
+
+    const handle = new RuntimeHandle(
+      session,
+      manager.getCwd(),
+      this.options.idleTimeoutMs,
+      this.options.eventCapacity,
+      (idleHandle) => void this.evict(idleHandle),
+      (changedHandle, previousId, sessionId, replacementManager) => {
+        this.rekey(changedHandle, previousId, sessionId, replacementManager);
+      },
+    );
+    this.handles.set(id, handle);
+    return handle;
+  }
+
+  private rekey(handle: RuntimeHandle, previousId: string, sessionId: string, manager: SessionManager): void {
+    this.repository.adoptManager(manager);
+    if (this.handles.get(previousId) === handle) {
+      this.handles.delete(previousId);
+    }
+
+    const existing = this.handles.get(sessionId);
+    if (existing && existing !== handle) {
+      void existing.dispose();
+    }
+    this.handles.set(sessionId, handle);
+  }
+
+  private async evict(handle: RuntimeHandle): Promise<void> {
+    if (!handle.isEvictable()) {
+      return;
+    }
+
+    for (const [id, candidate] of this.handles) {
+      if (candidate === handle) {
+        this.handles.delete(id);
+      }
+    }
+    await handle.dispose();
+  }
+}
+
+export async function createRuntimeRegistry(repository: SessionRepository): Promise<RuntimeRegistry> {
+  const modelRuntime = await ModelRuntime.create();
+  const agentDir = getAgentDir();
+
+  const createSession: RuntimeSessionFactory = async (manager) => {
+    const createRuntime = async ({
+      cwd,
+      sessionManager,
+      sessionStartEvent,
+    }: {
+      cwd: string;
+      sessionManager: SessionManager;
+      sessionStartEvent?: Parameters<typeof createAgentSessionFromServices>[0]["sessionStartEvent"];
+    }) => {
+      const services = await createAgentSessionServices({ cwd, agentDir, modelRuntime });
+      const created = await createAgentSessionFromServices({
+        services,
+        sessionManager,
+        ...(sessionStartEvent === undefined ? {} : { sessionStartEvent }),
+      });
+
+      return { ...created, services, diagnostics: services.diagnostics };
+    };
+    const runtime = await createAgentSessionRuntime(createRuntime, {
+      cwd: manager.getCwd(),
+      agentDir,
+      sessionManager: manager,
+    });
+
+    return new SdkRuntimeSession(runtime);
+  };
+
+  return new RuntimeRegistry(repository, { createSession });
+}