kiloServerManager.test.ts

import { ApprovalRequestId, ThreadId, TurnId, type ProviderRuntimeEvent } from "@t3tools/contracts";
import { describe, expect, it, vi } from "vitest";

import { KiloServerManager } from "./kiloServerManager.ts";
import {
  PROVIDER,
  type KiloClient,
  type KiloProviderSession,
  type KiloSessionContext,
  type SharedServerState,
} from "./kilo/types.ts";

vi.mock("./kilo/serverLifecycle.ts", async () => {
  const actual = await vi.importActual<typeof import("./kilo/serverLifecycle.ts")>(
    "./kilo/serverLifecycle.ts",
  );
  return {
    ...actual,
    ensureServer: vi.fn(),
  };
});

import * as ServerLifecycle from "./kilo/serverLifecycle.ts";

class TestKiloServerManager extends KiloServerManager {
  seedSession(context: KiloSessionContext) {
    (this as unknown as { sessions: Map<ThreadId, KiloSessionContext> }).sessions.set(
      context.threadId,
      context,
    );
  }
}

function createClient() {
  return {
    session: {
      get: vi.fn(async () => ({})),
      create: vi.fn(async () => ({})),
      promptAsync: vi.fn(async () => ({})),
      abort: vi.fn(async () => ({})),
      messages: vi.fn(async () => []),
      revert: vi.fn(async () => ({})),
      unrevert: vi.fn(async () => ({})),
    },
    permission: {
      reply: vi.fn(async () => ({})),
    },
    question: {
      reply: vi.fn(async () => ({})),
    },
    provider: {
      list: vi.fn(async () => ({ data: { all: [], connected: [] } })),
    },
    config: {
      providers: vi.fn(async () => ({ data: { providers: [] } })),
    },
    event: {
      subscribe: vi.fn(async () => ({ stream: (async function* () {})() })),
    },
  } satisfies KiloClient;
}

function createContext(client: KiloClient): KiloSessionContext {
  const now = new Date().toISOString();
  return {
    threadId: ThreadId.make("thread-kilo"),
    directory: process.cwd(),
    workspace: "/workspace/project",
    client,
    providerSessionId: "session-kilo",
    pendingPermissions: new Map([
      [
        ApprovalRequestId.make("approval-kilo"),
        {
          requestId: ApprovalRequestId.make("approval-kilo"),
          requestType: "exec_command_approval",
        },
      ],
    ]),
    pendingQuestions: new Map(),
    partStreamById: new Map(),
    messageIds: [],
    streamAbortController: new AbortController(),
    streamTask: Promise.resolve(),
    session: {
      provider: PROVIDER,
      status: "running",
      runtimeMode: "approval-required",
      threadId: ThreadId.make("thread-kilo"),
      createdAt: now,
      updatedAt: now,
      resumeCursor: { sessionId: "session-kilo" },
      activeTurnId: TurnId.make("turn-kilo"),
    } as KiloProviderSession,
    activeTurnId: TurnId.make("turn-kilo"),
    lastError: undefined,
  };
}

describe("KiloServerManager.respondToRequest", () => {
  it("aborts the active turn when the user cancels a pending approval", async () => {
    const manager = new TestKiloServerManager();
    const client = createClient();
    const context = createContext(client);
    const events: ProviderRuntimeEvent[] = [];
    manager.on("event", (event) => {
      events.push(event);
    });
    manager.seedSession(context);

    await manager.respondToRequest(
      context.threadId,
      ApprovalRequestId.make("approval-kilo"),
      "cancel",
    );

    expect(client.permission.reply).toHaveBeenCalledWith({
      requestID: "approval-kilo",
      workspace: "/workspace/project",
      reply: "reject",
    });
    expect(client.session.abort).toHaveBeenCalledWith({
      sessionID: "session-kilo",
      workspace: "/workspace/project",
    });
    expect(client.permission.reply.mock.invocationCallOrder[0]).toBeLessThan(
      client.session.abort.mock.invocationCallOrder[0] ?? Number.MAX_SAFE_INTEGER,
    );
    expect(context.activeTurnId).toBeUndefined();
    expect(context.session.status).toBe("ready");
    expect(events.some((event) => event.type === "turn.completed")).toBe(true);
  });

  it("does not abort the turn for a normal rejection", async () => {
    const manager = new TestKiloServerManager();
    const client = createClient();
    const context = createContext(client);
    manager.seedSession(context);

    await manager.respondToRequest(
      context.threadId,
      ApprovalRequestId.make("approval-kilo"),
      "decline",
    );

    expect(client.permission.reply).toHaveBeenCalledWith({
      requestID: "approval-kilo",
      workspace: "/workspace/project",
      reply: "reject",
    });
    expect(client.session.abort).not.toHaveBeenCalled();
    expect(context.activeTurnId).toBe(TurnId.make("turn-kilo"));
  });
});

describe("KiloServerManager.getOrStartServer", () => {
  it("kills the spawned server child on stopAll (manager-owned lifecycle)", async () => {
    // Regression: KiloTextGeneration previously called the standalone
    // `ensureServer` helper directly and stored the returned server in a
    // local closure variable, so the manager's `stopAll()` finalizer never
    // killed the spawned child. Routing through `getOrStartServer` keeps the
    // child owned by the manager.
    const childKill = vi.fn();
    const fakeServer = {
      baseUrl: "http://127.0.0.1:1",
      child: { kill: childKill, killed: false } as unknown,
      authHeader: undefined,
    } as unknown as SharedServerState;

    const ensureServerMock = vi.mocked(ServerLifecycle.ensureServer);
    ensureServerMock.mockReset();
    ensureServerMock.mockResolvedValue({ state: fakeServer, serverPromise: undefined });

    const manager = new KiloServerManager();
    const server = await manager.getOrStartServer({ binaryPath: "kilo" });
    expect(server).toBe(fakeServer);

    manager.stopAll();
    expect(childKill).toHaveBeenCalledTimes(1);
  });

  it("serializes concurrent starts so callers share one server", async () => {
    let started = 0;
    const ensureServerMock = vi.mocked(ServerLifecycle.ensureServer);
    ensureServerMock.mockReset();
    ensureServerMock.mockImplementation(async () => {
      started += 1;
      // Simulate an async start — both callers should see the same pending
      // promise and end up with the same SharedServerState.
      await new Promise((resolve) => setTimeout(resolve, 5));
      const fakeServer = {
        baseUrl: `http://127.0.0.1:${started}`,
        child: { kill: vi.fn(), killed: false } as unknown,
        authHeader: undefined,
      } as unknown as SharedServerState;
      return { state: fakeServer, serverPromise: undefined };
    });

    const manager = new KiloServerManager();
    const [a, b] = await Promise.all([
      manager.getOrStartServer({ binaryPath: "kilo" }),
      manager.getOrStartServer({ binaryPath: "kilo" }),
    ]);
    expect(started).toBe(1);
    expect(a).toBe(b);

    manager.stopAll();
  });
});