@seqyuan/annodex 0.1.10 → 0.1.12

package/lib/rpc-manager.ts

@@ -0,0 +1,668 @@
+/**
+ * RPC Manager — Codex app-server backend
+ *
+ * Replaces the original pi-agent backend. Manages CodexSession instances (backed by codex app-server)
+ * and routes commands from the Next.js API layer to the codex app-server.
+ *
+ * Session lifecycle:
+ *   1. POST /api/agent/new → startRpcSession() creates a new codex thread
+ *   2. POST /api/agent/[id] → getRpcSession()+send() routes commands
+ *   3. GET /api/agent/[id]/events → SSE stream from session.onEvent()
+ *   4. Idle timeout (10 min) → destroy()
+ */
+
+import { CodexSession, createCodexSession, resumeCodexSession, type AgentEvent, type EventListener } from "./codex-session";
+import {
+  getOrCreateCodexServer,
+  releaseCodexServerConnection,
+  codexServerEnvForModel,
+  type CodexServerConnection,
+} from "./codex-server";
+import { findProvider, normalizeRuntimeProviderId, readMergedConfig, resolveModel, type MergedConfig, type ProviderConfig } from "./annodex-config";
+import {
+  stripWidgetPromptBlocks,
+  withWidgetCapabilityPrompt,
+} from "./widget-prompt";
+import { stripAnalysisReportPrompt, withAnalysisReportPrompt } from "./report-prompt";
+import { readFileSync, existsSync, writeFileSync, mkdirSync } from "fs";
+import { dirname, join } from "path";
+import { homedir } from "os";
+import { findCodexSessionFile, getAgentDir } from "./session-reader";
+import { readSessionRuntime, upsertSessionRuntime } from "./session-runtime";
+import { readCodexUsageFromFile } from "./codex-usage";
+
+// ============================================================================
+// Types
+// ============================================================================
+
+type BusyOperation = "prompt" | "compact";
+
+export interface RuntimeSessionStatus {
+  sessionId: string;
+  sessionFile: string;
+  busy: boolean;
+  operations: BusyOperation[];
+  isStreaming: boolean;
+  isCompacting: boolean;
+}
+
+export interface RuntimeStatus {
+  busy: boolean;
+  activeSessions: number;
+  busySessions: RuntimeSessionStatus[];
+  sessions: RuntimeSessionStatus[];
+}
+
+function errorMessage(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
+function isThreadNotFoundError(error: unknown): boolean {
+  return /thread not found|not found/i.test(errorMessage(error));
+}
+
+// ============================================================================
+// CodexSessionWrapper — wraps CodexSession with command dispatch and lifecycle
+// ============================================================================
+
+export class CodexSessionWrapper {
+  private listeners: EventListener[] = [];
+  private cleanupEventSub: (() => void) | null = null;
+  private idleTimer: ReturnType<typeof setTimeout> | null = null;
+  private onDestroyCallback: (() => void) | null = null;
+  private busyOperations = new Set<BusyOperation>();
+  private _alive = true;
+  private _serverConnection: CodexServerConnection;
+  private _thinkingLevel = "auto";
+
+  constructor(
+    public inner: CodexSession,
+  ) {
+    this._serverConnection = inner.connection;
+  }
+
+  get sessionId(): string {
+    return this.inner.sessionId;
+  }
+
+  get sessionFile(): string {
+    return this.inner.sessionFile;
+  }
+
+  get cwd(): string {
+    return this.inner.cwd;
+  }
+
+  isAlive(): boolean {
+    return this._alive && this.inner.isAlive;
+  }
+
+  isBusy(): boolean {
+    return this.busyOperations.size > 0 || this.inner.isStreaming || this.inner.isCompacting;
+  }
+
+  getRuntimeStatus(): RuntimeSessionStatus {
+    return {
+      sessionId: this.sessionId,
+      sessionFile: this.sessionFile,
+      busy: this.isBusy(),
+      operations: [...this.busyOperations],
+      isStreaming: this.inner.isStreaming,
+      isCompacting: this.inner.isCompacting,
+    };
+  }
+
+  start(): void {
+    this.cleanupEventSub?.();
+    this.cleanupEventSub = this.inner.onEvent((event: AgentEvent) => {
+      this.resetIdleTimer();
+      for (const l of this.listeners) l(event);
+    });
+    this.resetIdleTimer();
+  }
+
+  private resetIdleTimer(): void {
+    if (this.idleTimer) clearTimeout(this.idleTimer);
+    this.idleTimer = setTimeout(() => {
+      if (this.isBusy()) {
+        this.resetIdleTimer();
+        return;
+      }
+      this.destroy();
+    }, 10 * 60 * 1000);
+  }
+
+  onEvent(listener: EventListener): () => void {
+    this.listeners.push(listener);
+    return () => {
+      const i = this.listeners.indexOf(listener);
+      if (i !== -1) this.listeners.splice(i, 1);
+    };
+  }
+
+  onDestroy(cb: () => void): void {
+    this.onDestroyCallback = cb;
+  }
+
+  async send(command: Record<string, unknown>): Promise<unknown> {
+    this.resetIdleTimer();
+    const type = command.type as string;
+
+    switch (type) {
+      case "prompt": {
+        const promptImages = command.images as Array<{ type: "image"; data: string; mimeType: string }> | undefined;
+        const message = command.message as string;
+        // Fire and forget — events come via onEvent subscription
+        this.trackBusy("prompt", this.inner.prompt(message, promptImages?.map((img) => ({
+          data: img.data,
+          mimeType: img.mimeType,
+        })))).catch(() => {});
+        return null;
+      }
+
+      case "abort": {
+        await this.inner.abort();
+        return null;
+      }
+
+      case "get_state": {
+        const usage = readCodexUsageFromFile(findCodexSessionFile(this.inner.sessionId), this.inner.model?.modelId);
+        return {
+          sessionId: this.inner.sessionId,
+          sessionFile: "",
+          isStreaming: this.inner.isStreaming,
+          isCompacting: this.inner.isCompacting,
+          autoCompactionEnabled: true,
+          autoRetryEnabled: true,
+          model: this.inner.model ? {
+            id: this.inner.model.modelId,
+            provider: this.inner.model.provider,
+            ...(usage.contextUsage?.contextWindow ? { contextWindow: usage.contextUsage.contextWindow } : {}),
+          } : undefined,
+          messageCount: 0,
+          pendingMessageCount: 0,
+          contextUsage: usage.contextUsage,
+          systemPrompt: this.inner.baseInstructions ?? "",
+          thinkingLevel: this._thinkingLevel,
+        };
+      }
+
+      case "set_model": {
+        const { provider, modelId } = command as { provider: string; modelId: string };
+        const cfg = readMergedConfig();
+        const normalizedProvider = normalizeRuntimeProviderId(cfg, provider);
+        const target = resolveModel(cfg, modelId, normalizedProvider ?? provider);
+        const targetProvider = target?.provider ?? findProvider(cfg, normalizedProvider ?? provider);
+        const targetProviderId = target?.provider.id ?? normalizedProvider ?? provider;
+        const targetModelId = target?.modelId ?? modelId;
+
+        if (targetProvider) {
+          const targetConnection = await getOrCreateCodexServer(this.cwd, codexServerEnvForModel(
+            targetProvider,
+            targetModelId,
+            this._thinkingLevel,
+          ));
+          if (targetConnection.serverKey !== this._serverConnection.serverKey) {
+            // Cross-provider switch: each codex app-server instance only knows
+            // its own threads, so the old thread won't be found on the new
+            // server. Create a fresh thread directly.
+            const previousConnection = this._serverConnection;
+            const previousInner = this.inner;
+            const previousSessionId = previousInner.sessionId;
+            let nextInner: CodexSession | undefined;
+            try {
+              nextInner = await resumeCodexSession(targetConnection, this.sessionId, this.cwd);
+            } catch {
+              // Thread not accessible from new provider instance — expected.
+            }
+            if (!nextInner) {
+              nextInner = await createCodexSession(targetConnection, {
+                cwd: this.cwd,
+                baseInstructions: previousInner.baseInstructions,
+                model: { provider: targetProviderId, modelId: targetModelId },
+              });
+            }
+            await nextInner.setModel(targetProviderId, targetModelId);
+            this.cleanupEventSub?.();
+            previousInner.destroy();
+            this.inner = nextInner;
+            this._serverConnection = targetConnection;
+            this.start();
+            releaseCodexServerConnection(previousConnection);
+            if (nextInner.sessionId !== previousSessionId) {
+              const registry = getRegistry();
+              aliasSession(previousSessionId, nextInner.sessionId);
+              registry.delete(previousSessionId);
+              registry.set(nextInner.sessionId, this);
+              this.onDestroyCallback = () => {
+                registry.delete(nextInner.sessionId);
+                clearAliasesFor(nextInner.sessionId);
+              };
+              upsertSessionRuntime(nextInner.sessionId, {
+                cwd: this.cwd,
+                provider: targetProviderId,
+                modelId: targetModelId,
+                thinkingLevel: this._thinkingLevel,
+              });
+              return { id: targetModelId, provider: targetProviderId, newSessionId: nextInner.sessionId };
+            }
+          } else {
+            releaseCodexServerConnection(targetConnection);
+            await this.inner.setModel(targetProviderId, targetModelId);
+          }
+        } else {
+          await this.inner.setModel(targetProviderId, targetModelId);
+        }
+        upsertSessionRuntime(this.sessionId, {
+          cwd: this.cwd,
+          provider: this.inner.model?.provider ?? targetProviderId,
+          modelId: this.inner.model?.modelId ?? targetModelId,
+        });
+        return { id: targetModelId, provider: targetProviderId };
+      }
+
+      case "rename": {
+        const name = typeof command.name === "string" ? command.name : "";
+        await this.inner.rename(name);
+        return null;
+      }
+
+      case "fork": {
+        const result = await this.inner.fork();
+        if ("newThreadId" in result) {
+          upsertSessionRuntime(result.newThreadId, {
+            cwd: this.cwd,
+            provider: this.inner.model?.provider,
+            modelId: this.inner.model?.modelId,
+            thinkingLevel: this._thinkingLevel,
+          });
+          this.destroy();
+          return { cancelled: false, newSessionId: result.newThreadId };
+        }
+        return { cancelled: true };
+      }
+
+      case "navigate_tree": {
+        const targetId = command.targetId as string;
+        const result = await this.inner.navigateTree(targetId);
+        return { cancelled: result.cancelled };
+      }
+
+      case "compact": {
+        await this.trackBusy("compact", this.inner.compact());
+        return null;
+      }
+
+      case "get_tools": {
+        // Codex has built-in tools (bash, edit, web_search, etc.)
+        // Return a basic set — codex manages its own tool activation
+        return [
+          { name: "read", description: "Read files", active: true },
+          { name: "bash", description: "Run shell commands", active: true },
+          { name: "edit", description: "Edit files", active: true },
+          { name: "write", description: "Create files", active: true },
+          { name: "grep", description: "Search files", active: true },
+          { name: "find", description: "List files", active: true },
+        ];
+      }
+
+      case "set_tools": {
+        // Codex doesn't support dynamic tool activation in the same way
+        // This is a no-op for now
+        return null;
+      }
+
+      case "set_thinking_level": {
+        this._thinkingLevel = typeof command.level === "string" ? command.level : "auto";
+        upsertSessionRuntime(this.sessionId, {
+          cwd: this.cwd,
+          thinkingLevel: this._thinkingLevel,
+          provider: this.inner.model?.provider,
+          modelId: this.inner.model?.modelId,
+        });
+        return null;
+      }
+
+      case "abort_compaction": {
+        // Codex manages compaction internally
+        return null;
+      }
+
+      case "set_auto_compaction":
+      case "set_auto_retry":
+        return null;
+
+      default:
+        throw new Error(`Unsupported command: ${type}`);
+    }
+  }
+
+  private async trackBusy<T>(operation: BusyOperation, promise: Promise<T>): Promise<T> {
+    this.busyOperations.add(operation);
+    this.resetIdleTimer();
+    try {
+      return await promise;
+    } finally {
+      this.busyOperations.delete(operation);
+      this.resetIdleTimer();
+    }
+  }
+
+  destroy(): void {
+    if (!this._alive) return;
+    this._alive = false;
+    if (this.idleTimer) clearTimeout(this.idleTimer);
+    this.cleanupEventSub?.();
+    this.inner.destroy();
+    releaseCodexServerConnection(this._serverConnection);
+    this.onDestroyCallback?.();
+  }
+}
+
+// ============================================================================
+// Session registry (replaces pi session map)
+// ============================================================================
+
+declare global {
+  var __annodexSessions: Map<string, CodexSessionWrapper> | undefined;
+  var __annodexStartLocks: Map<string, Promise<{ session: CodexSessionWrapper; realSessionId: string }>> | undefined;
+  var __annodexSessionAliases: Map<string, string> | undefined;
+}
+
+function getRegistry(): Map<string, CodexSessionWrapper> {
+  if (!globalThis.__annodexSessions) {
+    globalThis.__annodexSessions = new Map();
+    const cleanup = () => globalThis.__annodexSessions?.forEach((s) => s.destroy());
+    process.once("exit", cleanup);
+    process.once("SIGINT", cleanup);
+    process.once("SIGTERM", cleanup);
+  }
+  return globalThis.__annodexSessions;
+}
+
+function getAliases(): Map<string, string> {
+  if (!globalThis.__annodexSessionAliases) globalThis.__annodexSessionAliases = new Map();
+  return globalThis.__annodexSessionAliases;
+}
+
+function resolveAlias(sessionId: string): string {
+  let current = sessionId;
+  const seen = new Set<string>();
+  while (!seen.has(current)) {
+    seen.add(current);
+    const next = getAliases().get(current);
+    if (!next) return current;
+    current = next;
+  }
+  return current;
+}
+
+function aliasSession(requestedSessionId: string, realSessionId: string): void {
+  if (requestedSessionId && requestedSessionId !== realSessionId) {
+    getAliases().set(requestedSessionId, realSessionId);
+  }
+}
+
+function clearAliasesFor(realSessionId: string): void {
+  const aliases = getAliases();
+  aliases.delete(realSessionId);
+  for (const [from, to] of aliases) {
+    if (to === realSessionId) aliases.delete(from);
+  }
+}
+
+function getLocks(): Map<string, Promise<{ session: CodexSessionWrapper; realSessionId: string }>> {
+  if (!globalThis.__annodexStartLocks) globalThis.__annodexStartLocks = new Map();
+  return globalThis.__annodexStartLocks;
+}
+
+export function getRpcSession(sessionId: string): CodexSessionWrapper | undefined {
+  return getRegistry().get(resolveAlias(sessionId));
+}
+
+export function getRuntimeStatus(): RuntimeStatus {
+  const sessions = [...getRegistry().values()]
+    .filter((s) => s.isAlive())
+    .map((s) => s.getRuntimeStatus());
+  const busySessions = sessions.filter((s) => s.busy);
+  return {
+    busy: busySessions.length > 0,
+    activeSessions: sessions.length,
+    busySessions,
+    sessions,
+  };
+}
+
+// ============================================================================
+// buildSystemPrompt — replaces pi's system prompt construction
+// ============================================================================
+
+function buildSystemPrompt(cwd: string, config: MergedConfig): string {
+  const agentDir = getAgentDir();
+  const contextBlocks: string[] = [];
+
+  // 1. SOUL.md — global personality / identity (highest priority, first in prompt)
+  //    Location: ~/.config/annodex/SOUL.md
+  const soulPath = (() => {
+    const newPath = join(agentDir, "SOUL.md");
+    const oldPath = join(homedir(), ".pi", "agent", "SOUL.md");
+    // Migrate from old path on first access
+    if (!existsSync(newPath) && existsSync(oldPath)) {
+      try {
+        const content = readFileSync(oldPath, "utf-8");
+        if (!existsSync(dirname(newPath))) mkdirSync(dirname(newPath), { recursive: true });
+        writeFileSync(newPath, content, "utf-8");
+      } catch { /* skip migration */ }
+    }
+    const lower = join(agentDir, "soul.md");
+    return existsSync(newPath) ? newPath : lower;
+  })();
+  if (existsSync(soulPath)) {
+    try {
+      const content = readFileSync(soulPath, "utf-8").trim();
+      if (content) {
+        contextBlocks.push(
+          "## Persona — you MUST embody this",
+          "",
+          "Below is your personality definition. You MUST adopt this persona and tone " +
+          "in all replies. Avoid stiff, generic, or AI-assistant-style responses. " +
+          "Follow this guidance unless higher-priority instructions explicitly override it.",
+          "",
+          content,
+        );
+      }
+    } catch { /* ignore */ }
+  }
+
+  // 2. HARNESS.md — global behavioral constraints (output guardrails)
+  //    Location: ~/.config/annodex/HARNESS.md
+  const harnessPath = (() => {
+    const newPath = join(agentDir, "HARNESS.md");
+    const oldPath = join(homedir(), ".pi", "agent", "HARNESS.md");
+    if (!existsSync(newPath) && existsSync(oldPath)) {
+      try {
+        const content = readFileSync(oldPath, "utf-8");
+        if (!existsSync(dirname(newPath))) mkdirSync(dirname(newPath), { recursive: true });
+        writeFileSync(newPath, content, "utf-8");
+      } catch { /* skip migration */ }
+    }
+    const lower = join(agentDir, "harness.md");
+    return existsSync(newPath) ? newPath : lower;
+  })();
+  if (existsSync(harnessPath)) {
+    try {
+      const content = readFileSync(harnessPath, "utf-8").trim();
+      if (content) {
+        contextBlocks.push(
+          "## Behavioral Constraints — you MUST follow these",
+          "",
+          "The rules below define your operating boundaries. " +
+          "You MUST comply with every constraint. " +
+          "These are NOT suggestions — they are mandatory behavioral rules that " +
+          "supersede conflicting persona preferences.",
+          "",
+          content,
+        );
+      }
+    } catch { /* ignore */ }
+  }
+
+  // 3. AGENTS.md / CLAUDE.md — project-level instructions
+  //    Location: {cwd}/AGENTS.md or {cwd}/CLAUDE.md
+  const agentsPaths = [join(cwd, "AGENTS.md"), join(cwd, "CLAUDE.md")];
+  for (const agentsPath of agentsPaths) {
+    if (existsSync(agentsPath)) {
+      try {
+        const content = readFileSync(agentsPath, "utf-8").trim();
+        if (content) {
+          contextBlocks.push("## Project Instructions", "", content);
+          break; // only include one
+        }
+      } catch { /* ignore */ }
+    }
+  }
+
+  // 4. MEMORY.md — project-level self-updating memory
+  //    Location: {cwd}/MEMORY.md
+  const memoryPath = existsSync(join(cwd, "MEMORY.md"))
+    ? join(cwd, "MEMORY.md")
+    : join(cwd, "memory.md");
+  if (existsSync(memoryPath)) {
+    try {
+      const content = readFileSync(memoryPath, "utf-8").trim();
+      if (content) contextBlocks.push("## Project Memory", "", content);
+    } catch { /* ignore */ }
+  }
+
+  let basePrompt = contextBlocks.join("\n\n");
+
+  // 5. Apply generative UI / widget / report prompts (after persona context)
+  if (config.generativeUI) {
+    basePrompt = withAnalysisReportPrompt(withWidgetCapabilityPrompt(
+      stripAnalysisReportPrompt(stripWidgetPromptBlocks(basePrompt)),
+    ));
+  } else {
+    basePrompt = withAnalysisReportPrompt(
+      stripAnalysisReportPrompt(stripWidgetPromptBlocks(basePrompt)),
+    );
+  }
+
+  return basePrompt;
+}
+
+// ============================================================================
+// startRpcSession — create or resume a codex session
+// ============================================================================
+
+export async function startRpcSession(
+  sessionId: string,
+  sessionFile: string,
+  cwd: string,
+  toolNames?: string[],
+  modelInfo?: { provider: string; modelId: string },
+  thinkingLevel?: string,
+): Promise<{ session: CodexSessionWrapper; realSessionId: string }> {
+  const registry = getRegistry();
+  const locks = getLocks();
+  const effectiveSessionId = resolveAlias(sessionId);
+
+  // Check registry for existing session
+  const existing = registry.get(effectiveSessionId);
+  if (existing?.isAlive()) {
+    return { session: existing, realSessionId: effectiveSessionId };
+  }
+
+  // Check for in-flight start
+  const inflight = locks.get(effectiveSessionId);
+  if (inflight) return inflight;
+
+  const starting = (async () => {
+    const config = readMergedConfig();
+    const baseInstructions = buildSystemPrompt(cwd, config);
+
+    const storedRuntime = sessionFile ? readSessionRuntime(sessionId) : null;
+    const resumeModelInfo = modelInfo
+      ?? (storedRuntime?.modelId
+        ? { provider: normalizeRuntimeProviderId(config, storedRuntime.provider), modelId: storedRuntime.modelId }
+        : undefined);
+    const effectiveThinkingLevel = thinkingLevel ?? storedRuntime?.thinkingLevel;
+
+    // Only apply a model when the UI/session runtime has one. Without an
+    // explicit selection, Codex should use the user's ~/.codex/config.toml.
+    const resolved = resumeModelInfo ? resolveModel(config, resumeModelInfo.modelId, resumeModelInfo.provider) : null;
+    const providerConfig: ProviderConfig | undefined = resolved?.provider
+      ?? (resumeModelInfo?.provider ? findProvider(config, resumeModelInfo.provider) : undefined);
+    const modelId = resumeModelInfo?.modelId ?? resolved?.modelId;
+    const modelProviderId = resolved?.provider.id ?? resumeModelInfo?.provider;
+    // Get or create the codex app-server for this cwd
+    const connection = await getOrCreateCodexServer(
+      cwd,
+      providerConfig
+        ? codexServerEnvForModel(providerConfig, modelId, effectiveThinkingLevel)
+        : {
+          modelId,
+          thinkingLevel: effectiveThinkingLevel,
+        },
+    );
+
+    let codexSession: CodexSession;
+
+    if (!sessionFile) {
+      // New session — create a new codex thread
+      codexSession = await createCodexSession(connection, {
+        cwd,
+        baseInstructions,
+        acceptReturnedModel: Boolean(config.codexDefaults.model),
+        ...(modelId && modelProviderId ? { model: { provider: modelProviderId, modelId } } : {}),
+      });
+    } else {
+      // Existing session — resume the codex thread
+      try {
+        codexSession = await resumeCodexSession(connection, sessionId, cwd);
+      } catch (error) {
+        if (!isThreadNotFoundError(error)) throw error;
+        codexSession = await createCodexSession(connection, {
+          cwd,
+          baseInstructions,
+          acceptReturnedModel: Boolean(config.codexDefaults.model),
+          ...(modelId && modelProviderId ? { model: { provider: modelProviderId, modelId } } : {}),
+        });
+      }
+    }
+
+    const activeProvider = codexSession.model?.provider ?? modelProviderId;
+    const activeModelId = codexSession.model?.modelId ?? modelId;
+    if (activeProvider && activeModelId && (
+      !codexSession.model ||
+      codexSession.model.provider !== activeProvider ||
+      codexSession.model.modelId !== activeModelId
+    )) {
+      await codexSession.setModel(activeProvider, activeModelId);
+    }
+
+    const wrapper = new CodexSessionWrapper(codexSession);
+    const realSessionId = codexSession.sessionId;
+    aliasSession(sessionId, realSessionId);
+    upsertSessionRuntime(realSessionId, {
+      cwd,
+      provider: codexSession.model?.provider ?? activeProvider,
+      modelId: codexSession.model?.modelId ?? activeModelId,
+      thinkingLevel: effectiveThinkingLevel,
+    });
+    if (effectiveThinkingLevel) {
+      await wrapper.send({ type: "set_thinking_level", level: effectiveThinkingLevel });
+    }
+    wrapper.start();
+
+    wrapper.onDestroy(() => {
+      registry.delete(realSessionId);
+      clearAliasesFor(realSessionId);
+    });
+    registry.set(realSessionId, wrapper);
+
+    return { session: wrapper, realSessionId };
+  })().finally(() => locks.delete(effectiveSessionId));
+
+  locks.set(effectiveSessionId, starting);
+  return starting;
+}