@hassangameryt/freeflow 0.5.0 → 0.7.0

package/pi-extension/src/cognitive-routing/pi-routing-session.ts

@@ -0,0 +1,538 @@
+import {
+  COGNITIVE_ROUTING_INACTIVE_ENTRY,
+  CognitiveRoutingController,
+  type CognitiveRoutingControllerState,
+  type CognitiveRoutingMechanism,
+} from "./controller.js";
+import {
+  createPiSessionControllerHost,
+  parsePiSessionModelStateCommit,
+  PI_SESSION_MODEL_STATE_ENTRY,
+} from "./pi-session-control.js";
+import type {
+  CognitiveRoutingCapabilityState,
+  CognitiveRoutingProfileName,
+  CognitiveRoutingSessionStart,
+} from "./types.js";
+import { isPiFlowHost } from "../runtime/runtime-identity.js";
+import { readCapabilityState, refreshRuntimeContext } from "../runtime/runtime-context.js";
+
+export type ArmedCognitiveRoutingState = CognitiveRoutingSessionStart;
+
+export type ArmedCognitiveRoutingController = {
+  state(): CognitiveRoutingControllerState;
+  setManualProfile(
+    profile: CognitiveRoutingProfileName,
+    _mechanism?: string,
+  ): Promise<{ status: "active"; profile: CognitiveRoutingProfileName }>;
+  setAutomaticControl(_mechanism?: string): Promise<{ status: "automatic" }>;
+};
+
+export function armedCognitiveRoutingControllerState(
+  armed: Readonly<ArmedCognitiveRoutingState>,
+): CognitiveRoutingControllerState {
+  return {
+    effective: true,
+    controlMode: armed.control === "manual" ? `manual-${armed.profile}` : "automatic",
+    activeProfile: armed.profile,
+  };
+}
+
+export function armedCognitiveRoutingStateFor(
+  capabilityState: CognitiveRoutingCapabilityState,
+): ArmedCognitiveRoutingState {
+  return { ...capabilityState.sessionStart };
+}
+
+export function createArmedCognitiveRoutingController(
+  armed: ArmedCognitiveRoutingState,
+): ArmedCognitiveRoutingController {
+  return {
+    state: () => armedCognitiveRoutingControllerState(armed),
+    async setManualProfile(profile) {
+      armed.profile = profile;
+      armed.control = "manual";
+      return { status: "active", profile };
+    },
+    async setAutomaticControl() {
+      armed.control = "automatic";
+      return { status: "automatic" };
+    },
+  };
+}
+
+type SessionStartEvent = { reason?: string } | undefined;
+
+type SessionEntriesContext = {
+  sessionManager?: {
+    getBranch?: () => readonly unknown[];
+    getEntries?: () => readonly unknown[];
+  };
+};
+
+export type PiRoutingSessionDisposition = "available" | "inactive" | "blocked";
+
+type PiRoutingRuntimeState =
+  | CognitiveRoutingControllerState
+  | {
+      effective: false;
+      controlMode?: CognitiveRoutingControllerState["controlMode"];
+      activeProfile?: CognitiveRoutingProfileName;
+      runtimeStatus: "inactive" | "blocked";
+      runtimeReason?: string;
+    };
+
+export type PiRoutingSessionSnapshot = {
+  status: "available" | "active" | "armed" | "inactive" | "blocked";
+  disposition: PiRoutingSessionDisposition;
+  startupSelectionSuppressed: boolean;
+  reason?: string;
+  controllerState?: CognitiveRoutingControllerState;
+  runtimeState?: PiRoutingRuntimeState;
+};
+
+const THINKING_LEVELS = new Set(["off", "minimal", "low", "medium", "high", "xhigh", "max"]);
+
+function entriesFrom(ctx: SessionEntriesContext | undefined): readonly unknown[] {
+  try {
+    const entries = ctx?.sessionManager?.getBranch?.() ?? ctx?.sessionManager?.getEntries?.() ?? [];
+    return Array.isArray(entries) ? entries : [];
+  } catch {
+    return [];
+  }
+}
+
+function hasExplicitStartupSelection(argv: readonly string[]): boolean {
+  for (let index = 0; index < argv.length; index += 1) {
+    const argument = argv[index];
+    if (argument === "--model" && typeof argv[index + 1] === "string" && !argv[index + 1].startsWith("--")) {
+      return true;
+    }
+    if (argument.startsWith("--model=") && argument.length > "--model=".length) return true;
+    if (argument === "--thinking" && THINKING_LEVELS.has(argv[index + 1])) return true;
+    if (argument.startsWith("--thinking=") && THINKING_LEVELS.has(argument.slice("--thinking=".length))) {
+      return true;
+    }
+  }
+  return false;
+}
+
+function hasInactiveEntry(entries: readonly unknown[]): boolean {
+  return entries.some(
+    (entry) =>
+      Boolean(entry && typeof entry === "object") &&
+      (entry as { type?: unknown; customType?: unknown }).type === "custom" &&
+      (entry as { customType?: unknown }).customType === COGNITIVE_ROUTING_INACTIVE_ENTRY,
+  );
+}
+
+function hasNativeOverrideAfterRoutingState(entries: readonly unknown[]): boolean {
+  let latestNativePosition = -1;
+  let latestCommittedRoutingPosition = -1;
+  let latestInactivePosition = -1;
+  let modelChangeCount = 0;
+  let thinkingChangeCount = 0;
+
+  entries.forEach((entry: unknown, position) => {
+    if (!entry || typeof entry !== "object") return;
+    const record = entry as Record<string, unknown>;
+    if (record.type === "model_change") {
+      latestNativePosition = position;
+      modelChangeCount += 1;
+      return;
+    }
+    if (record.type === "thinking_level_change") {
+      latestNativePosition = position;
+      thinkingChangeCount += 1;
+      return;
+    }
+    if (record.type !== "custom") return;
+    if (record.customType === COGNITIVE_ROUTING_INACTIVE_ENTRY) {
+      latestInactivePosition = position;
+      return;
+    }
+    if (record.customType !== PI_SESSION_MODEL_STATE_ENTRY) return;
+    const commit = parsePiSessionModelStateCommit(record.data);
+    if (commit?.phase === "committed" && commit.status === "applied") {
+      latestCommittedRoutingPosition = position;
+    }
+  });
+
+  if (latestInactivePosition > latestCommittedRoutingPosition) return true;
+  if (latestCommittedRoutingPosition >= 0) return latestNativePosition > latestCommittedRoutingPosition;
+  return modelChangeCount > 1 || thinkingChangeCount > 1;
+}
+
+function hasPiSessionModelStateApi(pi: any): boolean {
+  return (
+    typeof pi?.appendEntry === "function" &&
+    typeof pi?.setModel === "function" &&
+    typeof pi?.setThinkingLevel === "function"
+  );
+}
+
+function hasCognitiveRoutingHost(pi: any, ctx: any): boolean {
+  if (
+    isPiFlowHost(pi?.host) &&
+    typeof pi?.appendEntryDurable === "function" &&
+    typeof pi?.acquireModelStateControl === "function"
+  ) {
+    return true;
+  }
+  return (
+    hasPiSessionModelStateApi(pi) &&
+    typeof ctx?.modelRegistry?.find === "function" &&
+    typeof ctx?.modelRegistry?.getApiKeyAndHeaders === "function"
+  );
+}
+
+function cognitiveRoutingControllerHost(pi: any, ctx: any): any | undefined {
+  if (isPiFlowHost(pi?.host) && typeof pi?.acquireModelStateControl === "function") return pi;
+  if (hasPiSessionModelStateApi(pi) && hasCognitiveRoutingHost(pi, ctx)) {
+    return createPiSessionControllerHost({ pi, ctx });
+  }
+  return undefined;
+}
+
+export class PiRoutingSession {
+  private disposition: PiRoutingSessionDisposition = "available";
+  private reason: string | undefined;
+  private startupSelectionSuppressed = false;
+  private readonly pi: any;
+  private readonly stockPi: boolean;
+  private newSessionBoundary = false;
+  private explicitReactivationPending = false;
+  private controllerState: CognitiveRoutingController | undefined;
+  private armedState: ArmedCognitiveRoutingState | undefined;
+
+  get controller(): CognitiveRoutingController | undefined {
+    return this.controllerState;
+  }
+
+  controllerForTool(): CognitiveRoutingController | undefined {
+    return this.controllerState;
+  }
+
+  hasController(): boolean {
+    return this.controllerState !== undefined;
+  }
+
+  async reconcileLiveState(reason = "Pi model state diverged"): Promise<any> {
+    return this.controllerState?.reconcileLiveState(reason) ?? { status: "ignored" };
+  }
+
+  async shutdown(reason?: string): Promise<any> {
+    const controller = this.controllerState;
+    this.reset();
+    return controller?.shutdown(reason as any);
+  }
+
+  get armed(): Readonly<ArmedCognitiveRoutingState> | undefined {
+    return this.armedState ? { ...this.armedState } : undefined;
+  }
+
+  clearController(): void {
+    this.controllerState = undefined;
+  }
+
+  clearArmed(): void {
+    this.armedState = undefined;
+  }
+
+  armFromCapability(capabilityState: CognitiveRoutingCapabilityState): void {
+    this.armedState = armedCognitiveRoutingStateFor(capabilityState);
+  }
+
+  cycleArmedManual(): { target: CognitiveRoutingProfileName } | undefined {
+    if (!this.armedState) return undefined;
+    const target = this.armedState.profile === "reasoning" ? "standard" : "reasoning";
+    this.armedState.profile = target;
+    this.armedState.control = "manual";
+    return { target };
+  }
+
+  cycleArmedAutomatic(): { kind: "control"; changed: boolean } | undefined {
+    if (!this.armedState) return undefined;
+    if (this.armedState.control === "automatic") return { kind: "control", changed: false };
+    this.armedState.control = "automatic";
+    return { kind: "control", changed: true };
+  }
+
+  armedController(): ArmedCognitiveRoutingController | undefined {
+    return this.armedState ? createArmedCognitiveRoutingController(this.armedState) : undefined;
+  }
+
+  commandController(): ArmedCognitiveRoutingController | CognitiveRoutingController | undefined {
+    return this.controllerState ?? this.armedController();
+  }
+
+  history(options: any = {}): any {
+    return this.controllerState?.history(options);
+  }
+
+  async observeNativeChange(reason: string, source: "set" | "cycle" | "restore" | undefined = undefined): Promise<any> {
+    return this.controllerState?.observeNativeChange(reason, source) ?? { status: "ignored" };
+  }
+
+  async setManualProfile(
+    profile: CognitiveRoutingProfileName,
+    mechanism: CognitiveRoutingMechanism = "profile-command",
+  ): Promise<any> {
+    if (this.controllerState) return this.controllerState.setManualProfile(profile, mechanism);
+    return this.armedController()?.setManualProfile(profile, mechanism) ?? { status: "inactive", reason: "not_active" };
+  }
+
+  async setAutomaticControl(mechanism: CognitiveRoutingMechanism = "profile-command"): Promise<any> {
+    if (this.controllerState) return this.controllerState.setAutomaticControl(mechanism);
+    return this.armedController()?.setAutomaticControl(mechanism) ?? { status: "inactive", reason: "not_active" };
+  }
+
+  async switchAutomaticProfile(
+    profile: CognitiveRoutingProfileName,
+    reason: string,
+    source: "agent" | "system" | "user" = "user",
+    mechanism: CognitiveRoutingMechanism = "profile-shortcut",
+  ): Promise<any> {
+    return (
+      this.controllerState?.switchAutomaticProfile(profile, reason, source, mechanism) ?? {
+        status: "blocked",
+        reason: "not_active",
+      }
+    );
+  }
+
+  async reconcileBranch(mechanism: CognitiveRoutingMechanism = "session-tree"): Promise<any> {
+    return this.controllerState?.reconcileBranch(mechanism);
+  }
+
+  constructor({ pi, stockPi }: { pi: any; stockPi: boolean }) {
+    this.pi = pi;
+    this.stockPi = stockPi;
+  }
+
+  snapshot(): PiRoutingSessionSnapshot {
+    const controllerState = this.controllerState?.state();
+    const armedRuntimeState = this.armedState ? armedCognitiveRoutingControllerState(this.armedState) : undefined;
+    const runtimeState = controllerState ?? armedRuntimeState ?? this.runtimeState();
+    const status =
+      this.disposition === "blocked"
+        ? "blocked"
+        : this.disposition === "inactive"
+          ? "inactive"
+          : controllerState?.runtimeStatus === "blocked"
+            ? "blocked"
+            : controllerState?.runtimeStatus === "inactive"
+              ? "inactive"
+              : controllerState?.effective === true
+                ? "active"
+                : this.armedState
+                  ? "armed"
+                  : "available";
+    return {
+      status,
+      disposition: this.disposition,
+      startupSelectionSuppressed: this.startupSelectionSuppressed,
+      ...(this.reason ? { reason: this.reason } : {}),
+      ...(controllerState ? { controllerState } : {}),
+      ...(runtimeState ? { runtimeState } : {}),
+    };
+  }
+
+  supportsCognitiveRoutingRuntime(): boolean {
+    return isPiFlowHost(this.pi?.host) || hasPiSessionModelStateApi(this.pi);
+  }
+
+  async reconcileController(
+    ctx: any,
+    capabilityState: any,
+    initialSessionStart?: CognitiveRoutingSessionStart,
+  ): Promise<CognitiveRoutingController | undefined> {
+    const cognitiveRouting = capabilityState?.cognitiveRouting;
+    const controllerPi = cognitiveRoutingControllerHost(this.pi, ctx);
+    if (!cognitiveRouting?.effective || this.startupSelectionSuppressed || !controllerPi) {
+      if (this.controllerState) await this.controllerState.deactivate();
+      this.controllerState = undefined;
+      return undefined;
+    }
+
+    const controller = new CognitiveRoutingController({
+      capabilityState: cognitiveRouting,
+      pi: controllerPi,
+      ctx,
+      liveStateAuthoritative: !isPiFlowHost(this.pi?.host),
+      ...(initialSessionStart ? { initialSessionStart } : {}),
+    });
+    // Pi emits model/thinking events while setModel() is still awaiting its transition.
+    // Publish the controller first so those events can observe transitionInFlight and ignore
+    // host notifications generated by this activation rather than treating them as overrides.
+    this.controllerState = controller;
+    try {
+      const recovered = await controller.recover();
+      let nextController: CognitiveRoutingController | undefined;
+      if (
+        recovered.status === "pending" ||
+        (recovered.status === "inactive" && recovered.reason === "native_override")
+      ) {
+        nextController = controller;
+      } else if (recovered.status === "active") {
+        nextController = controller;
+      } else {
+        const activated = await controller.activate();
+        nextController = activated.status === "active" ? controller : undefined;
+      }
+      this.controllerState = nextController;
+      return nextController;
+    } catch (error) {
+      if (this.controllerState === controller) this.controllerState = undefined;
+      throw error;
+    }
+  }
+
+  async applyLiveCapabilityState(
+    ctx: any,
+    options: { reconcileCognitiveRouting?: boolean; materializeCognitiveRouting?: boolean } = {},
+  ): Promise<CognitiveRoutingController | undefined> {
+    const capabilityState = await readCapabilityState(ctx.cwd, ctx, this.pi?.host);
+    await refreshRuntimeContext(capabilityState);
+    this.reconcileNativeEntries(ctx);
+    let nextController = this.controllerState;
+    if (!isPiFlowHost(this.pi?.host) && nextController) await nextController.reconcileLiveState();
+    if (
+      options.reconcileCognitiveRouting &&
+      (this.armedState === undefined || options.materializeCognitiveRouting === true) &&
+      (nextController === undefined || capabilityState.cognitiveRouting?.effective !== true)
+    ) {
+      nextController = await this.reconcileController(
+        ctx,
+        capabilityState,
+        options.materializeCognitiveRouting ? this.armedState : undefined,
+      );
+    }
+    this.controllerState = nextController;
+    const controllerState = nextController?.state();
+    if (controllerState?.effective === true) {
+      this.markActive();
+    } else if (this.isExplicitReactivationPending() && options.reconcileCognitiveRouting) {
+      this.cancelExplicitReactivation();
+    }
+    return nextController;
+  }
+
+  runtimeState(): PiRoutingRuntimeState | undefined {
+    if (this.disposition === "inactive") {
+      return { effective: false, runtimeStatus: "inactive", ...(this.reason ? { runtimeReason: this.reason } : {}) };
+    }
+    if (this.disposition === "blocked") {
+      return { effective: false, runtimeStatus: "blocked", ...(this.reason ? { runtimeReason: this.reason } : {}) };
+    }
+    return undefined;
+  }
+
+  canAutoActivate(): boolean {
+    return !this.startupSelectionSuppressed && (this.disposition === "available" || this.explicitReactivationPending);
+  }
+
+  isExplicitReactivationPending(): boolean {
+    return this.explicitReactivationPending;
+  }
+
+  resetForSession(ctx: SessionEntriesContext | undefined, event: SessionStartEvent, argv = process.argv): void {
+    const persistedInactive = hasInactiveEntry(entriesFrom(ctx));
+    this.disposition = "available";
+    this.reason = undefined;
+    this.newSessionBoundary = event?.reason !== "reload";
+    this.explicitReactivationPending = false;
+    this.armedState = undefined;
+    const hostStartupSelection =
+      !this.stockPi &&
+      ((ctx as any)?.modelStateProvenance?.explicitModel === true ||
+        (ctx as any)?.modelStateProvenance?.explicitThinking === true);
+    this.startupSelectionSuppressed =
+      hostStartupSelection || (hasExplicitStartupSelection(argv) && (this.newSessionBoundary || !persistedInactive));
+  }
+
+  ensureStartupInactive(ctx: SessionEntriesContext | undefined, effective: boolean): void {
+    if (!this.stockPi || !this.startupSelectionSuppressed || !effective) return;
+    const entries = entriesFrom(ctx);
+    const shouldPersistFresh = this.newSessionBoundary || !hasInactiveEntry(entries);
+    if (shouldPersistFresh && !this.persistInactive(ctx, "Pi startup model or thinking selection")) {
+      this.disposition = "blocked";
+      this.reason = "inactive_state_persistence_failed";
+      return;
+    }
+    this.disposition = "inactive";
+    this.reason = "Pi startup model or thinking selection";
+  }
+
+  reconcileNativeEntries(ctx: SessionEntriesContext | undefined): void {
+    if (!this.stockPi || this.startupSelectionSuppressed || this.explicitReactivationPending) return;
+    const entries = entriesFrom(ctx);
+    if (!hasNativeOverrideAfterRoutingState(entries)) return;
+    this.disposition = "inactive";
+    this.reason = "Pi model state override during recovery";
+    if (!hasInactiveEntry(entries) && !this.persistInactive(ctx, this.reason)) {
+      this.disposition = "blocked";
+      this.reason = "inactive_state_persistence_failed";
+    }
+  }
+
+  observeArmedNativeChange(ctx: SessionEntriesContext | undefined, reason: string): void {
+    if (!this.stockPi) return;
+    this.disposition = "inactive";
+    this.reason = reason;
+    if (!this.persistInactive(ctx, reason)) {
+      this.disposition = "blocked";
+      this.reason = "inactive_state_persistence_failed";
+    }
+  }
+
+  beginExplicitReactivation(): void {
+    this.startupSelectionSuppressed = false;
+    this.explicitReactivationPending = true;
+  }
+
+  cancelExplicitReactivation(): void {
+    this.explicitReactivationPending = false;
+  }
+
+  markActive(): void {
+    this.startupSelectionSuppressed = false;
+    this.explicitReactivationPending = false;
+    this.disposition = "available";
+    this.reason = undefined;
+  }
+
+  reset(): void {
+    this.startupSelectionSuppressed = false;
+    this.explicitReactivationPending = false;
+    this.disposition = "available";
+    this.reason = undefined;
+    this.newSessionBoundary = false;
+    this.controllerState = undefined;
+    this.armedState = undefined;
+  }
+
+  private persistInactive(ctx: SessionEntriesContext | undefined, reason: string): boolean {
+    const candidate = (ctx as any)?.model;
+    const thinkingLevel = (ctx as any)?.thinkingLevel;
+    if (
+      !candidate ||
+      typeof candidate.provider !== "string" ||
+      typeof candidate.id !== "string" ||
+      typeof thinkingLevel !== "string"
+    ) {
+      return false;
+    }
+    try {
+      this.pi.appendEntry(COGNITIVE_ROUTING_INACTIVE_ENTRY, {
+        version: 1,
+        source: "user",
+        reason,
+        pair: { provider: candidate.provider, modelId: candidate.id, thinkingLevel },
+      });
+      return true;
+    } catch {
+      return false;
+    }
+  }
+}