@seqyuan/annodex 0.1.11 → 0.1.12

package/lib/codex-session.ts

@@ -0,0 +1,590 @@
+/**
+ * Codex Session Wrapper
+ *
+ * Wraps a codex app-server thread as an annodex session.
+ * Maps codex thread events to the unified AgentEvent protocol used by the SSE stream.
+ *
+ * Event mapping:
+ *   codex app-server notifications → AgentEvent
+ *   - item/agentMessage/delta → message_update
+ *   - item/completed          → message_end / tool_execution_end
+ *   - turn/completed          → agent_end
+ */
+
+import type { CodexServerConnection, CodexThread, CodexThreadItem, CodexTurn, CodexUserInput } from "./codex-server";
+import { subagentRunsFromCollabItem } from "./subagent-progress";
+
+// ============================================================================
+// AgentEvent — the unified event protocol consumed by SSE and ChatWindow
+// ============================================================================
+
+export interface AgentEvent {
+  type: string;
+  [key: string]: unknown;
+}
+
+// ============================================================================
+// Codex thread raw event types (from app-server notifications)
+// ============================================================================
+
+interface CodexNotificationBase {
+  threadId: string;
+  turnId: string;
+}
+
+interface CodexItemNotification extends CodexNotificationBase {
+  item: CodexThreadItem;
+}
+
+interface CodexDeltaNotification extends CodexNotificationBase {
+  itemId: string;
+  delta: string;
+}
+
+interface CodexTurnCompletedNotification {
+  threadId: string;
+  turn: CodexTurn;
+}
+
+interface CodexThreadStatusNotification {
+  threadId: string;
+  status?: { type?: string; activeFlags?: string[] };
+}
+
+type CodexCompactedNotification = CodexNotificationBase;
+
+function userInputsFromMessage(message: string, images?: Array<{ data: string; mimeType: string }>): CodexUserInput[] {
+  const input: CodexUserInput[] = [];
+  if (message.trim()) {
+    input.push({ type: "text", text: message, text_elements: [] });
+  }
+  for (const image of images ?? []) {
+    const mime = image.mimeType || "image/png";
+    input.push({ type: "image", url: `data:${mime};base64,${image.data}` });
+  }
+  return input.length > 0 ? input : [{ type: "text", text: "", text_elements: [] }];
+}
+
+function itemText(item: CodexThreadItem): string {
+  if (item.type === "agentMessage" && typeof item.text === "string") return item.text;
+  if (item.type === "reasoning") return [...(item.summary ?? []), ...(item.content ?? [])].join("\n");
+  if (item.type === "plan" && typeof item.text === "string") return item.text;
+  return "";
+}
+
+function itemToolName(item: CodexThreadItem): string {
+  if (item.type === "commandExecution") return "bash";
+  if (item.type === "fileChange") return "edit";
+  if (item.type === "mcpToolCall") return item.tool || item.server || "mcp";
+  if (item.type === "dynamicToolCall") return item.tool || "tool";
+  if (item.type === "collabAgentToolCall") return "agents";
+  return item.type;
+}
+
+function itemToolResult(item: CodexThreadItem): string {
+  if (item.type === "commandExecution") {
+    const out = item.aggregatedOutput ?? "";
+    const exit = item.exitCode !== null && item.exitCode !== undefined ? `Exit code: ${item.exitCode}` : "";
+    return out ? (exit ? `${out}\n${exit}` : out) : exit;
+  }
+  if (item.type === "fileChange") return JSON.stringify(item.changes ?? []);
+  if (item.type === "mcpToolCall") {
+    if (item.error) return typeof item.error === "string" ? item.error : JSON.stringify(item.error);
+    if (item.result) return typeof item.result === "string" ? item.result : JSON.stringify(item.result);
+    return "";
+  }
+  if (item.type === "dynamicToolCall") return JSON.stringify(item.contentItems ?? []);
+  if (item.type === "collabAgentToolCall") return JSON.stringify(item.agentsStates ?? []);
+  return "";
+}
+
+function itemIsFailed(item: CodexThreadItem): boolean {
+  if ("status" in item && typeof item.status === "string") {
+    return item.status === "failed" || item.status === "declined";
+  }
+  if (item.type === "dynamicToolCall" && item.success === false) return true;
+  return false;
+}
+
+function itemIsProgressItem(item: CodexThreadItem): boolean {
+  return item.type === "commandExecution"
+    || item.type === "fileChange"
+    || item.type === "mcpToolCall"
+    || item.type === "dynamicToolCall"
+    || item.type === "collabAgentToolCall";
+}
+
+// ============================================================================
+// CodexSession — wraps a single thread
+// ============================================================================
+
+export type EventListener = (event: AgentEvent) => void;
+
+export class CodexSession {
+  private listeners: EventListener[] = [];
+  private cleanupFns: Array<() => void> = [];
+  private _alive = true;
+  private _isStreaming = false;
+  private _isCompacting = false;
+  private _accumulatedText = "";
+  private _activeTurnId: string | null = null;
+  private _activeAgentItemId: string | null = null;
+  private _currentModel: { provider: string; modelId: string } | null = null;
+
+  constructor(
+    public readonly threadId: string,
+    public readonly cwd: string,
+    public readonly connection: CodexServerConnection,
+    public readonly baseInstructions?: string,
+  ) {
+    this.setupNotifications();
+  }
+
+  // ---- Public state ----
+
+  get sessionId(): string {
+    return this.threadId;
+  }
+
+  get sessionFile(): string {
+    // Codex manages threads internally; we don't have .jsonl files
+    return "";
+  }
+
+  get isAlive(): boolean {
+    return this._alive;
+  }
+
+  get isStreaming(): boolean {
+    return this._isStreaming;
+  }
+
+  get isCompacting(): boolean {
+    return this._isCompacting;
+  }
+
+  get model(): { provider: string; modelId: string } | null {
+    return this._currentModel;
+  }
+
+  // ---- Event subscription ----
+
+  onEvent(listener: EventListener): () => void {
+    this.listeners.push(listener);
+    return () => {
+      const i = this.listeners.indexOf(listener);
+      if (i !== -1) this.listeners.splice(i, 1);
+    };
+  }
+
+  private emit(event: AgentEvent): void {
+    for (const l of this.listeners) {
+      try { l(event); } catch { /* ignore */ }
+    }
+  }
+
+  private assistantMeta(): { model: string; provider: string } {
+    return {
+      model: this._currentModel?.modelId ?? "",
+      provider: this._currentModel?.provider ?? "",
+    };
+  }
+
+  // ---- Notifications from codex app-server ----
+
+  private setupNotifications(): void {
+    const unsub1 = this.connection.onNotification(
+      "item/agentMessage/delta",
+      (params: unknown) => this.handleAgentMessageDelta(params as CodexDeltaNotification),
+    );
+    this.cleanupFns.push(unsub1);
+
+    const unsub2 = this.connection.onNotification(
+      "item/completed",
+      (params: unknown) => this.handleItemCompleted(params as CodexItemNotification),
+    );
+    this.cleanupFns.push(unsub2);
+
+    const unsub3 = this.connection.onNotification(
+      "turn/completed",
+      (params: unknown) => this.handleTurnCompleted(params as CodexTurnCompletedNotification),
+    );
+    this.cleanupFns.push(unsub3);
+
+    const unsub4 = this.connection.onNotification(
+      "thread/status/changed",
+      (params: unknown) => this.handleThreadStatusChanged(params as CodexThreadStatusNotification),
+    );
+    this.cleanupFns.push(unsub4);
+
+    const unsub5 = this.connection.onNotification(
+      "thread/compacted",
+      (params: unknown) => this.handleCompacted(params as CodexCompactedNotification),
+    );
+    this.cleanupFns.push(unsub5);
+
+    const unsub6 = this.connection.onNotification(
+      "item/started",
+      (params: unknown) => this.handleItemStarted(params as CodexItemNotification),
+    );
+    this.cleanupFns.push(unsub6);
+
+    const unsub7 = this.connection.onNotification(
+      "error",
+      (params: unknown) => this.handleError(params as { threadId?: string; error?: { message?: string; additionalDetails?: string | null } }),
+    );
+    this.cleanupFns.push(unsub7);
+  }
+
+  private eventBelongsToThread(event: { threadId?: string }): boolean {
+    return event.threadId === this.threadId;
+  }
+
+  private handleAgentMessageDelta(event: CodexDeltaNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    if (!this._isStreaming) return;
+    this._activeTurnId = event.turnId;
+    this._activeAgentItemId = event.itemId;
+    this._accumulatedText += event.delta ?? "";
+    this.emit({
+      type: "message_update",
+      message: {
+        role: "assistant",
+        content: [{ type: "text", text: this._accumulatedText }],
+        ...this.assistantMeta(),
+      },
+    });
+  }
+
+  private handleItemStarted(event: CodexItemNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    const item = event.item;
+    if (!itemIsProgressItem(item)) {
+      return;
+    }
+    this._activeTurnId = event.turnId;
+    if (item.type === "collabAgentToolCall") {
+      this.emit({
+        type: "subagent_progress",
+        runs: subagentRunsFromCollabItem(item, Date.now()),
+      });
+      return;
+    }
+    this.emit({
+      type: "tool_execution_start",
+      toolCallId: item.id,
+      toolName: itemToolName(item),
+    });
+  }
+
+  private handleItemCompleted(event: CodexItemNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    const item = event.item;
+    this._activeTurnId = event.turnId;
+
+    if (item.type === "agentMessage") {
+      const text = itemText(item) || this._accumulatedText;
+      this.emit({
+        type: "message_end",
+        message: {
+          role: "assistant",
+          content: [{ type: "text", text }],
+          ...this.assistantMeta(),
+        },
+      });
+      this._accumulatedText = "";
+      this._activeAgentItemId = null;
+    } else if (item.type === "reasoning") {
+      const text = itemText(item);
+      if (text) {
+        this.emit({
+          type: "message_update",
+          message: {
+            role: "assistant",
+            content: [{ type: "thinking", thinking: text }],
+            ...this.assistantMeta(),
+          },
+        });
+      }
+    } else if (item.type === "collabAgentToolCall") {
+      this.emit({
+        type: "subagent_progress",
+        runs: subagentRunsFromCollabItem(item, Date.now()),
+      });
+    } else if (item.type === "commandExecution" || item.type === "fileChange" || item.type === "mcpToolCall" || item.type === "dynamicToolCall") {
+      this.emit({
+        type: "tool_execution_end",
+        toolCallId: item.id,
+        toolName: itemToolName(item),
+        result: itemToolResult(item),
+        isError: itemIsFailed(item),
+      });
+    }
+  }
+
+  private handleTurnCompleted(event: CodexTurnCompletedNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    this._isStreaming = false;
+    this._activeTurnId = null;
+    if (this._accumulatedText) {
+      this.emit({
+        type: "message_end",
+        message: {
+          role: "assistant",
+          content: [{ type: "text", text: this._accumulatedText }],
+          ...this.assistantMeta(),
+        },
+      });
+      this._accumulatedText = "";
+    }
+    this.emit({
+      type: "agent_end",
+      ...(event.turn?.status === "failed" ? { error: event.turn.error?.message ?? "Turn failed" } : {}),
+    });
+  }
+
+  private handleThreadStatusChanged(event: CodexThreadStatusNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    if (event.status?.type === "idle" && this._isStreaming) {
+      this._isStreaming = false;
+      this._activeTurnId = null;
+      this.emit({ type: "agent_end" });
+    }
+  }
+
+  private handleCompacted(event: CodexCompactedNotification): void {
+    if (!this.eventBelongsToThread(event)) return;
+    this._isCompacting = false;
+    this.emit({ type: "compaction_end" });
+  }
+
+  private handleError(event: { threadId?: string; error?: { message?: string; additionalDetails?: string | null } }): void {
+    if (!this.eventBelongsToThread(event)) return;
+    const message = event.error?.message || event.error?.additionalDetails || "Codex turn failed";
+    this._isStreaming = false;
+    this._activeTurnId = null;
+    this.emit({ type: "agent_end", error: message });
+  }
+
+  // ---- Actions ----
+
+  async prompt(message: string, images?: Array<{ data: string; mimeType: string }>): Promise<void> {
+    this._isStreaming = true;
+    this._accumulatedText = "";
+
+    // Emit agent_start so the UI knows we're running
+    this.emit({ type: "agent_start" });
+
+    // Start message
+    this.emit({
+      type: "message_start",
+      message: {
+        role: "assistant",
+        content: [],
+        ...this.assistantMeta(),
+      },
+    });
+
+    try {
+      const result = await this.connection.request("turn/start", {
+        threadId: this.threadId,
+        input: userInputsFromMessage(message, images),
+        ...(this.cwd ? { cwd: this.cwd } : {}),
+        ...(this._currentModel?.modelId ? { model: this._currentModel.modelId } : {}),
+        ...(this._currentModel?.provider ? { modelProvider: this._currentModel.provider } : {}),
+        approvalPolicy: "never",
+        sandboxPolicy: {
+          type: "workspaceWrite",
+          writableRoots: [this.cwd],
+          readOnlyAccess: { type: "fullAccess" },
+          networkAccess: true,
+          excludeTmpdirEnvVar: false,
+          excludeSlashTmp: false,
+        },
+      }) as { turn?: { id?: string } };
+      this._activeTurnId = result.turn?.id ?? null;
+    } catch (error) {
+      this._isStreaming = false;
+      this.emit({
+        type: "agent_end",
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+
+  async abort(): Promise<void> {
+    if (this._activeTurnId) {
+      try {
+        await this.connection.request("turn/interrupt", {
+          threadId: this.threadId,
+          turnId: this._activeTurnId,
+        });
+      } catch {
+        // local state still needs to settle even if the server rejected interrupt
+      }
+    }
+    this._isStreaming = false;
+    this._activeTurnId = null;
+  }
+
+  async setModel(provider: string, modelId: string): Promise<void> {
+    this._currentModel = { provider, modelId };
+  }
+
+  /**
+   * Fork the thread — create a new thread from the current one at an entry point.
+   */
+  async fork(): Promise<{ newThreadId: string } | { cancelled: true }> {
+    try {
+      const result = await this.connection.request("thread/fork", {
+        threadId: this.threadId,
+        persistExtendedHistory: true,
+      });
+
+      const newThreadId = (result as { thread?: { id?: string } }).thread?.id;
+      if (!newThreadId) return { cancelled: true };
+
+      return { newThreadId };
+    } catch {
+      return { cancelled: true };
+    }
+  }
+
+  /**
+   * Navigate (rollback) to a specific point in the thread.
+   */
+  async navigateTree(targetId: string): Promise<{ cancelled: boolean }> {
+    try {
+      const thread = await this.connection.request("thread/read", {
+        threadId: this.threadId,
+        includeTurns: true,
+      }) as { thread?: CodexThread };
+      const turns = thread.thread?.turns ?? [];
+      const targetTurnIndex = turns.findIndex((turn) => turn.items.some((item) => item.id === targetId));
+      if (targetTurnIndex < 0 || targetTurnIndex >= turns.length - 1) return { cancelled: false };
+      const numTurns = turns.length - targetTurnIndex - 1;
+      await this.connection.request("thread/rollback", {
+        threadId: this.threadId,
+        numTurns,
+      });
+      return { cancelled: false };
+    } catch {
+      return { cancelled: true };
+    }
+  }
+
+  /**
+   * Compact the thread context
+   */
+  async compact(): Promise<void> {
+    this._isCompacting = true;
+    this.emit({ type: "compaction_start" });
+
+    try {
+      await this.connection.request("thread/compact/start", {
+        threadId: this.threadId,
+      });
+    } catch (error) {
+      this._isCompacting = false;
+      this.emit({
+        type: "compaction_end",
+        errorMessage: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+
+  async rename(name: string): Promise<void> {
+    try {
+      await this.connection.request("thread/name/set", {
+        threadId: this.threadId,
+        name,
+      });
+    } catch {
+      // ignore
+    }
+  }
+
+  // ---- Lifecycle ----
+
+  destroy(): void {
+    if (!this._alive) return;
+    this._alive = false;
+    this._isStreaming = false;
+
+    for (const cleanup of this.cleanupFns) {
+      try { cleanup(); } catch { /* ignore */ }
+    }
+    this.cleanupFns = [];
+    this.listeners = [];
+  }
+}
+
+// ============================================================================
+// Session creation helpers
+// ============================================================================
+
+export interface CreateSessionOptions {
+  cwd: string;
+  message?: string;
+  baseInstructions?: string;
+  toolNames?: string[];
+  model?: { provider: string; modelId: string };
+  acceptReturnedModel?: boolean;
+  images?: Array<{ type: "image"; data: string; mimeType: string }>;
+}
+
+export async function createCodexSession(
+  connection: CodexServerConnection,
+  opts: CreateSessionOptions,
+): Promise<CodexSession> {
+  const result = await connection.request("thread/start", {
+    cwd: opts.cwd,
+    ...(opts.baseInstructions ? { baseInstructions: opts.baseInstructions } : {}),
+    ...(opts.model ? { model: opts.model.modelId, modelProvider: opts.model.provider } : {}),
+    approvalPolicy: "never",
+    sandbox: "workspace-write",
+    experimentalRawEvents: false,
+    persistExtendedHistory: true,
+  }) as { thread?: { id?: string }; model?: string; modelProvider?: string };
+
+  const threadId = result.thread?.id;
+  if (!threadId) {
+    throw new Error("Failed to create codex thread: no thread.id returned");
+  }
+
+  const session = new CodexSession(threadId, opts.cwd, connection, opts.baseInstructions);
+
+  if (opts.model) {
+    session.setModel(opts.model.provider, opts.model.modelId);
+  } else if (opts.acceptReturnedModel && result.model) {
+    session.setModel(result.modelProvider ?? "codex", result.model);
+  }
+
+  return session;
+}
+
+/**
+ * Resume an existing codex thread.
+ */
+export async function resumeCodexSession(
+  connection: CodexServerConnection,
+  threadId: string,
+  cwd: string,
+): Promise<CodexSession> {
+  // Read thread info to get model etc.
+  let model: { provider: string; modelId: string } | null = null;
+  const details = await connection.request("thread/read", { threadId, includeTurns: false }) as {
+    thread?: { modelProvider?: string };
+    model?: string;
+    modelProvider?: string;
+  };
+  if (details?.model) {
+    model = { provider: details.modelProvider ?? details.thread?.modelProvider ?? "codex", modelId: details.model };
+  }
+
+  const session = new CodexSession(threadId, cwd, connection);
+  if (model) {
+    session.setModel(model.provider, model.modelId);
+  }
+  return session;
+}