kandev/apps/web/lib/ws/handlers/agent-session.ts

import type { StoreApi } from "zustand";
import { createDebugLogger } from "@/lib/debug/log";
import type { AppState } from "@/lib/state/store";
import type { WsHandlers } from "@/lib/ws/handlers/types";
import {
  sessionId as toSessionId,
  taskId as toTaskId,
  type SessionId,
  type TaskId,
  type TaskSessionState,
} from "@/lib/types/http";
import type { QueuedMessage } from "@/lib/state/slices/session/types";
import type { KanbanState, WorkflowSnapshotData } from "@/lib/state/slices/kanban/types";

const debug = createDebugLogger("session:state");

const TERMINAL_SESSION_STATES: ReadonlySet<TaskSessionState> = new Set([
  "COMPLETED",
  "CANCELLED",
  "FAILED",
]);

// States that imply agentctl is up and processing (or has fully processed) input.
// If we observe a session in one of these, agentctl must be ready even if we
// missed the agentctl_ready WS event (e.g. it fired before our subscription).
const AGENT_LIVE_STATES: ReadonlySet<TaskSessionState> = new Set(["RUNNING", "WAITING_FOR_INPUT"]);

export function isTerminalSessionState(state: TaskSessionState | undefined): boolean {
  return !!state && TERMINAL_SESSION_STATES.has(state);
}

function findSessionForTask(state: AppState, taskId: string, sessionId: string) {
  const byTask = state.taskSessionsByTask;
  const sessionsForTask = byTask?.itemsByTaskId?.[taskId];
  if (byTask?.loadedByTaskId?.[taskId]) {
    return sessionsForTask?.find((session) => session.id === sessionId) ?? null;
  }
  // Some task-update call sites use partial stores before taskSessions hydrates.
  const byId = state.taskSessions?.items?.[sessionId];
  if (byId) return byId.task_id === taskId ? byId : null;
  return sessionsForTask?.find((session) => session.id === sessionId) ?? null;
}

function isTaskSessionListHydrating(state: AppState, taskId: string): boolean {
  const byTask = state.taskSessionsByTask;
  if (!byTask) return true;
  if (byTask.loadingByTaskId?.[taskId]) return true;
  return !byTask.loadedByTaskId?.[taskId];
}

/**
 * Manual session selection pins a task-scoped session. Background WS events
 * may only override that pin once the pinned session is known terminal, or
 * when the terminal event is for the pinned session itself.
 */
export function shouldPreservePinnedSessionForTask(
  state: AppState,
  taskId: string,
  incoming?: { sessionId: string; newState: TaskSessionState | undefined },
): boolean {
  const pinnedSessionId = state.tasks.pinnedSessionId;
  if (!pinnedSessionId || state.tasks.activeTaskId !== taskId) return false;
  if (
    incoming?.sessionId === pinnedSessionId &&
    incoming.newState &&
    isTerminalSessionState(incoming.newState)
  ) {
    return false;
  }

  const pinnedSession = findSessionForTask(state, taskId, pinnedSessionId);
  if (!pinnedSession) {
    // Preserve missing rows only while the per-task list is still hydrating.
    // Once loaded, absence means the pinned session was deleted or went stale.
    return isTaskSessionListHydrating(state, taskId);
  }
  return !isTerminalSessionState(pinnedSession.state);
}

export function clearPinnedSessionIfOverridden(store: StoreApi<AppState>, sessionId: string): void {
  const pinnedSessionId = store.getState().tasks.pinnedSessionId;
  if (!pinnedSessionId || pinnedSessionId === sessionId) return;
  store.setState((state) => ({
    ...state,
    tasks: { ...state.tasks, pinnedSessionId: null },
  }));
}

/** Promote agentctl status to "ready" when the session enters a live state.
 *  Acts as a fallback for missed/late agentctl_ready WS events — the backend
 *  cannot reach RUNNING/WAITING_FOR_INPUT without a live agentctl. Never
 *  downgrades an existing "ready" entry. */
function maybePromoteAgentctlReady(
  store: StoreApi<AppState>,
  sessionId: string,
  newState: TaskSessionState | undefined,
  timestamp: string | undefined,
): void {
  if (!newState || !AGENT_LIVE_STATES.has(newState)) return;
  const current = store.getState().sessionAgentctl?.itemsBySessionId?.[sessionId];
  if (current?.status === "ready") return;
  store.getState().setSessionAgentctlStatus(sessionId, {
    status: "ready",
    agentExecutionId: current?.agentExecutionId,
    updatedAt: timestamp,
  });
}

/**
 * When the backend creates a new session for the active task (e.g., due to a
 * workflow step transition with a different agent profile), the chat UI should
 * follow the switch. Returns true when the caller should adopt the new session
 * as the task's active session.
 *
 * Adopts only when the current active session is missing, cross-task, or
 * already terminal — not while a live session for the same task is still
 * running (the backend only creates sessions during workflow step transitions
 * after stopping the previous one, but WS events may arrive out of order).
 */
export function shouldAdoptNewSession(
  state: AppState,
  taskId: string,
  newState: TaskSessionState | undefined,
): boolean {
  if (!newState || isTerminalSessionState(newState)) return false;
  if (state.tasks.activeTaskId !== taskId) return false;
  const activeSessionId = state.tasks.activeSessionId;
  if (activeSessionId) {
    const activeSession = state.taskSessions.items[activeSessionId];
    if (activeSession?.task_id === taskId && !isTerminalSessionState(activeSession.state)) {
      return false;
    }
  }
  return true;
}

/**
 * Pick the newest non-terminal session for a task. Used when the currently
 * active session just reached a terminal state — we want to hand focus to the
 * session that replaced it (typically created by a workflow step transition).
 */
export function pickReplacementSessionId(state: AppState, taskId: string): string | null {
  const sessions = state.taskSessionsByTask.itemsByTaskId[taskId];
  if (!sessions) return null;
  for (let i = sessions.length - 1; i >= 0; i -= 1) {
    const candidate = sessions[i];
    if (!isTerminalSessionState(candidate.state)) return candidate.id;
  }
  return null;
}

/** Ignore subscribe snapshots that were read before a newer state landed. */
export function isStaleSessionStateEvent(
  existing: { updated_at?: string } | null | undefined,
  payloadUpdatedAt: string | undefined,
): boolean {
  if (!payloadUpdatedAt || !existing?.updated_at) return false;
  const payloadTime = Date.parse(payloadUpdatedAt);
  const existingTime = Date.parse(existing.updated_at);
  if (Number.isNaN(payloadTime) || Number.isNaN(existingTime)) return false;
  // Strict less-than: equal timestamps are treated as not-stale so identical
  // events upsert idempotently rather than being silently dropped.
  return payloadTime < existingTime;
}

/** Build a session update object from the state_changed payload. */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function buildSessionUpdate(payload: any): Record<string, unknown> {
  const update: Record<string, unknown> = {};
  if (payload.new_state) update.state = payload.new_state;
  if (payload.agent_profile_id) update.agent_profile_id = payload.agent_profile_id;
  if (payload.review_status !== undefined) update.review_status = payload.review_status;
  if (payload.error_message !== undefined) update.error_message = payload.error_message;
  if (payload.agent_profile_snapshot)
    update.agent_profile_snapshot = payload.agent_profile_snapshot;
  if (payload.is_passthrough !== undefined) update.is_passthrough = payload.is_passthrough;
  if (payload.session_metadata !== undefined) update.metadata = payload.session_metadata;
  if (payload.task_environment_id) update.task_environment_id = payload.task_environment_id;
  if (payload.updated_at) update.updated_at = payload.updated_at;
  return update;
}

/** Upsert the session in the per-task sessions list from a WS event.
 *  Uses `upsertTaskSessionFromEvent` so the per-task list is not marked as
 *  fully loaded — partial event payloads must not gate the API hydration that
 *  fills in fields like agent_profile_id / repository_id / worktree_path. */
function upsertTaskSessionList(
  store: StoreApi<AppState>,
  taskId: TaskId,
  sessionId: SessionId,
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  payload: any,
  sessionUpdate: Record<string, unknown>,
): void {
  const newState = payload.new_state as TaskSessionState | undefined;
  const existing = store.getState().taskSessions.items[sessionId];
  if (!existing && !newState) return;

  store.getState().upsertTaskSessionFromEvent(taskId, {
    id: sessionId,
    task_id: taskId,
    state: (newState ?? existing?.state) as TaskSessionState,
    started_at: existing?.started_at ?? "",
    updated_at: (sessionUpdate.updated_at as string | undefined) ?? existing?.updated_at ?? "",
    ...(payload.agent_profile_id ? { agent_profile_id: payload.agent_profile_id } : {}),
    ...sessionUpdate,
  });
}

// Fan out the office refetch trigger only when the session's state
// actually changed. The WS layer fires `session.state_changed` for
// several adjacent reasons (agentctl status, context window, model
// updates) where `new_state` is undefined or unchanged; without this
// gate every one of those storms the dashboard-card re-render path.
function maybeFanOutOfficeRefetch(
  store: StoreApi<AppState>,
  newState: TaskSessionState | undefined,
  prevState: TaskSessionState | undefined,
): void {
  if (!newState || newState === prevState) return;
  const setOfficeTrigger = store.getState().setOfficeRefetchTrigger;
  if (!setOfficeTrigger) return;
  setOfficeTrigger("dashboard");
  setOfficeTrigger("agents");
}

function patchTaskPrimarySessionState(
  tasks: KanbanState["tasks"],
  taskId: string,
  sessionId: string,
  newState: TaskSessionState,
): KanbanState["tasks"] {
  let changed = false;
  const nextTasks = tasks.map((task) => {
    if (task.id !== taskId || task.primarySessionId !== sessionId) return task;
    if (task.primarySessionState === newState) return task;
    changed = true;
    return { ...task, primarySessionState: newState };
  });
  return changed ? nextTasks : tasks;
}

function patchSnapshotPrimarySessionState(
  snapshot: WorkflowSnapshotData,
  taskId: string,
  sessionId: string,
  newState: TaskSessionState,
): WorkflowSnapshotData {
  const tasks = patchTaskPrimarySessionState(snapshot.tasks, taskId, sessionId, newState);
  return tasks === snapshot.tasks ? snapshot : { ...snapshot, tasks };
}

function syncKanbanPrimarySessionState(
  store: StoreApi<AppState>,
  taskId: TaskId,
  sessionId: SessionId,
  newState: TaskSessionState | undefined,
): void {
  if (!newState) return;

  store.setState((state) => {
    const nextKanbanTasks = patchTaskPrimarySessionState(
      state.kanban.tasks,
      taskId,
      sessionId,
      newState,
    );
    let snapshotsChanged = false;
    const nextSnapshots = Object.fromEntries(
      Object.entries(state.kanbanMulti.snapshots).map(([workflowId, snapshot]) => {
        const nextSnapshot = patchSnapshotPrimarySessionState(
          snapshot,
          taskId,
          sessionId,
          newState,
        );
        if (nextSnapshot !== snapshot) snapshotsChanged = true;
        return [workflowId, nextSnapshot];
      }),
    );

    if (nextKanbanTasks === state.kanban.tasks && !snapshotsChanged) return state;

    return {
      ...state,
      kanban:
        nextKanbanTasks === state.kanban.tasks
          ? state.kanban
          : { ...state.kanban, tasks: nextKanbanTasks },
      kanbanMulti: snapshotsChanged
        ? {
            ...state.kanbanMulti,
            snapshots: nextSnapshots,
          }
        : state.kanbanMulti,
    };
  });
}

/** Extract context window data from payload metadata and store it. */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function extractContextWindow(store: StoreApi<AppState>, sessionId: string, payload: any): void {
  const metadata = payload.metadata;
  if (!metadata || typeof metadata !== "object") return;
  const contextWindow = (metadata as Record<string, unknown>).context_window;
  if (!contextWindow || typeof contextWindow !== "object") return;
  const cw = contextWindow as Record<string, unknown>;
  store.getState().setContextWindow(sessionId, {
    size: (cw.size as number) ?? 0,
    used: (cw.used as number) ?? 0,
    remaining: (cw.remaining as number) ?? 0,
    efficiency: (cw.efficiency as number) ?? 0,
    timestamp: new Date().toISOString(),
  });
}

/** Copy agentctl "ready" status from one session to another (same-task switch). */
function inheritAgentctlStatus(state: AppState, fromSessionId: string, toSessionId: string): void {
  const oldAgentctl = state.sessionAgentctl?.itemsBySessionId?.[fromSessionId];
  if (oldAgentctl?.status === "ready") {
    state.setSessionAgentctlStatus(toSessionId, oldAgentctl);
  }
}

/**
 * After a `session.state_changed` event, decide whether the chat UI should
 * follow a workflow-driven session switch. Covers both event orderings:
 *   1. New non-terminal session appears for the active task before the old
 *      one is torn down — adopt immediately.
 *   2. The current active session transitions to a terminal state — hand off
 *      to the newest non-terminal session for the same task, if any.
 */
function maybeAdoptSessionOnTransition(
  store: StoreApi<AppState>,
  taskId: string,
  sessionId: string,
  newState: TaskSessionState | undefined,
  wasKnownToStore: boolean,
): void {
  const state = store.getState();
  if (
    state.tasks.pinnedSessionId !== sessionId &&
    shouldPreservePinnedSessionForTask(state, taskId, { sessionId, newState })
  ) {
    return;
  }

  if (!wasKnownToStore && shouldAdoptNewSession(state, taskId, newState)) {
    const oldSessionId = state.tasks.activeSessionId;
    if (oldSessionId) inheritAgentctlStatus(state, oldSessionId, sessionId);
    clearPinnedSessionIfOverridden(store, sessionId);
    state.setActiveSessionAuto(taskId, sessionId);
    return;
  }

  const isActive = state.tasks.activeSessionId === sessionId;
  if (isActive && newState && isTerminalSessionState(newState)) {
    const replacement = pickReplacementSessionId(state, taskId);
    if (replacement && replacement !== sessionId) {
      inheritAgentctlStatus(state, sessionId, replacement);
      clearPinnedSessionIfOverridden(store, replacement);
      state.setActiveSessionAuto(taskId, replacement);
    }
  }
}

/** Seed the session→environment mapping from an agentctl event payload.
 *  Brand-new CREATED sessions never receive a `state_changed` carrying
 *  `task_environment_id` (no transition fires), so the env mapping would
 *  otherwise stay empty and env-routed shell terminals stall on
 *  "Connecting terminal...". Routes through `upsertTaskSessionFromEvent`
 *  which calls `syncEnvironmentMapping` and merges with any existing row. */
function syncEnvFromAgentctlPayload(
  store: StoreApi<AppState>,
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  payload: any,
): void {
  const rawTaskId = payload?.task_id;
  const rawSessionId = payload?.session_id;
  const envId = payload?.task_environment_id;
  if (!rawTaskId || !rawSessionId || !envId) return;
  const taskId = toTaskId(rawTaskId);
  const sessionId = toSessionId(rawSessionId);
  const existing = store.getState().taskSessions.items[sessionId];
  store.getState().upsertTaskSessionFromEvent(taskId, {
    id: sessionId,
    task_id: taskId,
    state: existing?.state ?? "CREATED",
    started_at: existing?.started_at ?? "",
    updated_at: existing?.updated_at ?? "",
    task_environment_id: envId,
    worktree_id: payload.worktree_id,
    worktree_path: payload.worktree_path,
    worktree_branch: payload.worktree_branch,
  });
}

/** Builds the partial-session patch applied for an agentctl_ready event.
 *  On sibling materialize we repoint worktree_path to the task root and keep
 *  the primary's id/branch; the initial ready event sets id/path/branch
 *  straight from the payload. */
function buildAgentctlReadySessionUpdate(
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  payload: any,
  isSibling: boolean,
): Record<string, unknown> {
  const update: Record<string, unknown> = {};
  if (isSibling) {
    if (payload.task_workspace_path) update.worktree_path = payload.task_workspace_path;
    return update;
  }
  if (payload.worktree_id) update.worktree_id = payload.worktree_id;
  if (payload.worktree_path) update.worktree_path = payload.worktree_path;
  if (payload.worktree_branch) update.worktree_branch = payload.worktree_branch;
  return update;
}

/** Adds the materialized worktree to the worktrees map + the per-session list. */
function recordAgentctlReadyWorktree(
  store: StoreApi<AppState>,
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  payload: any,
  existingSession: { repository_id?: string; worktree_path?: string; worktree_branch?: string },
): void {
  if (!payload.worktree_id) return;
  store.getState().setWorktree({
    id: payload.worktree_id,
    sessionId: payload.session_id,
    repositoryId: existingSession.repository_id ?? undefined,
    path: payload.worktree_path ?? existingSession.worktree_path ?? undefined,
    branch: payload.worktree_branch ?? existingSession.worktree_branch ?? undefined,
  });
  const existing =
    store.getState().sessionWorktreesBySessionId.itemsBySessionId[payload.session_id] ?? [];
  if (!existing.includes(payload.worktree_id)) {
    store.getState().setSessionWorktrees(payload.session_id, [...existing, payload.worktree_id]);
  }
}

/** Handle the agentctl_ready event: update session worktree info.
 *
 *  Two shapes share this event:
 *    1. Initial session ready — payload describes the session's primary
 *       worktree; we set worktree_id/path/branch on the session row.
 *    2. Sibling materialized (multi-branch add_branch flow) — payload
 *       describes a NEW worktree being added alongside the primary. The
 *       primary's worktree_id/branch must NOT be clobbered (they still own
 *       the chat/agent process); only worktree_path moves to the task root
 *       so the file browser repoints from "primary worktree" to "task root
 *       containing both worktree siblings". A commits refetch is bumped so
 *       the Commits panel re-queries with the new multi-repo subpaths
 *       (each commit then carries its repo/branch slug for grouping).
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function handleAgentctlReady(store: StoreApi<AppState>, payload: any): void {
  const existingSession = store.getState().taskSessions.items[payload.session_id];
  if (!existingSession) return;

  const isSibling =
    !!payload.worktree_id &&
    !!existingSession.worktree_id &&
    payload.worktree_id !== existingSession.worktree_id;

  const sessionUpdate = buildAgentctlReadySessionUpdate(payload, isSibling);
  if (Object.keys(sessionUpdate).length > 0) {
    store.getState().setTaskSession({ ...existingSession, ...sessionUpdate });
  }

  recordAgentctlReadyWorktree(store, payload, existingSession);

  if (isSibling) {
    // Drop the pre-multi-repo git-status snapshot — the backend just
    // transitioned this session from single-repo to multi-repo and the legacy
    // (empty-repo-name) tracker is gone. Without this the Changes panel keeps
    // surfacing the frozen snapshot until the user reloads the tab, masking
    // the per-repo updates streaming in for both the primary and the sibling.
    store.getState().clearLegacyGitStatusEntry(payload.session_id);
    store.getState().bumpSessionCommitsRefetch(payload.session_id);
  }
}

interface SessionFailureContext {
  taskId: TaskId;
  sessionId: SessionId;
  newState: TaskSessionState | undefined;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  payload: any;
  previousState: TaskSessionState | undefined;
}

/** Emits a session-failure notification for FAILED transitions, honoring suppress_toast and replay guards. */
function maybeNotifySessionFailure(store: StoreApi<AppState>, ctx: SessionFailureContext): void {
  const { taskId, sessionId, newState, payload, previousState } = ctx;
  if (newState !== "FAILED") return;

  // Only toast on observed transitions (previousState present and not FAILED).
  // If previousState is undefined we're learning about this session for the
  // first time — that's a snapshot of an already-failed session being replayed
  // by the backend (e.g. on page load / WS reconnect), not a fresh failure.
  if (
    payload.suppress_toast === true ||
    previousState === undefined ||
    previousState === "FAILED"
  ) {
    return;
  }

  store.getState().setSessionFailureNotification({
    sessionId,
    taskId,
    message: payload.error_message ? String(payload.error_message) : "Session failed unexpectedly",
  });
}

export function registerTaskSessionHandlers(store: StoreApi<AppState>): WsHandlers {
  return {
    "message.queue.status_changed": (message) => {
      const payload = message.payload;
      if (!payload?.session_id) {
        console.warn("[Queue] Missing session_id in queue status change event");
        return;
      }
      const sessionId = payload.session_id;
      const entries = (payload.entries as QueuedMessage[] | null | undefined) ?? [];
      const count = typeof payload.count === "number" ? payload.count : entries.length;
      const max = typeof payload.max === "number" ? payload.max : 0;
      store.getState().setQueueEntries(sessionId, entries, { count, max });
    },
    "session.state_changed": (message) => {
      const payload = message.payload;
      if (!payload?.task_id) return;
      const { task_id: rawTaskId, session_id: rawSessionId } = payload;
      const newState = payload.new_state as TaskSessionState | undefined;

      if (!rawSessionId) return;
      const taskId = toTaskId(rawTaskId);
      const sessionId = toSessionId(rawSessionId);

      const sessionUpdate = buildSessionUpdate(payload);
      const existingSession = store.getState().taskSessions.items[sessionId];

      if (isStaleSessionStateEvent(existingSession, payload.updated_at)) {
        debug("state_changed ignored stale snapshot", {
          sessionId,
          task_id: taskId,
          existingUpdatedAt: existingSession?.updated_at,
          payloadUpdatedAt: payload.updated_at,
          newState: newState ?? "-",
        });
        return;
      }

      debug("state_changed", {
        sessionId,
        // Logged before upsertTaskSessionList below, so on the first event for a
        // session the store has no row yet and the auto-resolver can't map it —
        // exactly the oldState="-" anchor line. taskId is already in scope, so
        // pass it directly (rendered as task_id=, matching the auto-annotation).
        task_id: taskId,
        oldState: existingSession?.state ?? "-",
        newState: newState ?? "-",
      });

      upsertTaskSessionList(store, taskId, sessionId, payload, sessionUpdate);
      syncKanbanPrimarySessionState(store, taskId, sessionId, newState);
      extractContextWindow(store, sessionId, payload);
      maybePromoteAgentctlReady(store, sessionId, newState, message.timestamp);

      maybeAdoptSessionOnTransition(store, taskId, sessionId, newState, !!existingSession);

      maybeNotifySessionFailure(store, {
        taskId,
        sessionId,
        newState,
        payload,
        previousState: existingSession?.state,
      });

      maybeFanOutOfficeRefetch(store, newState, existingSession?.state);
    },
    "session.agentctl_starting": (message) => {
      const payload = message.payload;
      if (!payload?.session_id) return;
      store.getState().setSessionAgentctlStatus(payload.session_id, {
        status: "starting",
        agentExecutionId: payload.agent_execution_id,
        updatedAt: message.timestamp,
      });
      syncEnvFromAgentctlPayload(store, payload);
    },
    "session.agentctl_ready": (message) => {
      const payload = message.payload;
      if (!payload?.session_id) return;
      store.getState().setSessionAgentctlStatus(payload.session_id, {
        status: "ready",
        agentExecutionId: payload.agent_execution_id,
        updatedAt: message.timestamp,
      });
      syncEnvFromAgentctlPayload(store, payload);
      handleAgentctlReady(store, payload);
    },
    "session.agentctl_error": (message) => {
      const payload = message.payload;
      if (!payload?.session_id) return;
      store.getState().setSessionAgentctlStatus(payload.session_id, {
        status: "error",
        agentExecutionId: payload.agent_execution_id,
        errorMessage: payload.error_message,
        updatedAt: message.timestamp,
      });
    },
  };
}