beaver_project/app-instance/frontend/lib/store.ts

import { create } from 'zustand';

import type {
  AuthUser,
  ChatMessage,
  ProcessArtifact,
  ProcessEvent,
  ProcessRun,
  ProcessWsEvent,
  Session,
  SessionProcessProjection,
  UiAgentDescriptor,
  UiMcpServerDescriptor,
} from '@/types';
import type { WsStatus } from '@/lib/api';

const ACTIVE_SESSION_STORAGE_KEY = 'beaver_active_session_id';

function getInitialSessionId(): string {
  if (typeof window === 'undefined') {
    return 'web:default';
  }
  const saved = window.localStorage.getItem(ACTIVE_SESSION_STORAGE_KEY)?.trim();
  return saved || 'web:default';
}

function persistSessionId(id: string): void {
  if (typeof window === 'undefined') {
    return;
  }
  window.localStorage.setItem(ACTIVE_SESSION_STORAGE_KEY, id);
}

interface ChatStore {
  user: AuthUser | null;
  isAuthLoading: boolean;
  sessionId: string;
  messages: ChatMessage[];
  isLoading: boolean;
  streamingContent: string;
  wsStatus: WsStatus;
  isThinking: boolean;
  beaverReady: boolean | null;
  sessions: Session[];
  processRuns: ProcessRun[];
  processEvents: ProcessEvent[];
  processArtifacts: ProcessArtifact[];
  selectedRunId: string | null;
  selectedArtifactId: string | null;
  agentRegistry: UiAgentDescriptor[];
  mcpRegistry: UiMcpServerDescriptor[];
  mcpToolRegistry: Array<{ server_id: string; tools: Array<Record<string, unknown>> }>;
  lastCancelAck: { runId: string; ok: boolean } | null;
  setUser: (user: AuthUser | null) => void;
  setIsAuthLoading: (loading: boolean) => void;
  setSessionId: (id: string) => void;
  setMessages: (msgs: ChatMessage[]) => void;
  addMessage: (msg: ChatMessage) => void;
  updateMessageFeedback: (
    runId: string,
    feedbackState: ChatMessage['feedback_state'],
    error?: string
  ) => void;
  setIsLoading: (loading: boolean) => void;
  setStreamingContent: (content: string) => void;
  appendStreamingContent: (chunk: string) => void;
  setSessions: (sessions: Session[]) => void;
  clearMessages: () => void;
  setWsStatus: (status: WsStatus) => void;
  setIsThinking: (thinking: boolean) => void;
  setBeaverReady: (ready: boolean | null) => void;
  resetProcessState: () => void;
  ingestProcessEvent: (event: ProcessWsEvent) => void;
  setSessionProcess: (sessionId: string, projection: SessionProcessProjection) => void;
  setSelectedRunId: (runId: string | null) => void;
  setSelectedArtifactId: (artifactId: string | null) => void;
  setAgentRegistry: (agents: UiAgentDescriptor[]) => void;
  setMcpRegistry: (servers: UiMcpServerDescriptor[]) => void;
  setMcpToolRegistry: (tools: Array<{ server_id: string; tools: Array<Record<string, unknown>> }>) => void;
}

function upsertRun(collection: ProcessRun[], run: ProcessRun): ProcessRun[] {
  const index = collection.findIndex((item) => item.run_id === run.run_id);
  if (index === -1) {
    return [...collection, run];
  }
  const next = [...collection];
  next[index] = {
    ...next[index],
    ...run,
    metadata: {
      ...(next[index].metadata ?? {}),
      ...(run.metadata ?? {}),
    },
  };
  return next;
}

function upsertArtifact(collection: ProcessArtifact[], artifact: ProcessArtifact): ProcessArtifact[] {
  const index = collection.findIndex((item) => item.artifact_id === artifact.artifact_id);
  if (index === -1) {
    return [...collection, artifact];
  }
  const next = [...collection];
  next[index] = artifact;
  return next;
}

function appendEvent(collection: ProcessEvent[], event: ProcessEvent): ProcessEvent[] {
  if (collection.some((item) => item.event_id === event.event_id)) {
    return collection;
  }
  return [...collection, event];
}

function createEventId(event: ProcessWsEvent): string {
  if (event.type === 'process_cancel_ack') {
    return `${event.type}:${event.run_id}`;
  }
  const suffix =
    'text' in event && typeof event.text === 'string'
      ? event.text
      : 'title' in event && typeof event.title === 'string'
        ? event.title
        : event.type;
  return `${event.type}:${event.run_id}:${event.created_at}:${suffix}`;
}

export const useChatStore = create<ChatStore>((set) => ({
  user: null,
  isAuthLoading: true,
  sessionId: getInitialSessionId(),
  messages: [],
  isLoading: false,
  streamingContent: '',
  wsStatus: 'disconnected',
  isThinking: false,
  beaverReady: null,
  sessions: [],
  processRuns: [],
  processEvents: [],
  processArtifacts: [],
  selectedRunId: null,
  selectedArtifactId: null,
  agentRegistry: [],
  mcpRegistry: [],
  mcpToolRegistry: [],
  lastCancelAck: null,

  setUser: (user) => set({ user }),
  setIsAuthLoading: (loading) => set({ isAuthLoading: loading }),
  setSessionId: (id) => {
    persistSessionId(id);
    set({ sessionId: id });
  },
  setMessages: (msgs) => set({ messages: msgs }),
  addMessage: (msg) => set((s) => ({ messages: [...s.messages, msg] })),
  updateMessageFeedback: (runId, feedbackState, error) =>
    set((s) => ({
      messages: s.messages.map((message) =>
        message.run_id === runId
          ? {
              ...message,
              feedback_state: feedbackState,
              feedback_error: error,
            }
          : message
      ),
    })),
  setIsLoading: (loading) => set({ isLoading: loading }),
  setStreamingContent: (content) => set({ streamingContent: content }),
  appendStreamingContent: (chunk) =>
    set((s) => ({ streamingContent: s.streamingContent + chunk })),
  setSessions: (sessions) => set({ sessions }),
  clearMessages: () => set({ messages: [], streamingContent: '' }),
  setWsStatus: (status) => set({ wsStatus: status }),
  setIsThinking: (thinking) => set({ isThinking: thinking }),
  setBeaverReady: (ready) => set({ beaverReady: ready }),
  resetProcessState: () =>
    set({
      processRuns: [],
      processEvents: [],
      processArtifacts: [],
      selectedRunId: null,
      selectedArtifactId: null,
      lastCancelAck: null,
    }),
  ingestProcessEvent: (event) =>
    set((state) => {
      if (event.type === 'process_cancel_ack') {
        return {
          lastCancelAck: {
            runId: event.run_id,
            ok: event.ok,
          },
        };
      }

      const eventId = createEventId(event);
      let nextRuns = state.processRuns;
      let nextArtifacts = state.processArtifacts;
      let nextSelectedRunId = state.selectedRunId;
      const nextEvents = appendEvent(state.processEvents, {
        event_id: eventId,
        run_id: event.run_id,
        parent_run_id: 'parent_run_id' in event ? event.parent_run_id ?? null : null,
        kind:
          event.type === 'process_run_started'
            ? 'run_started'
            : event.type === 'process_run_progress'
              ? 'run_progress'
              : event.type === 'process_run_message'
                ? 'run_message'
              : event.type === 'process_run_status'
                ? 'run_status'
                : event.type === 'process_run_artifact'
                  ? 'run_artifact'
                  : event.type === 'process_run_finished'
                    ? 'run_finished'
                    : 'run_cancelled',
        actor_type: event.actor_type,
        actor_id: event.actor_id,
        actor_name: event.actor_name,
        text:
          'text' in event && typeof event.text === 'string'
            ? event.text
            : 'summary' in event && typeof event.summary === 'string'
              ? event.summary
              : undefined,
        status: 'status' in event ? event.status : undefined,
        message_role: 'message_role' in event ? event.message_role : undefined,
        metadata: 'metadata' in event ? event.metadata : undefined,
        created_at: event.created_at,
      });

      if (event.type === 'process_run_started') {
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: event.parent_run_id ?? null,
          session_id: event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: event.title,
          source: event.source ?? null,
          status: event.status,
          started_at: event.created_at,
          metadata: event.metadata,
        });
      }

      if (event.type === 'process_run_status') {
        const current = nextRuns.find((item) => item.run_id === event.run_id);
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: current?.parent_run_id ?? event.parent_run_id ?? null,
          session_id: current?.session_id ?? event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title:
            current?.title || event.actor_name,
          source: current?.source ?? null,
          status: event.status,
          started_at:
            current?.started_at || event.created_at,
          metadata: event.metadata,
        });
      }

      if (event.type === 'process_run_progress') {
        const current = nextRuns.find((item) => item.run_id === event.run_id);
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: current?.parent_run_id ?? event.parent_run_id ?? null,
          session_id: current?.session_id ?? event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: current?.title || event.actor_name,
          source: current?.source ?? null,
          status: current?.status || 'running',
          started_at: current?.started_at || event.created_at,
          metadata: event.metadata,
        });
      }

      if (event.type === 'process_run_message') {
        const current = nextRuns.find((item) => item.run_id === event.run_id);
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: current?.parent_run_id ?? event.parent_run_id ?? null,
          session_id: current?.session_id ?? event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: current?.title || event.actor_name,
          source: current?.source ?? null,
          status: current?.status || 'running',
          started_at: current?.started_at || event.created_at,
          metadata: event.metadata,
        });
      }

      if (event.type === 'process_run_artifact') {
        nextArtifacts = upsertArtifact(nextArtifacts, {
          artifact_id: `${event.run_id}:${event.created_at}:${event.title}`,
          run_id: event.run_id,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: event.title,
          artifact_type: event.artifact_type,
          content: event.content,
          data: event.data,
          file_id: event.file_id,
          url: event.url,
          metadata: event.metadata,
          created_at: event.created_at,
        });
      }

      if (event.type === 'process_run_finished') {
        const current = nextRuns.find((item) => item.run_id === event.run_id);
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: current?.parent_run_id ?? null,
          session_id: current?.session_id ?? event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: current?.title || event.actor_name,
          source: current?.source ?? null,
          status: event.status,
          started_at: current?.started_at || event.created_at,
          finished_at: event.created_at,
          summary: event.summary ?? current?.summary ?? null,
          metadata: event.metadata,
        });
      }

      if (event.type === 'process_run_cancelled') {
        const current = nextRuns.find((item) => item.run_id === event.run_id);
        nextRuns = upsertRun(nextRuns, {
          run_id: event.run_id,
          parent_run_id: current?.parent_run_id ?? null,
          session_id: current?.session_id ?? event.session_id ?? state.sessionId,
          actor_type: event.actor_type,
          actor_id: event.actor_id,
          actor_name: event.actor_name,
          title: current?.title || event.actor_name,
          source: current?.source ?? null,
          status: 'cancelled',
          started_at: current?.started_at || event.created_at,
          finished_at: event.created_at,
          summary: current?.summary ?? null,
        });
      }

      return {
        processRuns: nextRuns,
        processEvents: nextEvents,
        processArtifacts: nextArtifacts,
        selectedRunId: nextSelectedRunId,
      };
    }),
  setSessionProcess: (sessionId, projection) =>
    set((state) => {
      const incomingRuns = projection.runs || [];
      const incomingEvents = projection.events || [];
      const incomingArtifacts = projection.artifacts || [];
      const incomingRunIds = new Set(incomingRuns.map((run) => run.run_id));
      const nextRuns = [
        ...state.processRuns.filter((run) => run.session_id !== sessionId && !incomingRunIds.has(run.run_id)),
        ...incomingRuns,
      ];
      const liveRunIds = new Set(nextRuns.map((run) => run.run_id));
      const incomingEventIds = new Set(incomingEvents.map((event) => event.event_id));
      const nextEvents = [
        ...state.processEvents.filter(
          (event) => liveRunIds.has(event.run_id) && !incomingEventIds.has(event.event_id)
        ),
        ...incomingEvents,
      ];
      const incomingArtifactIds = new Set(incomingArtifacts.map((artifact) => artifact.artifact_id));
      const nextArtifacts = [
        ...state.processArtifacts.filter(
          (artifact) => liveRunIds.has(artifact.run_id) && !incomingArtifactIds.has(artifact.artifact_id)
        ),
        ...incomingArtifacts,
      ];
      return {
        processRuns: nextRuns,
        processEvents: nextEvents,
        processArtifacts: nextArtifacts,
      };
    }),
  setSelectedRunId: (runId) => set({ selectedRunId: runId }),
  setSelectedArtifactId: (artifactId) => set({ selectedArtifactId: artifactId }),
  setAgentRegistry: (agents) => set({ agentRegistry: agents }),
  setMcpRegistry: (servers) => set({ mcpRegistry: servers }),
  setMcpToolRegistry: (tools) => set({ mcpToolRegistry: tools }),
}));