PaperClipAI/server/src/__tests__/idle-admission.test.ts

import express from "express";
import request from "supertest";
import { request as httpRequest } from "node:http";
import { createServer as createViteServer } from "vite";
import { afterEach, describe, expect, it } from "vitest";
import { drainHttpRequestWork, idleAdmissionMiddleware, trackIdleRequestHandlers } from "../middleware/idle-admission.js";
import { idleWorkSnapshot, startTaskDrain, stopTaskDrain, readTaskDrain, beginIdleTrackedWork } from "../services/task-admission.js";

function deferred() {
  let resolve!: () => void;
  const promise = new Promise<void>(done => { resolve = done; });
  return { promise, resolve };
}
function app() { const app = express(); app.use(idleAdmissionMiddleware); return app; }
afterEach(() => stopTaskDrain());

describe("idle admission", () => {
  it("preserves legacy drains, but blocks all new work during an owned idle hold", async () => {
    const server = app();
    server.all("/{*path}", (_req, res) => res.sendStatus(204));
    trackIdleRequestHandlers(server);
    startTaskDrain({ ttlMs: 60_000 });
    await request(server).post("/api/issues").expect(204);
    startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
    for (const path of ["/api/issues", "/api/chat-webhooks/test", "/mcp", "/api/auth/callback"]) {
      await request(server).get(path).expect(503).expect("Retry-After", "1");
      await request(server).post(path).expect(503);
    }
    await request(server).get("/api/health").expect(204);
    await request(server).get("/api/instance/task-drain/").expect(204);
    await request(server).post("/api/instance/task-drain").expect(204);
    stopTaskDrain();
    await request(server).post("/api/issues").expect(204);
    expect(idleWorkSnapshot().active).toBe(0);
  });

  it("counts a nested async handler after it sends a response", async () => {
    const server = app();
    const router = express.Router();
    const done = deferred();
    router.post("/work", async (_req, res) => { res.sendStatus(202); await done.promise; });
    server.use("/api", router);
    trackIdleRequestHandlers(server);
    await request(server).post("/api/work").expect(202);
    expect(idleWorkSnapshot().active).toBe(1);
    let drained = false;
    const drain = drainHttpRequestWork(server, new AbortController().signal).then(() => { drained = true; });
    await new Promise(resolve => setImmediate(resolve));
    expect(drained).toBe(false);
    startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
    await request(server).post("/api/work").expect(503);
    done.resolve();
    await drain;
    expect(drained).toBe(true);
    expect(idleWorkSnapshot().active).toBe(0);
  });

  it("tracks ordinary middleware with non-router stack metadata", async () => {
    const server = app(), done = deferred();
    const middleware = Object.assign(async (_req: express.Request, res: express.Response) => {
      res.sendStatus(202); await done.promise;
    }, { stack: { name: "middleware metadata" } });
    server.use(middleware);
    expect(() => trackIdleRequestHandlers(server)).not.toThrow();
    await request(server).get("/work").expect(202);
    expect(idleWorkSnapshot().active).toBe(1);
    done.resolve();
    await new Promise(resolve => setImmediate(resolve));
    expect(idleWorkSnapshot().active).toBe(0);
  });

  it("supports real Vite Connect route metadata without losing async work", async () => {
    const vite = await createViteServer({ configFile: false, appType: "custom",
      server: { middlewareMode: true, watch: null }, optimizeDeps: { noDiscovery: true, include: [] } });
    const server = app(), done = deferred();
    try {
      server.use(vite.middlewares);
      server.post("/work", async (_req, res) => { res.sendStatus(202); await done.promise; });
      expect(() => trackIdleRequestHandlers(server)).not.toThrow();
      await request(server).post("/work").expect(202);
      expect(idleWorkSnapshot().active).toBe(1);
      done.resolve();
      await new Promise(resolve => setImmediate(resolve));
      expect(idleWorkSnapshot().active).toBe(0);
    } finally { await vite.close(); }
  });

  it("tracks real Vite middleware without confusing Connect mount paths with Express routes", async () => {
    const { createServer } = await import("vite");
    const vite = await createServer({
      configFile: false,
      appType: "custom",
      server: { middlewareMode: true, hmr: false, watch: null },
      optimizeDeps: { noDiscovery: true, include: [] },
    });
    const server = app(), done = deferred();
    vite.middlewares.use("/vite-work", async (_req, res) => {
      res.statusCode = 202;
      res.end();
      await done.promise;
    });
    server.use(vite.middlewares);
    try {
      expect(() => trackIdleRequestHandlers(server)).not.toThrow();
      await request(server).get("/vite-work").expect(202);
      expect(idleWorkSnapshot().active).toBe(1);
      startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
      await request(server).get("/vite-work").expect(503);
      done.resolve();
      await new Promise(resolve => setImmediate(resolve));
      expect(idleWorkSnapshot().active).toBe(0);
    } finally {
      done.resolve();
      await vite.close();
    }
  });

  it("does not mistake a client disconnect for completed accepted work", async () => {
    const server = app();
    const entered = deferred(), done = deferred(), closed = deferred();
    server.post("/work", async (_req, res) => { res.once("close", closed.resolve); entered.resolve(); await done.promise; /* client has left */ });
    trackIdleRequestHandlers(server);
    const listener = server.listen(0, "127.0.0.1");
    await new Promise<void>(resolve => listener.once("listening", resolve));
    const address = listener.address() as { port: number };
    const client = httpRequest({ host: "127.0.0.1", port: address.port, path: "/work", method: "POST" });
    client.on("error", () => undefined);
    client.end();
    try {
      await entered.promise;
      client.destroy();
      await closed.promise;
      startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
      expect(idleWorkSnapshot().active).toBe(1);
      done.resolve();
      await new Promise(resolve => setImmediate(resolve));
      expect(idleWorkSnapshot().active).toBe(0);
    } finally {
      done.resolve(); client.destroy();
      await new Promise<void>(resolve => listener.close(() => resolve()));
    }
  });

  it("settles errors through async Express error middleware", async () => {
    const server = app(), done = deferred();
    server.get("/work", async () => { throw new Error("test"); });
    server.use(async (_err: Error, _req: express.Request, res: express.Response, _next: express.NextFunction) => {
      res.sendStatus(500); await done.promise;
    });
    trackIdleRequestHandlers(server);
    await request(server).get("/work").expect(500);
    expect(idleWorkSnapshot().active).toBe(1);
    done.resolve(); await new Promise(resolve => setImmediate(resolve));
    expect(idleWorkSnapshot().active).toBe(0);
  });

  it("uses unique owners even for simultaneous starts and expires the hold", () => {
    const first = startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
    const second = startTaskDrain({ purpose: "idle", ttlMs: 60_000 });
    expect(second.ownerId).not.toBe(first.ownerId);
    expect(readTaskDrain(second.expiresAt!)).toBeNull();
  });

  it("records even work that starts and finishes between observations", () => {
    const before = idleWorkSnapshot();
    const done = beginIdleTrackedWork(); done(); done();
    expect(idleWorkSnapshot()).toEqual({ active: before.active, generation: before.generation + 2 });
  });
});