kandev/apps/web/e2e/helpers/session.ts

import type { Page } from "@playwright/test";
import type { SeedData } from "../fixtures/test-base";
import type { ApiClient } from "./api-client";
import type { SessionPage } from "../pages/session-page";
import { pollUntil } from "./poll-until";

const DONE_STATES = ["COMPLETED", "WAITING_FOR_INPUT"];
const FAILED_STATES = ["FAILED", "CANCELLED"];

function sessionDoneOrThrow(session: { state: string; error_message?: string } | undefined) {
  if (!session) return false;
  if (FAILED_STATES.includes(session.state)) {
    throw new Error(
      `Session reached ${session.state}${session.error_message ? `: ${session.error_message}` : ""}`,
    );
  }
  return DONE_STATES.includes(session.state);
}

export async function waitForLatestSessionDone(
  apiClient: ApiClient,
  taskId: string,
  expectedCount: number,
  message: string,
  timeout = 120_000,
): Promise<void> {
  await pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(taskId);
      if (sessions.length < expectedCount) return false;
      // API returns sessions newest-first.
      const latest = sessions[0];
      return sessionDoneOrThrow(latest);
    },
    (done) => done,
    timeout,
    message,
  );
}

export async function waitForSessionDone(
  apiClient: ApiClient,
  taskId: string,
  sessionId: string,
  message: string,
  timeout = 120_000,
): Promise<void> {
  await pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(taskId);
      const session = sessions.find((s) => s.id === sessionId);
      return sessionDoneOrThrow(session);
    },
    (done) => done,
    timeout,
    message,
  );
}

export async function waitForWorkspacePath(
  apiClient: ApiClient,
  taskId: string,
  sessionId: string,
  timeout = 30_000,
): Promise<string> {
  return pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(taskId);
      const session = sessions.find((candidate) => candidate.id === sessionId);
      return session?.workspace_path ?? session?.worktree_path ?? "";
    },
    (workspacePath) => workspacePath.length > 0,
    timeout,
    "Waiting for task workspace path",
  );
}

export async function waitForAgentMessage(
  apiClient: ApiClient,
  sessionId: string,
  content: string,
  timeout = 60_000,
): Promise<void> {
  await pollUntil(
    async () => {
      const { messages } = await apiClient.listSessionMessages(sessionId);
      return messages.some(
        (message) => message.author_type === "agent" && message.content.includes(content),
      );
    },
    (found) => found,
    timeout,
    `Waiting for agent message containing ${JSON.stringify(content)}`,
  );
}

export async function waitForSessionState(
  apiClient: ApiClient,
  options: {
    taskId: string;
    sessionId: string;
    expectedState: string;
    message: string;
    timeout?: number;
  },
): Promise<void> {
  await pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(options.taskId);
      return sessions.find((session) => session.id === options.sessionId)?.state ?? null;
    },
    (state) => state === options.expectedState,
    options.timeout ?? 120_000,
    options.message,
  );
}

export async function waitForArchiveCancelledSession(
  apiClient: ApiClient,
  taskId: string,
  sessionId: string,
  message: string,
  timeout = 30_000,
): Promise<void> {
  await pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(taskId);
      const session = sessions.find((candidate) => candidate.id === sessionId);
      return {
        state: session?.state ?? null,
        errorMessage: session?.error_message ?? null,
      };
    },
    (result) => result.state === "CANCELLED" && result.errorMessage === "task tree archived",
    timeout,
    message,
  );
}

export async function waitForSessionEnvironment(
  apiClient: ApiClient,
  options: {
    taskId: string;
    sessionId: string;
    expectedEnvironmentId: string;
    message: string;
    timeout?: number;
  },
): Promise<void> {
  await pollUntil(
    async () => {
      const { sessions } = await apiClient.listTaskSessions(options.taskId);
      const session = sessions.find((s) => s.id === options.sessionId);
      return session?.task_environment_id ?? "";
    },
    (environmentId) => environmentId === options.expectedEnvironmentId,
    options.timeout ?? 60_000,
    options.message,
  );
}

export async function openTaskSession(page: Page, taskId: string): Promise<SessionPage> {
  const { SessionPage: SessionPageClass } = await import("../pages/session-page");
  await page.goto(`/t/${taskId}`);
  const session = new SessionPageClass(page);
  await session.waitForLoad();
  return session;
}

/**
 * Seed a task + session and navigate to it, waiting for the first (normal)
 * turn to complete. Follow-up prompts can then exercise retry flows from a
 * clean idle state.
 */
export async function seedIdleSession(
  testPage: Page,
  apiClient: ApiClient,
  seedData: SeedData,
  title: string,
): Promise<SessionPage> {
  const task = await apiClient.createTaskWithAgent(
    seedData.workspaceId,
    title,
    seedData.agentProfileId,
    {
      description: "/e2e:simple-message",
      workflow_id: seedData.workflowId,
      workflow_step_id: seedData.startStepId,
      repository_ids: [seedData.repositoryId],
    },
  );
  if (!task.session_id) throw new Error("createTaskWithAgent did not return a session_id");
  const session = await openTaskSession(testPage, task.id);
  await session.waitForChatIdle({ timeout: 30_000 });
  await session.composerReady();
  return session;
}