Research-Stack/4-Infrastructure/workers/s2-proxy/index.js

/**
 * s2-proxy Cloudflare Worker
 *
 * Proxies requests to the Semantic Scholar Graph API with:
 *   - Per-worker API key (stored as secret S2_API_KEY)
 *   - KV caching of responses (24h TTL, configurable)
 *   - Simple rate-limit header passthrough so coordinator knows to back off
 *
 * Deploy N copies of this worker (each with its own API key secret) to get
 * N × 1 req/sec aggregate throughput against S2.
 *
 * Endpoint:
 *   GET /fetch?path=/graph/v1/paper/{id}&fields=paperId,title,year,references,citations
 *
 * The coordinator calls  https://<worker>.workers.dev/fetch?path=<s2_path>
 * and gets back the raw S2 JSON (or a cached copy).
 *
 * Environment bindings required (wrangler.toml):
 *   S2_API_KEY  — secret (s2 api key for this worker instance)
 *   S2_CACHE    — KV namespace binding
 */

const S2_BASE = "https://api.semanticscholar.org";
const CACHE_TTL = 86400; // 24 hours

export default {
  async fetch(request, env, ctx) {
    // ── CORS preflight ────────────────────────────────────────────────────────
    if (request.method === "OPTIONS") {
      return new Response(null, {
        headers: {
          "Access-Control-Allow-Origin": "*",
          "Access-Control-Allow-Methods": "GET",
          "Access-Control-Allow-Headers": "Content-Type",
        },
      });
    }

    const url = new URL(request.url);

    // ── Health check ──────────────────────────────────────────────────────────
    if (url.pathname === "/health") {
      return new Response(JSON.stringify({ ok: true, key_set: !!env.S2_API_KEY }), {
        headers: { "Content-Type": "application/json" },
      });
    }

    // ── Main proxy endpoint ───────────────────────────────────────────────────
    if (url.pathname !== "/fetch") {
      return new Response("use /fetch?path=/graph/v1/paper/...", { status: 400 });
    }

    const s2Path = url.searchParams.get("path");
    if (!s2Path) {
      return new Response("missing path param", { status: 400 });
    }

    // Preserve any query params the caller passed (fields, limit, offset, etc.)
    const forwardParams = new URLSearchParams(url.search);
    forwardParams.delete("path");
    const s2Url = `${S2_BASE}${s2Path}${
      forwardParams.toString() ? "?" + forwardParams.toString() : ""
    }`;

    // ── KV cache lookup ───────────────────────────────────────────────────────
    const cacheKey = s2Url;
    if (env.S2_CACHE) {
      const cached = await env.S2_CACHE.get(cacheKey);
      if (cached !== null) {
        return new Response(cached, {
          headers: {
            "Content-Type": "application/json",
            "X-S2-Cache": "HIT",
            "Access-Control-Allow-Origin": "*",
          },
        });
      }
    }

    // ── Fetch from S2 ─────────────────────────────────────────────────────────
    const headers = { "User-Agent": "research-stack-crawler/1.0" };
    if (env.S2_API_KEY) {
      headers["x-api-key"] = env.S2_API_KEY;
    }

    let s2Resp;
    try {
      s2Resp = await fetch(s2Url, { headers });
    } catch (err) {
      return new Response(JSON.stringify({ error: String(err) }), {
        status: 502,
        headers: { "Content-Type": "application/json" },
      });
    }

    // Pass rate-limit headers back so coordinator can see them
    const outHeaders = {
      "Content-Type": "application/json",
      "X-S2-Cache": "MISS",
      "Access-Control-Allow-Origin": "*",
    };
    for (const h of ["x-ratelimit-limit", "x-ratelimit-remaining", "retry-after"]) {
      if (s2Resp.headers.has(h)) outHeaders[h] = s2Resp.headers.get(h);
    }

    const body = await s2Resp.text();

    // Cache successful responses only
    if (s2Resp.ok && env.S2_CACHE) {
      ctx.waitUntil(env.S2_CACHE.put(cacheKey, body, { expirationTtl: CACHE_TTL }));
    }

    return new Response(body, { status: s2Resp.status, headers: outHeaders });
  },
};