kandev/apps/web/hooks/domains/session/use-session-resumption.ts

import { useCallback, useEffect, useState, useRef } from "react";
import { getWebSocketClient } from "@/lib/ws/connection";
import { launchSession } from "@/lib/services/session-launch-service";
import { buildResumeRequest, buildRestoreWorkspaceRequest } from "@/lib/services/session-launch-helpers";
import { useAppStore } from "@/components/state-provider";
import type { TaskSessionState } from "@/lib/types/http";

export type SessionStatus = {
  session_id: string;
  task_id: string;
  state: string;
  agent_profile_id?: string;
  is_agent_running: boolean;
  is_resumable: boolean;
  needs_resume: boolean;
  needs_workspace_restore?: boolean;
  resume_reason?: string;
  acp_session_id?: string;
  worktree_path?: string;
  worktree_branch?: string;
  executor_id?: string;
  executor_type?: string;
  executor_name?: string;
  runtime?: string;
  is_remote_executor?: boolean;
  remote_state?: string;
  remote_name?: string;
  remote_created_at?: string;
  remote_checked_at?: string;
  remote_status_error?: string;
  error?: string;
};

export type ResumptionState = "idle" | "checking" | "resuming" | "resumed" | "running" | "error";

type ResumeResponse = {
  success: boolean;
  state?: string;
  worktree_path?: string;
  worktree_branch?: string;
  error?: string;
};

type ResumeStateSetter = {
  setResumptionState: (s: ResumptionState) => void;
  setError: (e: string | null) => void;
  setWorktreePath: (p: string | null) => void;
  setWorktreeBranch: (b: string | null) => void;
  setTaskSession: (s: {
    id: string;
    task_id: string;
    state: TaskSessionState;
    started_at: string;
    updated_at: string;
  }) => void;
};

type SessionLike = { started_at?: string; updated_at?: string } | null;

/** Apply a successful resume response to local state. */
function applyResumeResponse(
  resp: ResumeResponse,
  taskId: string,
  sessionId: string,
  session: SessionLike,
  setters: ResumeStateSetter,
): boolean {
  if (resp.success) {
    setters.setResumptionState("resumed");
    if (resp.state) {
      setters.setTaskSession({
        id: sessionId,
        task_id: taskId,
        state: resp.state as TaskSessionState,
        started_at: session?.started_at ?? "",
        updated_at: session?.updated_at ?? "",
      });
    }
    if (resp.worktree_path) setters.setWorktreePath(resp.worktree_path);
    if (resp.worktree_branch) setters.setWorktreeBranch(resp.worktree_branch);
    return true;
  }
  setters.setResumptionState("error");
  setters.setError(resp.error ?? "Failed to resume session");
  return false;
}

/** Launch a session via a request builder and apply the response. */
async function resumeViaLaunch(
  taskId: string,
  sessionId: string,
  session: SessionLike,
  setters: ResumeStateSetter,
  buildRequest: (taskId: string, sessionId: string) => { request: import("@/lib/services/session-launch-service").LaunchSessionRequest },
): Promise<void> {
  setters.setResumptionState("resuming");
  const { request } = buildRequest(taskId, sessionId);
  const launchResp = await launchSession(request);
  applyResumeResponse(
    {
      success: launchResp.success,
      state: launchResp.state,
      worktree_path: launchResp.worktree_path,
      worktree_branch: launchResp.worktree_branch,
    },
    taskId,
    sessionId,
    session,
    setters,
  );
}

type CheckAndResumeParams = {
  taskId: string;
  sessionId: string;
  session: SessionLike;
  setSessionStatus: (s: SessionStatus) => void;
  setters: ResumeStateSetter;
};

/** Apply session status fields to local state. */
function applyStatusToState(
  status: SessionStatus,
  taskId: string,
  sessionId: string,
  session: SessionLike,
  setters: ResumeStateSetter,
): void {
  setters.setWorktreePath(status.worktree_path ?? null);
  setters.setWorktreeBranch(status.worktree_branch ?? null);
  if (status.state) {
    setters.setTaskSession({
      id: sessionId,
      task_id: taskId,
      state: status.state as TaskSessionState,
      started_at: session?.started_at ?? "",
      updated_at: session?.updated_at ?? "",
    });
  }
}

async function checkAndResume({
  taskId,
  sessionId,
  session,
  setSessionStatus,
  setters,
}: CheckAndResumeParams): Promise<void> {
  const client = getWebSocketClient();
  if (!client) return;
  setters.setResumptionState("checking");
  setters.setError(null);
  try {
    const status = await client.request<SessionStatus>("task.session.status", {
      task_id: taskId,
      session_id: sessionId,
    });
    setSessionStatus(status);
    if (status.error) {
      setters.setResumptionState("error");
      setters.setError(status.error);
      return;
    }
    applyStatusToState(status, taskId, sessionId, session, setters);
    if (status.is_agent_running) {
      setters.setResumptionState("running");
    } else if (status.needs_resume && status.is_resumable) {
      await resumeViaLaunch(taskId, sessionId, session, setters, buildResumeRequest);
    } else if (status.needs_workspace_restore) {
      await resumeViaLaunch(taskId, sessionId, session, setters, buildRestoreWorkspaceRequest);
    } else {
      setters.setResumptionState("idle");
    }
  } catch (err) {
    setters.setResumptionState("error");
    setters.setError(err instanceof Error ? err.message : "Unknown error");
  }
}

interface UseSessionResumptionReturn {
  resumptionState: ResumptionState;
  sessionStatus: SessionStatus | null;
  error: string | null;
  taskSessionState: TaskSessionState | null;
  worktreePath: string | null;
  worktreeBranch: string | null;
  resumeSession: () => Promise<boolean>;
}

/**
 * Hook for handling session resumption on page reload.
 * When a sessionId is provided (from URL), it checks the session status
 * and automatically resumes if needed.
 */
export function useSessionResumption(
  taskId: string | null,
  sessionId: string | null,
): UseSessionResumptionReturn {
  const [resumptionState, setResumptionState] = useState<ResumptionState>("idle");
  const [sessionStatus, setSessionStatus] = useState<SessionStatus | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [worktreePath, setWorktreePath] = useState<string | null>(null);
  const [worktreeBranch, setWorktreeBranch] = useState<string | null>(null);
  const connectionStatus = useAppStore((state) => state.connection.status);
  const session = useAppStore((state) =>
    sessionId ? (state.taskSessions.items[sessionId] ?? null) : null,
  );
  const setTaskSession = useAppStore((state) => state.setTaskSession);
  const hasAttemptedResume = useRef(false);
  const remoteStatusRetryCount = useRef(0);

  const setters: ResumeStateSetter = {
    setResumptionState,
    setError,
    setWorktreePath,
    setWorktreeBranch,
    setTaskSession,
  };

  useEffect(() => {
    hasAttemptedResume.current = false;
    remoteStatusRetryCount.current = 0;
  }, [sessionId, taskId]);

  // Check session status and auto-resume if needed
  useEffect(() => {
    if (!taskId || !sessionId || connectionStatus !== "connected" || hasAttemptedResume.current)
      return;
    hasAttemptedResume.current = true;
    checkAndResume({ taskId, sessionId, session, setSessionStatus, setters });
  }, [taskId, sessionId, connectionStatus, setTaskSession, session]); // eslint-disable-line react-hooks/exhaustive-deps

  // Freshly created remote sessions may return status before runtime metadata is available.
  // Retry a few times so topbar/tooltips can show remote details without manual refresh.
  useEffect(() => {
    if (!taskId || !sessionId || connectionStatus !== "connected") return;
    if (!sessionStatus?.is_remote_executor) return;
    if (sessionStatus.remote_checked_at || sessionStatus.remote_status_error) return;
    if (remoteStatusRetryCount.current >= 3) return;

    const timer = window.setTimeout(async () => {
      const client = getWebSocketClient();
      if (!client) return;
      remoteStatusRetryCount.current += 1;
      try {
        const nextStatus = await client.request<SessionStatus>("task.session.status", {
          task_id: taskId,
          session_id: sessionId,
        });
        setSessionStatus(nextStatus);
      } catch {
        // Best-effort refresh only.
      }
    }, 1500);

    return () => window.clearTimeout(timer);
  }, [taskId, sessionId, connectionStatus, sessionStatus]);

  // Manual resume function
  const resumeSession = useCallback(async (): Promise<boolean> => {
    if (!taskId || !sessionId) return false;
    setResumptionState("resuming");
    setError(null);
    try {
      const { request } = buildResumeRequest(taskId, sessionId);
      const response = await launchSession(request);
      if (response.success) {
        setResumptionState("resumed");
        if (response.state) {
          setTaskSession({
            id: sessionId,
            task_id: taskId,
            state: response.state as TaskSessionState,
            started_at: session?.started_at ?? "",
            updated_at: session?.updated_at ?? "",
          });
        }
        if (response.worktree_path) setWorktreePath(response.worktree_path);
        if (response.worktree_branch) setWorktreeBranch(response.worktree_branch);
        return true;
      }
      setResumptionState("error");
      setError("Failed to resume session");
      return false;
    } catch (err) {
      setResumptionState("error");
      setError(err instanceof Error ? err.message : "Unknown error");
      return false;
    }
  }, [taskId, sessionId, session, setTaskSession, setWorktreePath, setWorktreeBranch]);

  return {
    resumptionState,
    sessionStatus,
    error,
    taskSessionState: session?.state ?? null,
    worktreePath,
    worktreeBranch,
    resumeSession,
  };
}