curriculum-project-hub/hub/test/unit/provider-proxy.test.ts

import { createServer } from "node:http";
import { once } from "node:events";
import { afterEach, describe, expect, it } from "vitest";
import { openProviderProxyLease, type AgentProviderLease } from "../../src/connections/providerProxy.js";

const leases: AgentProviderLease[] = [];
const upstreamServers: ReturnType<typeof createServer>[] = [];

afterEach(async () => {
  await Promise.all(leases.splice(0).map((lease) => lease.close()));
  await Promise.all(upstreamServers.splice(0).map(async (server) => {
    server.closeAllConnections();
    if (server.listening) {
      server.close();
      await once(server, "close");
    }
  }));
});

describe("run-scoped provider proxy", () => {
  it("keeps the Organization provider credential out of the Agent environment", async () => {
    const observed: Array<{ authorization: string | undefined; apiKey: string | undefined; body: string }> = [];
    const upstream = createServer((request, response) => {
      const chunks: Buffer[] = [];
      request.on("data", (chunk: Buffer) => chunks.push(chunk));
      request.on("end", () => {
        observed.push({
          authorization: request.headers.authorization,
          apiKey: header(request.headers["x-api-key"]),
          body: Buffer.concat(chunks).toString("utf8"),
        });
        response.writeHead(200, { "content-type": "application/json" });
        response.end(JSON.stringify({ ok: true }));
      });
    });
    upstreamServers.push(upstream);
    upstream.listen(0, "127.0.0.1");
    await once(upstream, "listening");
    const address = upstream.address();
    if (address === null || typeof address === "string") throw new Error("expected upstream TCP address");

    const lease = await openProviderProxyLease({
      baseUrl: `http://127.0.0.1:${address.port}/api`,
      authToken: "customer-provider-auth-token",
      anthropicApiKey: "customer-provider-api-key",
    });
    leases.push(lease);

    const serializedAgentEnv = JSON.stringify(lease.sdkEnv);
    expect(serializedAgentEnv).not.toContain("customer-provider-auth-token");
    expect(serializedAgentEnv).not.toContain("customer-provider-api-key");
    expect(serializedAgentEnv).not.toContain(String(address.port));
    expect(lease.sdkEnv).toMatchObject({
      ANTHROPIC_BASE_URL: expect.stringMatching(/^http:\/\/127\.0\.0\.1:\d+$/),
      ANTHROPIC_AUTH_TOKEN: expect.any(String),
      ANTHROPIC_API_KEY: "",
    });

    const response = await fetch(`${lease.sdkEnv.ANTHROPIC_BASE_URL}/v1/messages`, {
      method: "POST",
      headers: {
        authorization: `Bearer ${lease.sdkEnv.ANTHROPIC_AUTH_TOKEN}`,
        "content-type": "application/json",
        "anthropic-version": "2023-06-01",
      },
      body: JSON.stringify({ model: "test" }),
    });

    expect(response.status).toBe(200);
    await expect(response.json()).resolves.toEqual({ ok: true });
    expect(observed).toEqual([{
      authorization: "Bearer customer-provider-auth-token",
      apiKey: "customer-provider-api-key",
      body: JSON.stringify({ model: "test" }),
    }]);
  });

  it("rejects requests without the run capability before contacting upstream", async () => {
    let upstreamRequests = 0;
    const diagnostics: unknown[] = [];
    const upstream = createServer((_request, response) => {
      upstreamRequests += 1;
      response.end("unexpected");
    });
    upstreamServers.push(upstream);
    upstream.listen(0, "127.0.0.1");
    await once(upstream, "listening");
    const address = upstream.address();
    if (address === null || typeof address === "string") throw new Error("expected upstream TCP address");
    const lease = await openProviderProxyLease(
      {
        baseUrl: `http://127.0.0.1:${address.port}`,
        authToken: "customer-secret",
        anthropicApiKey: "",
      },
      { onDiagnostic: (diagnostic) => diagnostics.push(diagnostic) },
    );
    leases.push(lease);

    const response = await fetch(`${lease.sdkEnv.ANTHROPIC_BASE_URL}/v1/messages`, {
      method: "POST",
      body: "{}",
    });

    expect(response.status).toBe(401);
    expect(upstreamRequests).toBe(0);
    expect(diagnostics).toEqual([{ code: "provider_proxy_unauthorized", category: "authorization" }]);
  });

  it("does not forward provider credentials across an upstream redirect", async () => {
    let redirectedRequests = 0;
    const diagnostics: unknown[] = [];
    const redirectTarget = createServer((_request, response) => {
      redirectedRequests += 1;
      response.end("leaked");
    });
    upstreamServers.push(redirectTarget);
    redirectTarget.listen(0, "127.0.0.1");
    await once(redirectTarget, "listening");
    const targetAddress = redirectTarget.address();
    if (targetAddress === null || typeof targetAddress === "string") throw new Error("expected redirect target address");

    const upstream = createServer((_request, response) => {
      response.writeHead(307, { location: `http://127.0.0.1:${targetAddress.port}/capture` });
      response.end();
    });
    upstreamServers.push(upstream);
    upstream.listen(0, "127.0.0.1");
    await once(upstream, "listening");
    const upstreamAddress = upstream.address();
    if (upstreamAddress === null || typeof upstreamAddress === "string") throw new Error("expected upstream address");
    const lease = await openProviderProxyLease(
      {
        baseUrl: `http://127.0.0.1:${upstreamAddress.port}`,
        authToken: "customer-secret",
        anthropicApiKey: "",
      },
      { onDiagnostic: (diagnostic) => diagnostics.push(diagnostic) },
    );
    leases.push(lease);

    const response = await fetch(`${lease.sdkEnv.ANTHROPIC_BASE_URL}/v1/messages`, {
      headers: { authorization: `Bearer ${lease.sdkEnv.ANTHROPIC_AUTH_TOKEN}` },
      redirect: "manual",
    });

    expect(response.status).toBe(502);
    await expect(response.text()).resolves.toBe("provider redirect refused");
    expect(redirectedRequests).toBe(0);
    expect(diagnostics).toEqual([{ code: "provider_proxy_redirect_refused", category: "redirect" }]);
  });
});

function header(value: string | string[] | undefined): string | undefined {
  return Array.isArray(value) ? value[0] : value;
}