kandev/apps/web/lib/ws/handlers/tasks.ts

import type { StoreApi } from "zustand";
import type { AppState } from "@/lib/state/store";
import type { WsHandlers } from "@/lib/ws/handlers/types";
import type { KanbanState } from "@/lib/state/slices/kanban/types";
import { cleanupTaskStorage } from "@/lib/local-storage";
import { removeRecentTask } from "@/lib/recent-tasks";
import { useContextFilesStore } from "@/lib/state/context-files-store";
import { toKanbanTask, type TaskLike } from "@/lib/kanban/map-task";
import { sessionId as toSessionId } from "@/lib/types/http";
import { mergeTaskRepositoryFields } from "@/lib/ws/handlers/task-repositories";
import { softNavigate } from "@/lib/routing/client-router";
import { isTaskDetailPath } from "@/lib/links";
import {
  clearPinnedSessionIfOverridden,
  shouldPreservePinnedSessionForTask,
} from "@/lib/ws/handlers/agent-session";

type KanbanTask = KanbanState["tasks"][number];

function mergeTaskUpdate(existing: KanbanTask | undefined, nextTask: KanbanTask): KanbanTask {
  if (!existing) return nextTask;
  return {
    ...nextTask,
    ...mergeTaskRepositoryFields(existing, nextTask),
  };
}

function upsertTask(tasks: KanbanTask[], nextTask: KanbanTask): KanbanTask[] {
  const existing = tasks.find((task) => task.id === nextTask.id);
  const merged = mergeTaskUpdate(existing, nextTask);
  return existing
    ? tasks.map((task) => (task.id === nextTask.id ? merged : task))
    : [...tasks, merged];
}

function upsertMultiTask(state: AppState, workflowId: string, task: KanbanTask): AppState {
  const snapshot = state.kanbanMulti.snapshots[workflowId];
  if (!snapshot) return state;
  return {
    ...state,
    kanbanMulti: {
      ...state.kanbanMulti,
      snapshots: {
        ...state.kanbanMulti.snapshots,
        [workflowId]: {
          ...snapshot,
          tasks: upsertTask(snapshot.tasks, task),
        },
      },
    },
  };
}

type TaskEventPayload = TaskLike & {
  workflow_id: string;
  old_workflow_id?: string | null;
  is_ephemeral?: boolean;
  archived_at?: string | null;
};

/** Upsert a task in both single-kanban and multi-kanban snapshots. */
function upsertTaskInBothKanbans(
  state: AppState,
  wfId: string,
  payload: TaskEventPayload,
): AppState {
  // Skip ephemeral tasks - they should never be added to kanban
  if (payload.is_ephemeral) {
    return state;
  }

  const nextTask = toKanbanTask(payload);
  let next = state;

  if (state.kanban.workflowId === wfId) {
    next = { ...next, kanban: { ...next.kanban, tasks: upsertTask(next.kanban.tasks, nextTask) } };
  }

  if (state.kanbanMulti.snapshots[wfId]) {
    next = upsertMultiTask(next, wfId, nextTask);
  }

  return next;
}

/** Look up a task across both single-kanban and multi-kanban snapshots. */
function findTaskInState(state: AppState, taskId: string): KanbanTask | undefined {
  const fromKanban = state.kanban.tasks.find((t) => t.id === taskId);
  if (fromKanban) return fromKanban;
  for (const snapshot of Object.values(state.kanbanMulti.snapshots)) {
    const found = snapshot.tasks.find((t) => t.id === taskId);
    if (found) return found;
  }
  return undefined;
}

/** Remove a task from both single-kanban and multi-kanban snapshots. */
function removeTaskFromBothKanbans(state: AppState, wfId: string, taskId: string): AppState {
  let next = state;
  if (state.kanban.workflowId === wfId) {
    next = {
      ...next,
      kanban: { ...next.kanban, tasks: next.kanban.tasks.filter((t) => t.id !== taskId) },
    };
  }
  const snapshot = state.kanbanMulti.snapshots[wfId];
  if (snapshot) {
    next = {
      ...next,
      kanbanMulti: {
        ...next.kanbanMulti,
        snapshots: {
          ...next.kanbanMulti.snapshots,
          [wfId]: { ...snapshot, tasks: snapshot.tasks.filter((t) => t.id !== taskId) },
        },
      },
    };
  }
  return next;
}

/**
 * Soft-redirect away from a deleted task's page. Only fires when the user is
 * currently parked on that task's route (`/t/<id>` or the compatibility
 * `/tasks/<id>`), so a background deletion of some other task never yanks the
 * user elsewhere.
 */
function redirectAwayFromDeletedTask(deletedId: string): void {
  if (typeof window === "undefined") return;
  if (!isTaskDetailPath(window.location.pathname, deletedId)) return;
  softNavigate("/", "replace");
}

export function registerTasksHandlers(store: StoreApi<AppState>): WsHandlers {
  return {
    "task.created": (message) => {
      // Skip ephemeral tasks (e.g., quick chat) - they shouldn't appear on the Kanban board
      if (message.payload.is_ephemeral) return;
      store.setState((state) =>
        upsertTaskInBothKanbans(state, message.payload.workflow_id, message.payload),
      );
    },
    "task.updated": (message) => {
      // Skip ephemeral tasks (e.g., quick chat) - they shouldn't appear on the Kanban board
      if (message.payload.is_ephemeral) return;

      // Capture the previous primary session id BEFORE the upsert so we can
      // detect a primary-session swap (e.g. workflow profile switch reusing a
      // different session) and follow focus to the new primary.
      const beforeState = store.getState();
      const taskId = message.payload.task_id;
      const previousPrimary = findTaskInState(beforeState, taskId)?.primarySessionId ?? null;

      store.setState((state) => {
        const wfId = message.payload.workflow_id;
        const oldWfId = message.payload.old_workflow_id;
        let next = state;

        if (oldWfId && oldWfId !== wfId) {
          next = removeTaskFromBothKanbans(next, oldWfId, taskId);
        }

        if (message.payload.archived_at) {
          return removeTaskFromBothKanbans(next, wfId, taskId);
        }

        return upsertTaskInBothKanbans(next, wfId, message.payload);
      });

      // Follow focus to the new primary when:
      //  - the user is currently viewing this task,
      //  - the user was sitting on the previous primary,
      //  - they do NOT have a non-terminal pinned session for this task, and
      //  - the primary actually changed.
      // This makes workflow profile switches transparent for unpinned users
      // without yanking users off a live session they deliberately selected.
      const afterState = store.getState();
      const newPrimary = findTaskInState(afterState, taskId)?.primarySessionId ?? null;
      if (
        newPrimary &&
        newPrimary !== previousPrimary &&
        afterState.tasks.activeTaskId === taskId &&
        afterState.tasks.activeSessionId === previousPrimary &&
        !shouldPreservePinnedSessionForTask(afterState, taskId)
      ) {
        clearPinnedSessionIfOverridden(store, newPrimary);
        afterState.setActiveSessionAuto(taskId, newPrimary);
      }
    },
    "task.deleted": (message) => {
      const deletedId = message.payload.task_id;
      const wfId = message.payload.workflow_id;
      removeRecentTask(deletedId);

      const currentState = store.getState();
      const sessionIds = (currentState.taskSessionsByTask.itemsByTaskId[deletedId] ?? []).map(
        (s) => s.id,
      );
      const task = currentState.kanban.tasks.find((t) => t.id === deletedId);
      if (task?.primarySessionId) {
        const primaryId = toSessionId(task.primarySessionId);
        if (!sessionIds.includes(primaryId)) {
          sessionIds.push(primaryId);
        }
      }
      const envIds = Array.from(
        new Set(
          sessionIds
            .map((sid) => currentState.environmentIdBySessionId[sid])
            .filter((eid): eid is string => Boolean(eid)),
        ),
      );
      cleanupTaskStorage(deletedId, sessionIds, envIds);
      // Keep the in-memory sidebar pin/order arrays in sync — without this,
      // a later togglePinnedTask / setSidebarTaskOrder would persist the
      // stale state (still containing the deleted ID) back to localStorage.
      currentState.removeTaskFromSidebarPrefs(deletedId);
      for (const sid of sessionIds) {
        useContextFilesStore.getState().clearSession(sid);
      }

      const wasActive = currentState.tasks.activeTaskId === deletedId;

      store.setState((state) => {
        const isActive = state.tasks.activeTaskId === deletedId;
        let next = removeTaskFromBothKanbans(state, wfId, deletedId);
        if (isActive) {
          next = { ...next, tasks: { ...next.tasks, activeTaskId: null, activeSessionId: null } };
        }
        if (next.tasks.lastSessionByTaskId[deletedId]) {
          const { [deletedId]: _, ...rest } = next.tasks.lastSessionByTaskId;
          next = { ...next, tasks: { ...next.tasks, lastSessionByTaskId: rest } };
        }
        return next;
      });

      // Capture the route match before any redirect mutates the pathname. This
      // covers a fresh load where the browser is parked on the task's route
      // (`/t/<id>` or `/tasks/<id>`) but TaskPageContent hasn't hydrated
      // `activeTaskId` yet, so `wasActive` is still false.
      const onDeletedRoute =
        typeof window !== "undefined" && isTaskDetailPath(window.location.pathname, deletedId);

      // Only react to genuine auto-deletions, which the backend tags with a
      // reason (e.g. a review task whose PR was approved). User-initiated deletes
      // carry no reason: their local delete flow (useTaskRemoval) owns
      // navigation by switching to the next task, so redirecting here would
      // preempt it and strand the user on the home route. For auto-deletions we
      // move off the now-dead route (helper is route-guarded) and explain why.
      if (message.payload.reason && (wasActive || onDeletedRoute)) {
        redirectAwayFromDeletedTask(deletedId);
        store.getState().setTaskDeletedNotification({
          taskId: deletedId,
          title: message.payload.title,
          reason: message.payload.reason,
        });
      }
    },
    "task.state_changed": (message) => {
      // Skip ephemeral tasks (e.g., quick chat) - they shouldn't appear on the Kanban board
      if (message.payload.is_ephemeral) return;

      store.setState((state) =>
        upsertTaskInBothKanbans(state, message.payload.workflow_id, message.payload),
      );
    },
  };
}