oh-my-pi-telegram/lib/lifecycle.ts

/**
 * Telegram lifecycle hook registration helpers
 * Zones: pi agent lifecycle, telegram session
 * Binds prepared Telegram lifecycle runtimes to pi extension lifecycle events
 */

import * as BusFollower from "./bus-follower.ts";
import * as Queue from "./queue.ts";
import type {
  AgentEndEvent,
  AgentSettledEvent,
  AgentStartEvent,
  AssistantMessageEvent,
  BeforeAgentStartEvent,
  ExtensionAPI,
  ExtensionContext,
  InputEvent,
  SessionBeforeCompactEvent,
  SessionCompactEvent,
  SessionShutdownEvent,
  SessionStartEvent,
  ToolExecutionEndEvent,
  ToolExecutionStartEvent,
  ToolExecutionUpdateEvent,
} from "./pi.ts";
import { isTerminalAgentEnd } from "./pi.ts";

let resetTransportReplyDedupFn: (() => void) | undefined;

export function setResetTransportReplyDedup(fn: () => void): void {
  resetTransportReplyDedupFn = fn;
}

export function createAgentStartDedupHook(
  inner: (event: AgentStartEvent, ctx: ExtensionContext) => Promise<void>,
): (event: AgentStartEvent, ctx: ExtensionContext) => Promise<void> {
  return async (event, ctx) => {
    if (resetTransportReplyDedupFn) resetTransportReplyDedupFn();
    return inner(event, ctx);
  };
}

export interface TelegramBeforeAgentStartResult {
  systemPrompt?: string;
}

type TelegramBeforeAgentStartReturn =
  | Promise<TelegramBeforeAgentStartResult | undefined>
  | TelegramBeforeAgentStartResult
  | undefined;

type TelegramLifecycleModel = ExtensionContext["model"];
type TelegramLifecycleMessage = AgentEndEvent["messages"][number];

export interface TelegramLifecycleRegistrationDeps {
  isSessionActive?: (ctx: ExtensionContext) => boolean;
  onInput?: (event: InputEvent, ctx: ExtensionContext) => Promise<void> | void;
  onSessionStart: (
    event: SessionStartEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
  onSessionShutdown: (
    event: SessionShutdownEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
  onSessionBeforeCompact?: (
    event: SessionBeforeCompactEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onSessionCompact?: (
    event: SessionCompactEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onBeforeAgentStart: (
    event: BeforeAgentStartEvent,
    ctx: ExtensionContext,
  ) => TelegramBeforeAgentStartReturn;
  onModelSelect: (
    event: { model: TelegramLifecycleModel },
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onAgentStart: (
    event: AgentStartEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
  onToolExecutionStart: (
    event: ToolExecutionStartEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onToolExecutionUpdate?: (
    event: ToolExecutionUpdateEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onToolExecutionEnd: (
    event: ToolExecutionEndEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  onMessageStart: (
    event: { message: TelegramLifecycleMessage },
    ctx: ExtensionContext,
  ) => Promise<void>;
  onMessageUpdate: (
    event: {
      message: TelegramLifecycleMessage;
      assistantMessageEvent?: AssistantMessageEvent;
    },
    ctx: ExtensionContext,
  ) => Promise<void>;
  onAgentEnd: (event: AgentEndEvent, ctx: ExtensionContext) => Promise<void>;
  onAgentSettled?: (
    event: AgentSettledEvent,
    ctx: ExtensionContext,
  ) => Promise<void> | void;
  /**
   * Hosts that never emit `agent_settled` (omp) get it synthesized from the
   * first terminal `agent_end`. Leave unset on hosts that emit it natively so
   * the settle point is not announced early.
   */
  settleOnTerminalAgentEnd?: boolean;
}

export interface TelegramSessionLifecycleHooks {
  onSessionStart: (
    event: SessionStartEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
  onSessionShutdown: (
    event: SessionShutdownEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
}

export interface TelegramSessionContextStore<TContext> {
  get: () => TContext | undefined;
  getGeneration: () => number;
  isCurrent: (ctx: TContext, generation?: number) => boolean;
  set: (ctx: TContext) => number;
  clear: (ctx?: TContext) => boolean;
}

export function createTelegramSessionContextStore<TContext>(
  options: {
    getIdentity?: (ctx: TContext) => unknown;
  } = {},
): TelegramSessionContextStore<TContext> {
  let currentContext: TContext | undefined;
  let currentIdentity: unknown;
  let generation = 0;
  const objectIdentityGenerations = new WeakMap<object, number>();
  const primitiveIdentityGenerations = new Map<unknown, number>();
  const resolveIdentity = (ctx: TContext): unknown =>
    options.getIdentity?.(ctx) ?? ctx;
  const getIdentityGeneration = (identity: unknown): number | undefined =>
    (typeof identity === "object" && identity !== null) ||
    typeof identity === "function"
      ? objectIdentityGenerations.get(identity as object)
      : primitiveIdentityGenerations.get(identity);
  const setIdentityGeneration = (identity: unknown, value: number): void => {
    if (
      (typeof identity === "object" && identity !== null) ||
      typeof identity === "function"
    ) {
      objectIdentityGenerations.set(identity as object, value);
    } else {
      primitiveIdentityGenerations.set(identity, value);
    }
  };
  return {
    get: () => currentContext,
    getGeneration: () => generation,
    isCurrent: (ctx, expectedGeneration) => {
      if (currentContext === undefined) return false;
      const identity = resolveIdentity(ctx);
      return (
        identity === currentIdentity &&
        getIdentityGeneration(identity) === generation &&
        (expectedGeneration === undefined || generation === expectedGeneration)
      );
    },
    set: (ctx) => {
      currentContext = ctx;
      currentIdentity = resolveIdentity(ctx);
      generation += 1;
      setIdentityGeneration(currentIdentity, generation);
      return generation;
    },
    clear: (ctx) => {
      if (ctx !== undefined) {
        const identity = resolveIdentity(ctx);
        if (
          identity !== currentIdentity ||
          getIdentityGeneration(identity) !== generation
        ) {
          return false;
        }
      }
      if (currentContext === undefined) return false;
      currentContext = undefined;
      currentIdentity = undefined;
      generation += 1;
      return true;
    },
  };
}

export function createTelegramSessionContextTracker(
  store: Pick<TelegramSessionContextStore<ExtensionContext>, "set" | "clear">,
): TelegramSessionLifecycleHooks {
  return {
    onSessionStart: async (_event, ctx) => {
      store.set(ctx);
    },
    onSessionShutdown: async (_event, ctx) => {
      store.clear(ctx);
    },
  };
}

export function createTelegramSessionGenerationFence(
  store: TelegramSessionContextStore<ExtensionContext>,
  hooks: TelegramSessionLifecycleHooks,
): TelegramSessionLifecycleHooks {
  return {
    async onSessionStart(event, ctx) {
      const generation = store.set(ctx);
      await hooks.onSessionStart(event, ctx);
      if (!store.isCurrent(ctx, generation)) return;
    },
    async onSessionShutdown(event, ctx) {
      const generation = store.getGeneration();
      if (!store.isCurrent(ctx, generation)) return;
      await hooks.onSessionShutdown(event, ctx);
      store.clear(ctx);
    },
  };
}

export interface TelegramBridgeSessionServiceRuntime {
  resumeGroupedInput(ctx: ExtensionContext): void;
  suspendGroupedInput(): void;
  delivery: {
    onSessionStart(): Promise<void>;
    onSessionShutdown(): Promise<void>;
  };
  polling: {
    onSessionStart(
      event: SessionStartEvent,
      ctx: ExtensionContext,
    ): Promise<void>;
  };
  capabilityMonitor: { start(ctx: ExtensionContext): void; stop(): void };
  queueWatchdog: { start(ctx: ExtensionContext): void; stop(): void };
}

export interface TelegramBridgeSessionLifecycleAssemblyDeps<
  TQueueItem,
  TModel = unknown,
> {
  contextStore: TelegramSessionContextStore<ExtensionContext>;
  queue: Omit<
    Queue.TelegramSessionLifecycleRuntimeDeps<
      ExtensionContext,
      TQueueItem,
      TModel
    >,
    "isSessionActive" | "stopPolling" | "clearPendingMediaGroups"
  >;
  follower: Omit<
    BusFollower.TelegramBusFollowerSessionRefreshHookDeps<ExtensionContext>,
    "isSessionActive"
  > &
    BusFollower.TelegramBusFollowerSessionReplacementSuspenderDeps;
  services: TelegramBridgeSessionServiceRuntime;
}

export function createTelegramBridgeSessionLifecycleAssembly<
  TQueueItem,
  TModel = unknown,
>(
  deps: TelegramBridgeSessionLifecycleAssemblyDeps<TQueueItem, TModel>,
): TelegramSessionLifecycleHooks {
  const isSessionActive = deps.contextStore.isCurrent;
  const suspendForReplacement =
    BusFollower.createTelegramBusFollowerSessionReplacementSuspender({
      registrationState: deps.follower.registrationState,
      instanceId: deps.follower.instanceId,
      suspendPolling: deps.follower.suspendPolling,
      recordRuntimeEvent: deps.follower.recordRuntimeEvent,
    });
  const queueLifecycle = Queue.createTelegramSessionLifecycleRuntime({
    ...deps.queue,
    isSessionActive,
    stopPolling: suspendForReplacement,
    clearPendingMediaGroups: deps.services.suspendGroupedInput,
  });
  const servicesLifecycle = appendTelegramLifecycleHooks(
    queueLifecycle,
    {
      async onSessionStart(event, ctx) {
        deps.services.resumeGroupedInput(ctx);
        await deps.services.delivery.onSessionStart();
        await deps.services.polling.onSessionStart(event, ctx);
        deps.services.capabilityMonitor.start(ctx);
        deps.services.queueWatchdog.start(ctx);
      },
      async onSessionShutdown() {
        await deps.services.delivery.onSessionShutdown();
        deps.services.queueWatchdog.stop();
        deps.services.capabilityMonitor.stop();
      },
    },
    isSessionActive,
  );
  const followerLifecycle = appendTelegramLifecycleHooks(
    servicesLifecycle,
    {
      onSessionStart: BusFollower.createTelegramBusFollowerSessionRefreshHook({
        registrationState: deps.follower.registrationState,
        registrationRuntime: deps.follower.registrationRuntime,
        getLeaderState: deps.follower.getLeaderState,
        isSessionActive,
        updateStatus: deps.follower.updateStatus,
        recordRuntimeEvent: deps.follower.recordRuntimeEvent,
      }),
    },
    isSessionActive,
  );
  return createTelegramSessionGenerationFence(
    deps.contextStore,
    followerLifecycle,
  );
}

type TelegramLifecycleTimer = number | ReturnType<typeof setTimeout>;

function unrefTelegramLifecycleTimer(timer: TelegramLifecycleTimer): void {
  if (!timer || typeof timer !== "object") return;
  if (typeof timer.unref === "function") timer.unref();
}

export interface TelegramCompactionObserverRuntimeDeps<TContext> {
  isContextActive?: (ctx: TContext) => boolean;
  setCompactionInProgress: (inProgress: boolean) => void;
  updateStatus: (ctx: TContext) => void;
  startTypingLoop?: (ctx: TContext) => boolean | void;
  stopTypingLoop?: () => void;
  requestDeferredDispatchNextQueuedTelegramTurn: (
    dispatch: (ctx: TContext) => void,
  ) => void;
  dispatchNextQueuedTelegramTurn: (ctx: TContext) => void;
  recordRuntimeEvent?: (category: string, error: unknown) => void;
  onCompactionAbandoned?: () => void;
  timeoutMs?: number;
  setTimer?: (callback: () => void, ms: number) => TelegramLifecycleTimer;
  clearTimer?: (timer: TelegramLifecycleTimer) => void;
}

export interface TelegramCompactionObserverRuntime<TContext> {
  onSessionBeforeCompact: (
    event: SessionBeforeCompactEvent,
    ctx: TContext,
  ) => void;
  onSessionCompact: (event: SessionCompactEvent, ctx: TContext) => void;
  onSessionShutdown: () => void;
}

export function createTelegramCompactionObserverRuntime<TContext>(
  deps: TelegramCompactionObserverRuntimeDeps<TContext>,
): TelegramCompactionObserverRuntime<TContext> {
  const timeoutMs = deps.timeoutMs ?? 300_000;
  const setTimer = deps.setTimer ?? setTimeout;
  const clearTimer = deps.clearTimer ?? clearTimeout;
  let fallbackTimer: TelegramLifecycleTimer | undefined;
  let typingStartedByObserver = false;
  const clearFallbackTimer = (): void => {
    if (!fallbackTimer) return;
    clearTimer(fallbackTimer);
    fallbackTimer = undefined;
  };
  const requestDispatch = (): void => {
    deps.requestDeferredDispatchNextQueuedTelegramTurn(
      deps.dispatchNextQueuedTelegramTurn,
    );
  };
  return {
    onSessionBeforeCompact: (_event, ctx) => {
      if (deps.isContextActive && !deps.isContextActive(ctx)) return;
      deps.setCompactionInProgress(true);
      const typingStartResult = deps.startTypingLoop?.(ctx);
      typingStartedByObserver =
        !!deps.startTypingLoop && typingStartResult !== false;
      deps.updateStatus(ctx);
      clearFallbackTimer();
      fallbackTimer = setTimer(() => {
        fallbackTimer = undefined;
        if (deps.isContextActive && !deps.isContextActive(ctx)) return;
        deps.setCompactionInProgress(false);
        if (typingStartedByObserver) deps.stopTypingLoop?.();
        typingStartedByObserver = false;
        deps.updateStatus(ctx);
        deps.recordRuntimeEvent?.(
          "compact",
          new Error("Compaction observer timed out"),
        );
        deps.onCompactionAbandoned?.();
        requestDispatch();
      }, timeoutMs);
      unrefTelegramLifecycleTimer(fallbackTimer);
    },
    onSessionCompact: (_event, ctx) => {
      clearFallbackTimer();
      if (deps.isContextActive && !deps.isContextActive(ctx)) return;
      deps.setCompactionInProgress(false);
      if (typingStartedByObserver) deps.stopTypingLoop?.();
      typingStartedByObserver = false;
      deps.updateStatus(ctx);
      requestDispatch();
    },
    onSessionShutdown: () => {
      clearFallbackTimer();
      if (typingStartedByObserver) deps.stopTypingLoop?.();
      typingStartedByObserver = false;
    },
  };
}

export interface TelegramMessageActivityTypingDeps<TContext> {
  hasActiveTurn: () => boolean;
  startTypingLoop: (ctx: TContext) => void;
  onMessageStart: TelegramLifecycleRegistrationDeps["onMessageStart"];
  onMessageUpdate: TelegramLifecycleRegistrationDeps["onMessageUpdate"];
  recordRuntimeEvent?: (
    category: string,
    error: unknown,
    details?: Record<string, unknown>,
  ) => void;
}

export function createTelegramMessageActivityTypingHooks<
  TContext extends ExtensionContext,
>(
  deps: TelegramMessageActivityTypingDeps<TContext>,
): Pick<
  TelegramLifecycleRegistrationDeps,
  "onMessageStart" | "onMessageUpdate"
> {
  const ensureTyping = (ctx: TContext): void => {
    if (deps.hasActiveTurn()) deps.startTypingLoop(ctx);
  };
  const handleMessageActivity = async <TEvent>(
    phase: "start" | "update",
    event: TEvent,
    ctx: ExtensionContext,
    inner: (event: TEvent, ctx: ExtensionContext) => Promise<void>,
  ): Promise<void> => {
    const typedCtx = ctx as TContext;
    ensureTyping(typedCtx);
    try {
      await inner(event, ctx);
    } catch (error) {
      deps.recordRuntimeEvent?.("message-activity", error, { phase });
    } finally {
      ensureTyping(typedCtx);
    }
  };
  return {
    onMessageStart: (event, ctx) =>
      handleMessageActivity("start", event, ctx, deps.onMessageStart),
    onMessageUpdate: (event, ctx) =>
      handleMessageActivity("update", event, ctx, deps.onMessageUpdate),
  };
}

export function createDedupAgentStartHook(
  dedup: { reset(): void },
  inner: (event: AgentStartEvent, ctx: ExtensionContext) => Promise<void>,
): (event: AgentStartEvent, ctx: ExtensionContext) => Promise<void> {
  return async (event, ctx) => {
    dedup.reset();
    await inner(event, ctx);
  };
}

export interface TelegramExtraLifecycleHooks {
  onSessionStart?: (
    event: SessionStartEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
  onSessionShutdown?: (
    event: SessionShutdownEvent,
    ctx: ExtensionContext,
  ) => Promise<void>;
}

export function appendTelegramLifecycleHooks(
  base: TelegramSessionLifecycleHooks,
  extra: TelegramExtraLifecycleHooks,
  isSessionActive?: (ctx: ExtensionContext) => boolean,
): TelegramSessionLifecycleHooks {
  return {
    onSessionStart: async (event, ctx) => {
      await base.onSessionStart(event, ctx);
      if (isSessionActive?.(ctx) === false) return;
      await extra.onSessionStart?.(event, ctx);
    },
    onSessionShutdown: async (event, ctx) => {
      await base.onSessionShutdown(event, ctx);
      if (isSessionActive?.(ctx) === false) return;
      await extra.onSessionShutdown?.(event, ctx);
    },
  };
}

export function registerTelegramLifecycleHooks(
  pi: ExtensionAPI,
  deps: TelegramLifecycleRegistrationDeps,
): void {
  const isActive = (ctx: ExtensionContext): boolean =>
    deps.isSessionActive?.(ctx) !== false;
  pi.on("input", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onInput?.(event, ctx);
  });
  pi.on("session_start", async (event, ctx) => {
    await deps.onSessionStart(event, ctx);
  });
  pi.on("session_shutdown", async (event, ctx) => {
    await deps.onSessionShutdown(event, ctx);
  });
  pi.on("session_before_compact", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onSessionBeforeCompact?.(event, ctx);
  });
  pi.on("session_compact", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onSessionCompact?.(event, ctx);
  });
  pi.on("before_agent_start", async (event, ctx) => {
    return deps.onBeforeAgentStart(event, ctx);
  });
  pi.on("model_select", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onModelSelect(event, ctx);
  });
  pi.on("agent_start", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onAgentStart(event, ctx);
  });
  pi.on("tool_execution_start", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onToolExecutionStart(event, ctx);
  });
  pi.on("tool_execution_update", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onToolExecutionUpdate?.(event, ctx);
  });
  pi.on("tool_execution_end", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onToolExecutionEnd(event, ctx);
  });
  pi.on("message_start", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onMessageStart(event, ctx);
  });
  pi.on("message_update", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onMessageUpdate(event, ctx);
  });
  pi.on("agent_end", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onAgentEnd(event, ctx);
    if (deps.settleOnTerminalAgentEnd && isTerminalAgentEnd(event)) {
      await deps.onAgentSettled?.({ type: "agent_settled" }, ctx);
    }
  });
  pi.on("agent_settled", async (event, ctx) => {
    if (!isActive(ctx)) return;
    await deps.onAgentSettled?.(event, ctx);
  });
}