1815 lines · 61.9 KB · TypeScript
RawDownload
import * as Http from "node:http";import fs from "node:fs";import os from "node:os";import path from "node:path"; import * as NodeServices from "@effect/platform-node/NodeServices";import { Effect, Exit, Layer, PlatformError, PubSub, Scope, Stream } from "effect";import { describe, expect, it, afterEach, vi } from "vitest";import { createServer } from "./wsServer";import WebSocket from "ws";import { ServerConfig, type ServerConfigShape } from "./config";import { makeServerProviderLayer, makeServerRuntimeServicesLayer } from "./serverLayers"; import {  DEFAULT_TERMINAL_ID,  EDITORS,  EventId,  ORCHESTRATION_WS_CHANNELS,  ORCHESTRATION_WS_METHODS,  ProviderItemId,  ThreadId,  TurnId,  WS_CHANNELS,  WS_METHODS,  type WebSocketResponse,  type ProviderRuntimeEvent,  type ServerProviderStatus,  type KeybindingsConfig,  type ResolvedKeybindingsConfig,  type WsPushChannel,  type WsPushMessage,  type WsPush,} from "@t3tools/contracts";import { compileResolvedKeybindingRule, DEFAULT_KEYBINDINGS } from "./keybindings";import type {  TerminalClearInput,  TerminalCloseInput,  TerminalEvent,  TerminalOpenInput,  TerminalResizeInput,  TerminalSessionSnapshot,  TerminalWriteInput,} from "@t3tools/contracts";import { TerminalManager, type TerminalManagerShape } from "./terminal/Services/Manager";import { makeSqlitePersistenceLive, SqlitePersistenceMemory } from "./persistence/Layers/Sqlite";import { SqlClient, SqlError } from "effect/unstable/sql";import { ProviderService, type ProviderServiceShape } from "./provider/Services/ProviderService";import { ProviderHealth, type ProviderHealthShape } from "./provider/Services/ProviderHealth";import { Open, type OpenShape } from "./open";import { GitManager, type GitManagerShape } from "./git/Services/GitManager.ts";import type { GitCoreShape } from "./git/Services/GitCore.ts";import { GitCore } from "./git/Services/GitCore.ts";import { GitCommandError, GitManagerError } from "./git/Errors.ts";import { MigrationError } from "@effect/sql-sqlite-bun/SqliteMigrator";import { AnalyticsService } from "./telemetry/Services/AnalyticsService.ts"; const asEventId = (value: string): EventId => EventId.makeUnsafe(value);const asProviderItemId = (value: string): ProviderItemId => ProviderItemId.makeUnsafe(value);const asThreadId = (value: string): ThreadId => ThreadId.makeUnsafe(value);const asTurnId = (value: string): TurnId => TurnId.makeUnsafe(value); const defaultOpenService: OpenShape = {  openBrowser: () => Effect.void,  openInEditor: () => Effect.void,}; const defaultProviderStatuses: ReadonlyArray<ServerProviderStatus> = [  {    provider: "codex",    status: "ready",    available: true,    authStatus: "authenticated",    checkedAt: "2026-01-01T00:00:00.000Z",  },]; const defaultProviderHealthService: ProviderHealthShape = {  getStatuses: Effect.succeed(defaultProviderStatuses),}; class MockTerminalManager implements TerminalManagerShape {  private readonly sessions = new Map<string, TerminalSessionSnapshot>();  private readonly listeners = new Set<(event: TerminalEvent) => void>();   private key(threadId: string, terminalId: string): string {    return `${threadId}\u0000${terminalId}`;  }   emitEvent(event: TerminalEvent): void {    for (const listener of this.listeners) {      listener(event);    }  }   subscriptionCount(): number {    return this.listeners.size;  }   readonly open: TerminalManagerShape["open"] = (input: TerminalOpenInput) =>    Effect.sync(() => {      const now = new Date().toISOString();      const terminalId = input.terminalId ?? DEFAULT_TERMINAL_ID;      const snapshot: TerminalSessionSnapshot = {        threadId: input.threadId,        terminalId,        cwd: input.cwd,        status: "running",        pid: 4242,        history: "",        exitCode: null,        exitSignal: null,        updatedAt: now,      };      this.sessions.set(this.key(input.threadId, terminalId), snapshot);      queueMicrotask(() => {        this.emitEvent({          type: "started",          threadId: input.threadId,          terminalId,          createdAt: now,          snapshot,        });      });      return snapshot;    });   readonly write: TerminalManagerShape["write"] = (input: TerminalWriteInput) =>    Effect.sync(() => {      const terminalId = input.terminalId ?? DEFAULT_TERMINAL_ID;      const existing = this.sessions.get(this.key(input.threadId, terminalId));      if (!existing) {        throw new Error(`Unknown terminal thread: ${input.threadId}`);      }      queueMicrotask(() => {        this.emitEvent({          type: "output",          threadId: input.threadId,          terminalId,          createdAt: new Date().toISOString(),          data: input.data,        });      });    });   readonly resize: TerminalManagerShape["resize"] = (_input: TerminalResizeInput) => Effect.void;   readonly clear: TerminalManagerShape["clear"] = (input: TerminalClearInput) =>    Effect.sync(() => {      const terminalId = input.terminalId ?? DEFAULT_TERMINAL_ID;      queueMicrotask(() => {        this.emitEvent({          type: "cleared",          threadId: input.threadId,          terminalId,          createdAt: new Date().toISOString(),        });      });    });   readonly restart: TerminalManagerShape["restart"] = (input: TerminalOpenInput) =>    Effect.sync(() => {      const now = new Date().toISOString();      const terminalId = input.terminalId ?? DEFAULT_TERMINAL_ID;      const snapshot: TerminalSessionSnapshot = {        threadId: input.threadId,        terminalId,        cwd: input.cwd,        status: "running",        pid: 5252,        history: "",        exitCode: null,        exitSignal: null,        updatedAt: now,      };      this.sessions.set(this.key(input.threadId, terminalId), snapshot);      queueMicrotask(() => {        this.emitEvent({          type: "restarted",          threadId: input.threadId,          terminalId,          createdAt: now,          snapshot,        });      });      return snapshot;    });   readonly close: TerminalManagerShape["close"] = (input: TerminalCloseInput) =>    Effect.sync(() => {      if (input.terminalId) {        this.sessions.delete(this.key(input.threadId, input.terminalId));        return;      }      for (const key of this.sessions.keys()) {        if (key.startsWith(`${input.threadId}\u0000`)) {          this.sessions.delete(key);        }      }    });   readonly subscribe: TerminalManagerShape["subscribe"] = (listener) =>    Effect.sync(() => {      this.listeners.add(listener);      return () => {        this.listeners.delete(listener);      };    });   readonly dispose: TerminalManagerShape["dispose"] = Effect.void;} // ---------------------------------------------------------------------------// WebSocket test harness//// Incoming messages are split into two channels://   - pushChannel: server push envelopes (type === "push")//   - responseChannel: request/response envelopes (have an "id" field)//// This means sendRequest never has to skip push messages and waitForPush// never has to skip response messages, eliminating a class of ordering bugs.// --------------------------------------------------------------------------- interface MessageChannel<T> {  queue: T[];  waiters: Array<{    resolve: (value: T) => void;    reject: (error: Error) => void;    timeoutId: ReturnType<typeof setTimeout> | null;  }>;} interface SocketChannels {  push: MessageChannel<WsPush>;  response: MessageChannel<WebSocketResponse>;} const channelsBySocket = new WeakMap<WebSocket, SocketChannels>(); function enqueue<T>(channel: MessageChannel<T>, item: T) {  const waiter = channel.waiters.shift();  if (waiter) {    if (waiter.timeoutId !== null) clearTimeout(waiter.timeoutId);    waiter.resolve(item);    return;  }  channel.queue.push(item);} function dequeue<T>(channel: MessageChannel<T>, timeoutMs: number): Promise<T> {  const queued = channel.queue.shift();  if (queued !== undefined) {    return Promise.resolve(queued);  }   return new Promise((resolve, reject) => {    const waiter = {      resolve,      reject,      timeoutId: setTimeout(() => {        const index = channel.waiters.indexOf(waiter);        if (index >= 0) channel.waiters.splice(index, 1);        reject(new Error(`Timed out waiting for WebSocket message after ${timeoutMs}ms`));      }, timeoutMs) as ReturnType<typeof setTimeout>,    };    channel.waiters.push(waiter);  });} function isWsPushEnvelope(message: unknown): message is WsPush {  if (typeof message !== "object" || message === null) return false;  if (!("type" in message) || !("channel" in message)) return false;  return (message as { type?: unknown }).type === "push";} function asWebSocketResponse(message: unknown): WebSocketResponse | null {  if (typeof message !== "object" || message === null) return null;  if (!("id" in message)) return null;  const id = (message as { id?: unknown }).id;  if (typeof id !== "string") return null;  return message as WebSocketResponse;} function connectWsOnce(port: number, token?: string): Promise<WebSocket> {  return new Promise((resolve, reject) => {    const query = token ? `?token=${encodeURIComponent(token)}` : "";    const ws = new WebSocket(`ws://127.0.0.1:${port}/${query}`);    const channels: SocketChannels = {      push: { queue: [], waiters: [] },      response: { queue: [], waiters: [] },    };    channelsBySocket.set(ws, channels);     ws.on("message", (raw) => {      const parsed = JSON.parse(String(raw));      if (isWsPushEnvelope(parsed)) {        enqueue(channels.push, parsed);      } else {        const response = asWebSocketResponse(parsed);        if (response) {          enqueue(channels.response, response);        }      }    });     ws.once("open", () => resolve(ws));    ws.once("error", () => reject(new Error("WebSocket connection failed")));  });} async function connectWs(port: number, token?: string, attempts = 5): Promise<WebSocket> {  let lastError: unknown = new Error("WebSocket connection failed");   for (let attempt = 0; attempt < attempts; attempt += 1) {    try {      return await connectWsOnce(port, token);    } catch (error) {      lastError = error;      if (attempt < attempts - 1) {        await new Promise((resolve) => setTimeout(resolve, 25));      }    }  }   throw lastError;} /** Connect and wait for the server.welcome push. Returns [ws, welcomeData]. */async function connectAndAwaitWelcome(  port: number,  token?: string,): Promise<[WebSocket, WsPushMessage<typeof WS_CHANNELS.serverWelcome>]> {  const ws = await connectWs(port, token);  const welcome = await waitForPush(ws, WS_CHANNELS.serverWelcome);  return [ws, welcome];} async function sendRequest(  ws: WebSocket,  method: string,  params?: unknown,): Promise<WebSocketResponse> {  const channels = channelsBySocket.get(ws);  if (!channels) throw new Error("WebSocket not initialized");   const id = crypto.randomUUID();  const body =    method === ORCHESTRATION_WS_METHODS.dispatchCommand      ? { _tag: method, command: params }      : params && typeof params === "object" && !Array.isArray(params)        ? { _tag: method, ...(params as Record<string, unknown>) }        : { _tag: method };  ws.send(JSON.stringify({ id, body }));   // Response channel only contains responses — no push filtering needed  while (true) {    const response = await dequeue(channels.response, 60_000);    if (response.id === id || response.id === "unknown") {      return response;    }  }} async function waitForPush<C extends WsPushChannel>(  ws: WebSocket,  channel: C,  predicate?: (push: WsPushMessage<C>) => boolean,  maxMessages = 120,  idleTimeoutMs = 5_000,): Promise<WsPushMessage<C>> {  const channels = channelsBySocket.get(ws);  if (!channels) throw new Error("WebSocket not initialized");   for (let remaining = maxMessages; remaining > 0; remaining--) {    const push = await dequeue(channels.push, idleTimeoutMs);    if (push.channel !== channel) continue;    const typed = push as WsPushMessage<C>;    if (!predicate || predicate(typed)) return typed;  }  throw new Error(`Timed out waiting for push on ${channel}`);} async function rewriteKeybindingsAndWaitForPush(  ws: WebSocket,  keybindingsPath: string,  contents: string,  predicate: (push: WsPushMessage<typeof WS_CHANNELS.serverConfigUpdated>) => boolean,  attempts = 3,): Promise<WsPushMessage<typeof WS_CHANNELS.serverConfigUpdated>> {  let lastError: unknown;  for (let attempt = 0; attempt < attempts; attempt++) {    fs.writeFileSync(keybindingsPath, contents, "utf8");    try {      return await waitForPush(ws, WS_CHANNELS.serverConfigUpdated, predicate, 20, 3_000);    } catch (error) {      lastError = error;    }  }  throw lastError;} async function requestPath(  port: number,  requestPath: string,): Promise<{ statusCode: number; body: string }> {  return new Promise((resolve, reject) => {    const req = Http.request(      {        hostname: "127.0.0.1",        port,        path: requestPath,        method: "GET",      },      (res) => {        const chunks: Buffer[] = [];        res.on("data", (chunk) => {          chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));        });        res.on("end", () => {          resolve({            statusCode: res.statusCode ?? 0,            body: Buffer.concat(chunks).toString("utf8"),          });        });      },    );    req.once("error", reject);    req.end();  });} function compileKeybindings(bindings: KeybindingsConfig): ResolvedKeybindingsConfig {  const resolved: Array<ResolvedKeybindingsConfig[number]> = [];  for (const binding of bindings) {    const compiled = compileResolvedKeybindingRule(binding);    if (!compiled) {      throw new Error(`Unexpected invalid keybinding in test setup: ${binding.command}`);    }    resolved.push(compiled);  }  return resolved;} const DEFAULT_RESOLVED_KEYBINDINGS = compileKeybindings([...DEFAULT_KEYBINDINGS]);const VALID_EDITOR_IDS = new Set(EDITORS.map((editor) => editor.id)); function expectAvailableEditors(value: unknown): void {  expect(Array.isArray(value)).toBe(true);  for (const editorId of value as unknown[]) {    expect(typeof editorId).toBe("string");    expect(VALID_EDITOR_IDS.has(editorId as (typeof EDITORS)[number]["id"])).toBe(true);  }} describe("WebSocket Server", () => {  let server: Http.Server | null = null;  let serverScope: Scope.Closeable | null = null;  const connections: WebSocket[] = [];  const tempDirs: string[] = [];   function makeTempDir(prefix: string): string {    const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix));    tempDirs.push(dir);    return dir;  }   async function createTestServer(    options: {      persistenceLayer?: Layer.Layer<        SqlClient.SqlClient,        SqlError.SqlError | MigrationError | PlatformError.PlatformError      >;      cwd?: string;      autoBootstrapProjectFromCwd?: boolean;      logWebSocketEvents?: boolean;      devUrl?: string;      authToken?: string;      stateDir?: string;      staticDir?: string;      providerLayer?: Layer.Layer<ProviderService, never>;      providerHealth?: ProviderHealthShape;      open?: OpenShape;      gitManager?: GitManagerShape;      gitCore?: Pick<GitCoreShape, "listBranches" | "initRepo" | "pullCurrentBranch">;      terminalManager?: TerminalManagerShape;    } = {},  ): Promise<Http.Server> {    if (serverScope) {      throw new Error("Test server is already running");    }     const stateDir = options.stateDir ?? makeTempDir("t3code-ws-state-");    const scope = await Effect.runPromise(Scope.make("sequential"));    const persistenceLayer = options.persistenceLayer ?? SqlitePersistenceMemory;    const providerLayer = options.providerLayer ?? makeServerProviderLayer();    const providerHealthLayer = Layer.succeed(      ProviderHealth,      options.providerHealth ?? defaultProviderHealthService,    );    const openLayer = Layer.succeed(Open, options.open ?? defaultOpenService);    const serverConfigLayer = Layer.succeed(ServerConfig, {

Lines 1–500 of 1815.