pi-commandcode-provider/tests/helpers.ts

import { createServer, type IncomingHttpHeaders, type Server } from "node:http";

import {
  createStreamCommandCode,
  type AssistantMessageEvent,
  type AssistantMessageEventStreamLike,
  type ContextLike,
  type CoreDependencies,
  type ModelLike,
  type Usage,
} from "../src/core.ts";

export function createTestEventStream(): AssistantMessageEventStreamLike {
  const events: AssistantMessageEvent[] = [];
  const waiters: Array<() => void> = [];
  let ended = false;

  const wake = () => {
    const waiter = waiters.shift();
    if (waiter) waiter();
  };

  return {
    push(event: AssistantMessageEvent) {
      events.push(event);
      wake();
    },
    end() {
      ended = true;
      while (waiters.length > 0) wake();
    },
    [Symbol.asyncIterator]() {
      let index = 0;
      return {
        async next(): Promise<IteratorResult<AssistantMessageEvent>> {
          while (index >= events.length && !ended) {
            await new Promise<void>((resolve) => waiters.push(resolve));
          }
          if (index < events.length) {
            const value = events[index];
            index += 1;
            return { done: false, value };
          }
          return { done: true, value: undefined };
        },
      };
    },
  };
}

export async function collectEvents(
  stream: AssistantMessageEventStreamLike,
  timeoutMs = 2_000,
): Promise<AssistantMessageEvent[]> {
  const events: AssistantMessageEvent[] = [];

  const collect = async () => {
    for await (const event of stream) {
      events.push(event);
      if (event.type === "done" || event.type === "error") break;
    }
    return events;
  };

  return await Promise.race([
    collect(),
    new Promise<AssistantMessageEvent[]>((_, reject) => {
      setTimeout(
        () =>
          reject(
            new Error(
              `Timed out collecting stream events after ${timeoutMs}ms`,
            ),
          ),
        timeoutMs,
      );
    }),
  ]);
}

export function makeModel(overrides: Partial<ModelLike> = {}): ModelLike {
  return {
    id: "deepseek/deepseek-v4-flash",
    api: "commandcode-custom",
    provider: "commandcode",
    maxTokens: 384_000,
    ...overrides,
  };
}

export function makeContext(overrides: Partial<ContextLike> = {}): ContextLike {
  return {
    systemPrompt: "You are a test assistant.",
    messages: [{ role: "user", content: "hello" }],
    tools: [],
    ...overrides,
  };
}

export interface TestDepsResult {
  streamCommandCode: ReturnType<typeof createStreamCommandCode>;
  calculatedUsages: Usage[];
}

export function createTestDeps(
  overrides: Partial<CoreDependencies> = {},
): TestDepsResult {
  const calculatedUsages: Usage[] = [];
  const streamCommandCode = createStreamCommandCode({
    createStream: createTestEventStream,
    calculateCost: (_model, usage) => {
      calculatedUsages.push({
        ...usage,
        cost: { ...usage.cost },
      });
    },
    env: {},
    authPaths: [],
    now: () => new Date("2026-05-05T12:00:00Z").getTime(),
    uuid: () => "00000000-0000-4000-8000-000000000000",
    cwd: () => "/repo",
    ...overrides,
  });
  return { streamCommandCode, calculatedUsages };
}

type SuccessPlan = {
  type: "success";
  status?: number;
  events?: string[];
  chunks?: string[];
  delays?: number[];
  hangAfterLast?: boolean;
};

type ErrorPlan = {
  type: "error";
  status: number;
  body: string;
};

export type ResponsePlan = SuccessPlan | ErrorPlan;

function headersToRecord(headers: IncomingHttpHeaders): Record<string, string> {
  const out: Record<string, string> = {};
  for (const [key, value] of Object.entries(headers)) {
    if (typeof value === "string") out[key] = value;
    else if (Array.isArray(value)) out[key] = value.join(", ");
  }
  return out;
}

export interface MockCommandCodeServer {
  baseUrl(): string;
  mockResponse(plan: ResponsePlan): void;
  reset(): void;
  close(): Promise<void>;
  lastRequestBody(): unknown;
  lastRequestHeaders(): Record<string, string>;
  requestCount(): number;
  responseClosedBeforeEnd(): boolean;
}

export async function startMockCommandCodeServer(): Promise<MockCommandCodeServer> {
  let nextPlan: ResponsePlan = { type: "success", events: [] };
  let lastBody: unknown;
  let lastHeaders: Record<string, string> = {};
  let requests = 0;
  let closedBeforeEnd = false;
  let port = 0;

  const server: Server = createServer((req, res) => {
    if (req.method !== "POST" || req.url !== "/alpha/generate") {
      res.writeHead(404);
      res.end("Not found");
      return;
    }

    requests += 1;
    lastHeaders = headersToRecord(req.headers);
    let body = "";
    req.on("data", (chunk: Buffer) => {
      body += chunk.toString("utf-8");
    });
    req.on("end", () => {
      try {
        const parsed: unknown = JSON.parse(body);
        lastBody = parsed;
      } catch {
        lastBody = undefined;
      }

      const plan = nextPlan;
      if (plan.type === "error") {
        res.writeHead(plan.status, { "Content-Type": "text/plain" });
        res.end(plan.body);
        return;
      }

      res.writeHead(plan.status ?? 200, {
        "Content-Type": "text/plain; charset=utf-8",
        "Transfer-Encoding": "chunked",
      });

      let ended = false;
      res.on("close", () => {
        if (!ended) closedBeforeEnd = true;
      });

      const chunks =
        plan.chunks ?? (plan.events ?? []).map((event) => `${event}\n`);
      const delays = plan.delays ?? chunks.map(() => 0);
      let index = 0;

      const sendNext = () => {
        if (index >= chunks.length) {
          if (!plan.hangAfterLast) {
            ended = true;
            res.end();
          }
          return;
        }

        res.write(chunks[index]);
        index += 1;
        if (index < chunks.length) {
          setTimeout(sendNext, delays[index] ?? 0);
        } else if (!plan.hangAfterLast) {
          ended = true;
          res.end();
        }
      };

      sendNext();
    });
  });

  await new Promise<void>((resolve) => {
    server.listen(0, () => {
      const address = server.address();
      if (typeof address === "object" && address) port = address.port;
      resolve();
    });
  });

  return {
    baseUrl: () => `http://127.0.0.1:${port}`,
    mockResponse(plan: ResponsePlan) {
      nextPlan = plan;
    },
    reset() {
      nextPlan = { type: "success", events: [] };
      lastBody = undefined;
      lastHeaders = {};
      requests = 0;
      closedBeforeEnd = false;
    },
    close() {
      return new Promise<void>((resolve) => server.close(() => resolve()));
    },
    lastRequestBody: () => lastBody,
    lastRequestHeaders: () => lastHeaders,
    requestCount: () => requests,
    responseClosedBeforeEnd: () => closedBeforeEnd,
  };
}

export function objectAt(value: unknown, path: readonly string[]): unknown {
  let current = value;
  for (const key of path) {
    if (Array.isArray(current)) {
      const index = Number(key);
      if (!Number.isInteger(index)) return undefined;
      current = current[index];
      continue;
    }
    if (typeof current !== "object" || current === null) return undefined;
    current = Object.getOwnPropertyDescriptor(current, key)?.value;
  }
  return current;
}