ThothII/backend/src/bridge/session-bridge.ts

import type { RpcClient } from "../rpc/rpc-client.js";

const GENERIC_MODEL_FAILURE =
  "Model request failed. Check provider connectivity, then Resume the session.";
const SUBSCRIPTION_MODEL_FAILURE =
  "The selected model is unavailable for the current subscription. Choose another model and start a new session.";
const PHASE_STARTED_NOTIFICATION_PREFIX = "__tht_phase_started__:";

function phaseStartedNotification(message: unknown): string | null {
  if (typeof message !== "string" || !message.startsWith(PHASE_STARTED_NOTIFICATION_PREFIX)) {
    return null;
  }
  const phase = message.slice(PHASE_STARTED_NOTIFICATION_PREFIX.length);
  return /^F[1-8]$/.test(phase) ? phase : "";
}

function safeModelFailure(error: unknown): string {
  const detail = typeof error === "string" ? error : "";
  const isSubscriptionFailure =
    /\b429\b/.test(detail) &&
    /(subscription plan|code["':\s]+1311|does not yet include access)/i.test(detail);

  return isSubscriptionFailure ? SUBSCRIPTION_MODEL_FAILURE : GENERIC_MODEL_FAILURE;
}

export type ToolActivity = {
  kind: "tool";
  toolCallId: string;
  toolName: string;
  status: "running" | "completed" | "failed";
};

export type TokenUsage = {
  input: number;
  cacheRead: number;
  output: number;
  totalTokens: number;
  contextWindow: number;
};

export type ClientEvent =
  | { type: "ui_request"; ui_request: any }
  | { type: "text_delta"; text: string }
  | { type: "activity_delta"; text: string }
  | { type: "activity_event"; activity: ToolActivity }
  | { type: "usage"; usage: TokenUsage }
  | { type: "info"; [k: string]: any }
  | { type: "system_event"; event: string; phase?: string };

export type TurnState = "idle" | "running" | "waiting" | "failed";

export class SessionBridge {
  private state: TurnState = "idle";
  private pending: any = null;
  // Pi (rpc-mode createDialogPromise) assegna a ogni ctx.ui.input un id RPC PROPRIO
  // (crypto.randomUUID) e correla extension_ui_response su quell'id — NON sull'id interno
  // del descriptor (che viaggia opaco nel `title`). Va memorizzato e rimandato indietro,
  // altrimenti Pi scarta la risposta e ctx.ui.input non si risolve mai (stuck senza output).
  private pendingPiId: string | null = null;
  private contextWindow = 0;
  private cbs = new Set<(e: ClientEvent) => void>();

  constructor(private rpc: RpcClient) {
    rpc.on("event", (m) => {
      if (m.type === "extension_ui_request" && m.method === "input") {
        let descriptor: unknown;
        try { descriptor = JSON.parse(m.title as string); } catch { return; }
        this.pending = descriptor;
        this.pendingPiId = m.id as string | null;
        this.state = "waiting";
        this.fan({ type: "ui_request", ui_request: descriptor });
      } else if (m.type === "message_end" && m.message?.role === "assistant") {
        this.emitUsage(m.message.usage);
        if (m.message.stopReason === "error") {
          this.markFailed();
          this.fan({
            type: "info",
            level: "error",
            text: safeModelFailure(m.message.errorMessage),
          });
        }
      } else if (m.type === "extension_ui_request" && m.method === "notify") {
        const phase = phaseStartedNotification(m.message);
        if (phase) this.fan({ type: "system_event", event: "phase_started", phase });
        else if (phase === null) {
          this.fan({ type: "info", level: m.notifyType ?? "info", text: m.message ?? "" });
        }
      } else if (m.type === "message_update" && m.assistantMessageEvent?.type === "text_delta") {
        this.fan({ type: "text_delta", text: m.assistantMessageEvent.delta ?? "" });
      } else if (m.type === "message_update" && m.assistantMessageEvent?.type === "thinking_delta") {
        this.fan({ type: "activity_delta", text: m.assistantMessageEvent.delta ?? "" });
      } else if (m.type === "text_delta") {
        this.fan({ type: "text_delta", text: m.text ?? "" });
      } else if (m.type === "tool_execution_start") {
        this.emitToolActivity(m, "start");
      } else if (m.type === "tool_execution_end") {
        this.emitToolActivity(m, "end");
      // Tool updates and raw payloads remain intentionally dropped.
      } else if (m.type === "system_event") {
        if (typeof m.event === "string" && m.event.trim() !== "") {
          this.fan({ type: "system_event", event: m.event });
        }
      } else if (m.type === "agent_end") {
        if (this.state !== "failed" && !this.pending) this.state = "idle";
        this.fan({ type: "system_event", event: "agent_end" });
      } else if (m.type === "agent_start") {
        this.state = "running";
        this.fan({ type: "system_event", event: "agent_start" });
      } else if (m.type === "turn_end") {
        this.fan({ type: "system_event", event: "turn_end" });
      }
    });
  }

  private emitToolActivity(message: any, lifecycle: "start" | "end"): void {
    const toolCallId = message.toolCallId;
    if (typeof toolCallId !== "string" || toolCallId.trim() === "") return;

    const toolName =
      typeof message.toolName === "string" && message.toolName.trim() !== ""
        ? message.toolName
        : "Tool";
    const status =
      lifecycle === "start" ? "running" : message.isError === true ? "failed" : "completed";

    this.fan({
      type: "activity_event",
      activity: { kind: "tool", toolCallId, toolName, status },
    });
  }

  private emitUsage(usage: any): void {
    if (!usage || this.contextWindow <= 0) return;
    const number = (value: unknown): number =>
      typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : 0;
    const input = number(usage.input);
    const cacheRead = number(usage.cacheRead);
    const output = number(usage.output);
    const reportedTotal = number(usage.totalTokens);
    this.fan({
      type: "usage",
      usage: {
        input,
        cacheRead,
        output,
        totalTokens: reportedTotal || input + cacheRead + output,
        contextWindow: this.contextWindow,
      },
    });
  }

  private fan(e: ClientEvent) { for (const cb of this.cbs) cb(e); }

  turnState(): TurnState { return this.state; }

  beginTurn(): void { this.state = "running"; }

  setContextWindow(value: unknown): void {
    if (typeof value === "number" && Number.isFinite(value) && value > 0) {
      this.contextWindow = value;
    }
  }

  /** Record a backend-detected failure; callers own any sanitized client message. */
  markFailed(): void { this.state = "failed"; }

  onClientEvent(cb: (e: ClientEvent) => void): void { this.cbs.add(cb); }

  /** Eventi generati dal backend stesso (es. exit inatteso del child Pi), non da Pi. */
  emitClientEvent(e: ClientEvent): void { this.fan(e); }

  /**
   * Deliver a reviewer response to Pi. Accepts ONLY a response matching the pending
   * descriptor: a stale/duplicate response (old gate id, double submit) would otherwise
   * be sent with the CURRENT gate's RPC id, flip the state to running, and leave the
   * real gate waiting. Returns false when rejected so the route can 409.
   */
  respond(uiResponse: object & { id: string }): boolean {
    if (!this.pending || uiResponse.id !== (this.pending as { id?: unknown }).id) return false;
    // Correla sull'id RPC di Pi; `value` porta l'uiResponse (con l'id del descriptor) cosi'
    // il check interno del gate (resp.id === descriptor.id) regge.
    const piId = this.pendingPiId ?? uiResponse.id;
    this.state = "running";
    this.rpc.send({ type: "extension_ui_response", id: piId, value: JSON.stringify(uiResponse) });
    this.pending = null;
    this.pendingPiId = null;
    return true;
  }

  steer(text: string): void {
    this.state = "running";
    this.rpc.send({ type: "steer", message: text });
  }

  pendingWidget(): object | null { return this.pending; }
}