720 lines · 22.2 KB · TypeScript
RawDownload
import {  type AuthSessionRole,  type EnvironmentId,  type OrchestrationEvent,  type OrchestrationReadModel,  type ServerConfig,  type TerminalEvent,  ThreadId,} from "@t3tools/contracts";import { type QueryClient } from "@tanstack/react-query";import { Throttler } from "@tanstack/react-pacer";import {  createKnownEnvironment,  getKnownEnvironmentWsBaseUrl,  scopedProjectKey,  scopedThreadKey,  scopeProjectRef,  scopeThreadRef,} from "@t3tools/client-runtime"; import {  markPromotedDraftThreadByRef,  markPromotedDraftThreadsByRef,  useComposerDraftStore,} from "~/composerDraftStore";import { ensureLocalApi } from "~/localApi";import { collectActiveTerminalThreadIds } from "~/lib/terminalStateCleanup";import { deriveOrchestrationBatchEffects } from "~/orchestrationEventEffects";import { projectQueryKeys } from "~/lib/projectReactQuery";import { providerQueryKeys } from "~/lib/providerReactQuery";import { getPrimaryKnownEnvironment } from "../primary";import {  bootstrapRemoteBearerSession,  fetchRemoteEnvironmentDescriptor,  fetchRemoteSessionState,  resolveRemoteWebSocketConnectionUrl,} from "../remote/api";import { resolveRemotePairingTarget } from "../remote/target";import {  getSavedEnvironmentRecord,  hasSavedEnvironmentRegistryHydrated,  listSavedEnvironmentRecords,  persistSavedEnvironmentRecord,  readSavedEnvironmentBearerToken,  removeSavedEnvironmentBearerToken,  type SavedEnvironmentRecord,  useSavedEnvironmentRegistryStore,  useSavedEnvironmentRuntimeStore,  waitForSavedEnvironmentRegistryHydration,  writeSavedEnvironmentBearerToken,} from "./catalog";import { createEnvironmentConnection, type EnvironmentConnection } from "./connection";import {  useStore,  selectProjectsAcrossEnvironments,  selectThreadByRef,  selectThreadsAcrossEnvironments,} from "~/store";import { useTerminalStateStore } from "~/terminalStateStore";import { useUiStateStore } from "~/uiStateStore";import { WsTransport } from "../../rpc/wsTransport";import { createWsRpcClient, type WsRpcClient } from "../../rpc/wsRpcClient"; type EnvironmentServiceState = {  readonly queryClient: QueryClient;  readonly queryInvalidationThrottler: Throttler<() => void>;  refCount: number;  stop: () => void;}; const environmentConnections = new Map<EnvironmentId, EnvironmentConnection>();const environmentConnectionListeners = new Set<() => void>(); let activeService: EnvironmentServiceState | null = null;let needsProviderInvalidation = false; function emitEnvironmentConnectionRegistryChange() {  for (const listener of environmentConnectionListeners) {    listener();  }} function getRuntimeErrorFields(error: unknown) {  return {    lastError: error instanceof Error ? error.message : String(error),    lastErrorAt: new Date().toISOString(),  } as const;} function isoNow(): string {  return new Date().toISOString();} function setRuntimeConnecting(environmentId: EnvironmentId) {  useSavedEnvironmentRuntimeStore.getState().patch(environmentId, {    connectionState: "connecting",    lastError: null,    lastErrorAt: null,  });} function setRuntimeConnected(environmentId: EnvironmentId) {  const connectedAt = isoNow();  useSavedEnvironmentRuntimeStore.getState().patch(environmentId, {    connectionState: "connected",    authState: "authenticated",    connectedAt,    disconnectedAt: null,    lastError: null,    lastErrorAt: null,  });  useSavedEnvironmentRegistryStore.getState().markConnected(environmentId, connectedAt);} function setRuntimeDisconnected(environmentId: EnvironmentId, reason?: string | null) {  useSavedEnvironmentRuntimeStore.getState().patch(environmentId, {    connectionState: "disconnected",    disconnectedAt: isoNow(),    ...(reason && reason.trim().length > 0      ? {          lastError: reason,          lastErrorAt: isoNow(),        }      : {}),  });} function setRuntimeError(environmentId: EnvironmentId, error: unknown) {  useSavedEnvironmentRuntimeStore.getState().patch(environmentId, {    connectionState: "error",    ...getRuntimeErrorFields(error),  });} function coalesceOrchestrationUiEvents(  events: ReadonlyArray<OrchestrationEvent>,): OrchestrationEvent[] {  if (events.length < 2) {    return [...events];  }   const coalesced: OrchestrationEvent[] = [];  for (const event of events) {    const previous = coalesced.at(-1);    if (      previous?.type === "thread.message-sent" &&      event.type === "thread.message-sent" &&      previous.payload.threadId === event.payload.threadId &&      previous.payload.messageId === event.payload.messageId    ) {      coalesced[coalesced.length - 1] = {        ...event,        payload: {          ...event.payload,          attachments: event.payload.attachments ?? previous.payload.attachments,          createdAt: previous.payload.createdAt,          text:            !event.payload.streaming && event.payload.text.length > 0              ? event.payload.text              : previous.payload.text + event.payload.text,        },      };      continue;    }     coalesced.push(event);  }   return coalesced;} function reconcileSnapshotDerivedState() {  const storeState = useStore.getState();  const threads = selectThreadsAcrossEnvironments(storeState);  const projects = selectProjectsAcrossEnvironments(storeState);   useUiStateStore.getState().syncProjects(    projects.map((project) => ({      key: scopedProjectKey(scopeProjectRef(project.environmentId, project.id)),      cwd: project.cwd,    })),  );  useUiStateStore.getState().syncThreads(    threads.map((thread) => ({      key: scopedThreadKey(scopeThreadRef(thread.environmentId, thread.id)),      seedVisitedAt: thread.updatedAt ?? thread.createdAt,    })),  );  markPromotedDraftThreadsByRef(    threads.map((thread) => scopeThreadRef(thread.environmentId, thread.id)),  );   const activeThreadKeys = collectActiveTerminalThreadIds({    snapshotThreads: threads.map((thread) => ({      key: scopedThreadKey(scopeThreadRef(thread.environmentId, thread.id)),      deletedAt: null,      archivedAt: thread.archivedAt,    })),    draftThreadKeys: useComposerDraftStore.getState().listDraftThreadKeys(),  });  useTerminalStateStore.getState().removeOrphanedTerminalStates(activeThreadKeys);} export function shouldApplyTerminalEvent(input: {  serverThreadArchivedAt: string | null | undefined;  hasDraftThread: boolean;}): boolean {  if (input.serverThreadArchivedAt !== undefined) {    return input.serverThreadArchivedAt === null;  }   return input.hasDraftThread;} function applyRecoveredEventBatch(  events: ReadonlyArray<OrchestrationEvent>,  environmentId: EnvironmentId,) {  if (events.length === 0) {    return;  }   const batchEffects = deriveOrchestrationBatchEffects(events);  const uiEvents = coalesceOrchestrationUiEvents(events);  const needsProjectUiSync = events.some(    (event) =>      event.type === "project.created" ||      event.type === "project.meta-updated" ||      event.type === "project.deleted",  );   if (batchEffects.needsProviderInvalidation) {    needsProviderInvalidation = true;    void activeService?.queryInvalidationThrottler.maybeExecute();  }   useStore.getState().applyOrchestrationEvents(uiEvents, environmentId);  if (needsProjectUiSync) {    const projects = selectProjectsAcrossEnvironments(useStore.getState());    useUiStateStore.getState().syncProjects(      projects.map((project) => ({        key: scopedProjectKey(scopeProjectRef(project.environmentId, project.id)),        cwd: project.cwd,      })),    );  }   const needsThreadUiSync = events.some(    (event) => event.type === "thread.created" || event.type === "thread.deleted",  );  if (needsThreadUiSync) {    const threads = selectThreadsAcrossEnvironments(useStore.getState());    useUiStateStore.getState().syncThreads(      threads.map((thread) => ({        key: scopedThreadKey(scopeThreadRef(thread.environmentId, thread.id)),        seedVisitedAt: thread.updatedAt ?? thread.createdAt,      })),    );  }   const draftStore = useComposerDraftStore.getState();  for (const threadId of batchEffects.promoteDraftThreadIds) {    markPromotedDraftThreadByRef(scopeThreadRef(environmentId, threadId));  }  for (const threadId of batchEffects.clearDeletedThreadIds) {    draftStore.clearDraftThread(scopeThreadRef(environmentId, threadId));    useUiStateStore      .getState()      .clearThreadUi(scopedThreadKey(scopeThreadRef(environmentId, threadId)));  }  for (const threadId of batchEffects.removeTerminalStateThreadIds) {    useTerminalStateStore.getState().removeTerminalState(scopeThreadRef(environmentId, threadId));  }} function createEnvironmentConnectionHandlers() {  return {    applyEventBatch: applyRecoveredEventBatch,    syncSnapshot: (snapshot: OrchestrationReadModel, environmentId: EnvironmentId) => {      useStore.getState().syncServerReadModel(snapshot, environmentId);      reconcileSnapshotDerivedState();    },    applyTerminalEvent: (event: TerminalEvent, environmentId: EnvironmentId) => {      const threadRef = scopeThreadRef(environmentId, ThreadId.makeUnsafe(event.threadId));      const serverThread = selectThreadByRef(useStore.getState(), threadRef);      const hasDraftThread =        useComposerDraftStore.getState().getDraftThreadByRef(threadRef) !== null;      if (        !shouldApplyTerminalEvent({          serverThreadArchivedAt: serverThread?.archivedAt,          hasDraftThread,        })      ) {        return;      }      useTerminalStateStore.getState().applyTerminalEvent(threadRef, event);    },  };} function createPrimaryEnvironmentClient(  knownEnvironment: ReturnType<typeof getPrimaryKnownEnvironment>,) {  const wsBaseUrl = getKnownEnvironmentWsBaseUrl(knownEnvironment);  if (!wsBaseUrl) {    throw new Error(      `Unable to resolve websocket URL for ${knownEnvironment?.label ?? "primary environment"}.`,    );  }   return createWsRpcClient(new WsTransport(wsBaseUrl));} function createSavedEnvironmentClient(  record: SavedEnvironmentRecord,  bearerToken: string,): WsRpcClient {  useSavedEnvironmentRuntimeStore.getState().ensure(record.environmentId);   return createWsRpcClient(    new WsTransport(      () =>        resolveRemoteWebSocketConnectionUrl({          wsBaseUrl: record.wsBaseUrl,          httpBaseUrl: record.httpBaseUrl,          bearerToken,        }),      {        onAttempt: () => {          setRuntimeConnecting(record.environmentId);        },        onOpen: () => {          setRuntimeConnected(record.environmentId);        },        onError: (message: string) => {          useSavedEnvironmentRuntimeStore.getState().patch(record.environmentId, {            connectionState: "error",            lastError: message,            lastErrorAt: isoNow(),          });        },        onClose: (details: { readonly code: number; readonly reason: string }) => {          setRuntimeDisconnected(record.environmentId, details.reason);        },      },    ),  );} async function refreshSavedEnvironmentMetadata(  record: SavedEnvironmentRecord,  bearerToken: string,  client: WsRpcClient,  roleHint?: AuthSessionRole | null,  configHint?: ServerConfig | null,): Promise<void> {  const [serverConfig, sessionState] = await Promise.all([    configHint ? Promise.resolve(configHint) : client.server.getConfig(),    fetchRemoteSessionState({      httpBaseUrl: record.httpBaseUrl,      bearerToken,    }),  ]);   useSavedEnvironmentRuntimeStore.getState().patch(record.environmentId, {    authState: sessionState.authenticated ? "authenticated" : "requires-auth",    descriptor: serverConfig.environment,    serverConfig,    role: sessionState.authenticated ? (sessionState.role ?? roleHint ?? null) : null,  });} function registerConnection(connection: EnvironmentConnection): EnvironmentConnection {  const existing = environmentConnections.get(connection.environmentId);  if (existing && existing !== connection) {    throw new Error(`Environment ${connection.environmentId} already has an active connection.`);  }  environmentConnections.set(connection.environmentId, connection);  emitEnvironmentConnectionRegistryChange();  return connection;} async function removeConnection(environmentId: EnvironmentId): Promise<boolean> {  const connection = environmentConnections.get(environmentId);  if (!connection) {    return false;  }   environmentConnections.delete(environmentId);  emitEnvironmentConnectionRegistryChange();  await connection.dispose();  return true;} function createPrimaryEnvironmentConnection(): EnvironmentConnection {  const knownEnvironment = getPrimaryKnownEnvironment();  if (!knownEnvironment?.environmentId) {    throw new Error("Unable to resolve the primary environment.");  }   const existing = environmentConnections.get(knownEnvironment.environmentId);  if (existing) {    return existing;  }   return registerConnection(    createEnvironmentConnection({      kind: "primary",      knownEnvironment,      client: createPrimaryEnvironmentClient(knownEnvironment),      ...createEnvironmentConnectionHandlers(),    }),  );} async function ensureSavedEnvironmentConnection(  record: SavedEnvironmentRecord,  options?: {    readonly client?: WsRpcClient;    readonly bearerToken?: string;    readonly role?: AuthSessionRole | null;    readonly serverConfig?: ServerConfig | null;  },): Promise<EnvironmentConnection> {  const existing = environmentConnections.get(record.environmentId);  if (existing) {    return existing;  }   const bearerToken =    options?.bearerToken ?? (await readSavedEnvironmentBearerToken(record.environmentId));  if (!bearerToken) {    useSavedEnvironmentRuntimeStore.getState().patch(record.environmentId, {      authState: "requires-auth",      role: null,      connectionState: "disconnected",      lastError: "Saved environment is missing its saved credential. Pair it again.",      lastErrorAt: isoNow(),    });    throw new Error("Saved environment is missing its saved credential.");  }   const client = options?.client ?? createSavedEnvironmentClient(record, bearerToken);  const knownEnvironment = createKnownEnvironment({    id: record.environmentId,    label: record.label,    source: "manual",    target: {      httpBaseUrl: record.httpBaseUrl,      wsBaseUrl: record.wsBaseUrl,    },  });  const connection = createEnvironmentConnection({    kind: "saved",    knownEnvironment: {      ...knownEnvironment,      environmentId: record.environmentId,    },    client,    refreshMetadata: async () => {      await refreshSavedEnvironmentMetadata(record, bearerToken, client);    },    onConfigSnapshot: (config) => {      useSavedEnvironmentRuntimeStore.getState().patch(record.environmentId, {        descriptor: config.environment,        serverConfig: config,      });    },    onWelcome: (payload) => {      useSavedEnvironmentRuntimeStore.getState().patch(record.environmentId, {        descriptor: payload.environment,      });    },    ...createEnvironmentConnectionHandlers(),  });   registerConnection(connection);   try {    await refreshSavedEnvironmentMetadata(      record,      bearerToken,      client,      options?.role ?? null,      options?.serverConfig ?? null,    );    return connection;  } catch (error) {    setRuntimeError(record.environmentId, error);    await removeConnection(record.environmentId).catch(() => false);    throw error;  }} async function syncSavedEnvironmentConnections(  records: ReadonlyArray<SavedEnvironmentRecord>,): Promise<void> {  const expectedEnvironmentIds = new Set(records.map((record) => record.environmentId));  const staleEnvironmentIds = [...environmentConnections.values()]    .filter((connection) => connection.kind === "saved")

Lines 1–500 of 720.