@slashfi/agents-sdk 0.12.0 → 0.14.0

package/src/key-manager.test.ts

@@ -0,0 +1,273 @@
+import { describe, test, expect, afterEach } from "bun:test";
+import { createKeyManager, type KeyStore, type StoredKey, type KeyManager } from "./key-manager";
+import { jwtVerify, createLocalJWKSet } from "jose";
+
+// In-memory KeyStore for testing (no DB needed)
+function createMemoryKeyStore(): KeyStore & { keys: StoredKey[] } {
+  const keys: StoredKey[] = [];
+  return {
+    keys,
+    async loadKeys() {
+      return keys.filter((k) => k.status !== "revoked");
+    },
+    async insertKey(key: StoredKey) {
+      keys.push(key);
+    },
+    async deprecateAllActive() {
+      for (const k of keys) {
+        if (k.status === "active") k.status = "deprecated";
+      }
+    },
+    async cleanupExpired() {
+      const now = Date.now();
+      const before = keys.length;
+      const remaining = keys.filter((k) => k.expiresAt.getTime() > now);
+      keys.length = 0;
+      keys.push(...remaining);
+      return before - remaining.length;
+    },
+    // Transaction support: snapshot + rollback on error
+    async transaction<T>(fn: () => Promise<T>): Promise<T> {
+      const snapshot = keys.map((k) => ({ ...k }));
+      try {
+        return await fn();
+      } catch (err) {
+        // Rollback
+        keys.length = 0;
+        keys.push(...snapshot);
+        throw err;
+      }
+    },
+  };
+}
+
+describe("KeyManager", () => {
+  let km: KeyManager;
+
+  afterEach(() => {
+    km?.stop();
+  });
+
+  test("creates initial key on startup", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      checkIntervalMs: 60_000,
+    });
+
+    const jwks = km.getJwks();
+    expect(jwks.keys).toHaveLength(1);
+    expect(jwks.keys[0].alg).toBe("ES256");
+    expect(jwks.keys[0].kid).toMatch(/^key-/);
+    // No private key exposed
+    expect(jwks.keys[0].d).toBeUndefined();
+  });
+
+  test("signs a valid JWT", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      checkIntervalMs: 60_000,
+    });
+
+    const token = await km.signJwt({ sub: "test-service" });
+    const parts = token.split(".");
+    expect(parts).toHaveLength(3);
+
+    const payload = JSON.parse(
+      Buffer.from(parts[1], "base64url").toString()
+    );
+    expect(payload.sub).toBe("test-service");
+    expect(payload.iss).toBe("http://test:3000");
+    expect(payload.exp - payload.iat).toBe(300); // default 5 min TTL
+  });
+
+  test("token verifies against JWKS", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      checkIntervalMs: 60_000,
+    });
+
+    const token = await km.signJwt({ sub: "verify-me" });
+    const jwks = km.getJwks();
+    const JWKS = createLocalJWKSet(jwks);
+    const { payload } = await jwtVerify(token, JWKS);
+    expect(payload.sub).toBe("verify-me");
+  });
+
+  test("custom token TTL", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      tokenTtlSeconds: 60,
+      checkIntervalMs: 60_000,
+    });
+
+    const token = await km.signJwt({ sub: "short-lived" });
+    const payload = JSON.parse(
+      Buffer.from(token.split(".")[1], "base64url").toString()
+    );
+    expect(payload.exp - payload.iat).toBe(60);
+  });
+
+  // ---- Rotation tests ----
+
+  test("rotation: creates new key when threshold exceeded", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      rotationThresholdMs: 0, // immediate rotation
+      checkIntervalMs: 60_000,
+    });
+
+    const initialKid = km.getJwks().keys.find(
+      (k) => store.keys.find((sk) => sk.kid === k.kid)?.status === "active"
+    )?.kid;
+
+    // Force rotation
+    await km.rotate();
+
+    const activeKeys = store.keys.filter((k) => k.status === "active");
+    expect(activeKeys).toHaveLength(1);
+    expect(activeKeys[0].kid).not.toBe(initialKid);
+
+    const deprecatedKeys = store.keys.filter((k) => k.status === "deprecated");
+    expect(deprecatedKeys.length).toBeGreaterThanOrEqual(1);
+  });
+
+  test("rotation: deprecated keys still in JWKS for verification", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      rotationThresholdMs: 0,
+      checkIntervalMs: 60_000,
+    });
+
+    await km.rotate();
+
+    const jwksKids = km.getJwks().keys.map((k) => k.kid);
+    const nonRevoked = store.keys.filter((k) => k.status !== "revoked");
+    for (const key of nonRevoked) {
+      expect(jwksKids).toContain(key.kid);
+    }
+  });
+
+  test("rotation: pre-rotation tokens still verify", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      rotationThresholdMs: 0,
+      checkIntervalMs: 60_000,
+    });
+
+    const preRotationToken = await km.signJwt({ sub: "pre-rotate" });
+    await km.rotate();
+
+    // Pre-rotation token should still verify (deprecated key still in JWKS)
+    const jwks = km.getJwks();
+    const JWKS = createLocalJWKSet(jwks);
+    const { payload } = await jwtVerify(preRotationToken, JWKS);
+    expect(payload.sub).toBe("pre-rotate");
+  });
+
+  test("rotation: exactly 1 active key", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      rotationThresholdMs: 0,
+      checkIntervalMs: 60_000,
+    });
+
+    await km.rotate();
+    await km.rotate();
+
+    const activeKeys = store.keys.filter((k) => k.status === "active");
+    expect(activeKeys).toHaveLength(1);
+  });
+
+  test("rotation: tokens fail verification after key expires and is cleaned up", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      rotationThresholdMs: 0,
+      keyLifetimeMs: 1, // 1ms — keys expire almost immediately
+      checkIntervalMs: 60_000,
+    });
+
+    const token = await km.signJwt({ sub: "will-expire" });
+
+    // Token should verify now (key is active)
+    const JWKS = createLocalJWKSet(km.getJwks());
+    const { payload } = await jwtVerify(token, JWKS);
+    expect(payload.sub).toBe("will-expire");
+
+    // Wait for key to expire, then rotate (which cleans up expired keys)
+    await new Promise((r) => setTimeout(r, 5));
+    await km.rotate();
+
+    // Old key should be gone from JWKS — verification should fail
+    const jwksAfter = km.getJwks();
+    const JWKS2 = createLocalJWKSet(jwksAfter);
+    let failed = false;
+    try {
+      await jwtVerify(token, JWKS2);
+    } catch {
+      failed = true;
+    }
+    expect(failed).toBe(true);
+  });
+
+    // ---- Transaction tests ----
+
+  test("transaction: rotate is atomic (rollback on failure)", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      checkIntervalMs: 60_000,
+    });
+
+    const initialKid = store.keys[0].kid;
+
+    // Monkey-patch insertKey to fail mid-transaction
+    const origInsert = store.insertKey;
+    store.insertKey = async () => { throw new Error("simulated failure"); };
+
+    // Rotation should fail, but state should be rolled back
+    try { await km.rotate(); } catch {}
+
+    // Original key should still be active (not deprecated)
+    const active = store.keys.filter((k) => k.status === "active");
+    expect(active).toHaveLength(1);
+    expect(active[0].kid).toBe(initialKid);
+
+    // Restore
+    store.insertKey = origInsert;
+  });
+
+  // ---- enableRotation option ----
+
+  test("enableRotation: false skips background interval", async () => {
+    const store = createMemoryKeyStore();
+    km = await createKeyManager({
+      store,
+      issuer: "http://test:3000",
+      enableRotation: false,
+    });
+
+    // Should still have an initial key
+    expect(km.getJwks().keys).toHaveLength(1);
+    // But stop() should be a no-op (no interval to clear)
+    km.stop();
+  });
+});

package/src/key-manager.ts

@@ -0,0 +1,257 @@
+/**
+ * JWKS Key Manager
+ *
+ * Manages ES256 signing keys with automatic rotation and revocation.
+ * Store-agnostic — provide a KeyStore implementation for your DB.
+ *
+ * Features:
+ * - Automatic key rotation on a configurable schedule
+ * - Key lifecycle: active → deprecated → revoked → cleaned up
+ * - Multi-instance safe: checks DB before rotating (another instance may have already rotated)
+ * - Periodic background checks (configurable interval)
+ * - Exposes JWKS for /.well-known/jwks.json
+ * - Signs JWTs with the active key
+ */
+
+import {
+  generateKeyPair,
+  exportJWK,
+  importJWK,
+  SignJWT,
+  type JWK,
+} from "jose";
+
+// ── Types ──
+
+export type KeyStatus = "active" | "deprecated" | "revoked";
+
+export interface StoredKey {
+  kid: string;
+  alg: string;
+  status: KeyStatus;
+  publicJwk: JWK;
+  privateJwk: JWK;
+  createdAt: Date;
+  expiresAt: Date;
+}
+
+interface CachedKey extends StoredKey {
+  privateKey: CryptoKey;
+}
+
+/**
+ * Pluggable store interface for key persistence.
+ * Implement this for your database (CockroachDB, Postgres, SQLite, etc.)
+ */
+export interface KeyStore {
+  /** Load all non-revoked keys */
+  loadKeys(): Promise<StoredKey[]>;
+  /** Insert a new key */
+  insertKey(key: StoredKey): Promise<void>;
+  /** Set all active keys to deprecated */
+  deprecateAllActive(): Promise<void>;
+  /** Delete expired keys, return count deleted */
+  cleanupExpired(): Promise<number>;
+  /**
+   * Run operations atomically. Implementations with transaction support
+   * should wrap the callback in a DB transaction to ensure rotate() is
+   * atomic (load + deprecate + insert + cleanup all succeed or all fail).
+   * For stores without real tx support, pass a no-op wrapper that runs fn sequentially.
+   */
+  transaction<T>(fn: () => Promise<T>): Promise<T>;
+}
+
+export interface KeyManager {
+  /** Get the JWKS for /.well-known/jwks.json */
+  getJwks(): { keys: JWK[] };
+  /** Sign a JWT with the active key */
+  signJwt(claims: Record<string, unknown>): Promise<string>;
+  /** Force a key rotation */
+  rotate(): Promise<void>;
+  /** Stop the background rotation check */
+  stop(): void;
+}
+
+export interface KeyManagerOptions {
+  /** Key store implementation */
+  store: KeyStore;
+  /** Issuer URL for the iss claim */
+  issuer: string;
+  /** How often to check if rotation is needed (default: 5 min) */
+  checkIntervalMs?: number;
+  /** Max age of an active key before rotation (default: 1 hour) */
+  rotationThresholdMs?: number;
+  /** How long deprecated keys stay in JWKS for verification (default: 2 hours) */
+  keyLifetimeMs?: number;
+  /** Token TTL in seconds (default: 300 = 5 min) */
+  tokenTtlSeconds?: number;
+  /** Enable background key rotation (default: true). Set to false on read-only replicas. */
+  enableRotation?: boolean;
+}
+
+// ── Constants ──
+
+const FIVE_MINUTES = 5 * 60 * 1000;
+const ONE_HOUR = 60 * 60 * 1000;
+const TWO_HOURS = 2 * ONE_HOUR;
+const ALG = "ES256";
+
+// ── Key generation ──
+
+function generateKid(): string {
+  return `key-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+}
+
+async function generateNewKey(keyLifetimeMs: number): Promise<StoredKey> {
+  const { privateKey, publicKey } = await generateKeyPair(ALG, { extractable: true });
+  const kid = generateKid();
+
+  const publicJwk = await exportJWK(publicKey);
+  publicJwk.kid = kid;
+  publicJwk.alg = ALG;
+  publicJwk.use = "sig";
+
+  const privateJwk = await exportJWK(privateKey);
+  privateJwk.kid = kid;
+  privateJwk.alg = ALG;
+
+  return {
+    kid,
+    alg: ALG,
+    status: "active",
+    publicJwk,
+    privateJwk,
+    createdAt: new Date(),
+    expiresAt: new Date(Date.now() + keyLifetimeMs),
+  };
+}
+
+async function toCachedKey(stored: StoredKey): Promise<CachedKey> {
+  const privateKey = await importJWK(stored.privateJwk, ALG) as CryptoKey;
+  return { ...stored, privateKey };
+}
+
+// ── Key Manager ──
+
+export async function createKeyManager(opts: KeyManagerOptions): Promise<KeyManager> {
+  const {
+    store,
+    issuer,
+    checkIntervalMs = FIVE_MINUTES,
+    rotationThresholdMs = ONE_HOUR,
+    keyLifetimeMs = TWO_HOURS,
+    tokenTtlSeconds = 300,
+    enableRotation = true,
+  } = opts;
+
+  let keys: CachedKey[] = [];
+
+  /** Generate a new key, deprecate old ones, cleanup expired, refresh cache — all in one transaction */
+  async function rotate(): Promise<void> {
+    await store.transaction(async () => {
+      const newKey = await generateNewKey(keyLifetimeMs);
+      await store.deprecateAllActive();
+      await store.insertKey(newKey);
+      await store.cleanupExpired();
+      const updated = await store.loadKeys();
+      keys = await Promise.all(updated.map(toCachedKey));
+    });
+  }
+
+  /** Check if rotation is needed and rotate if so — all within a single transaction */
+  async function checkAndRotate(): Promise<void> {
+    // Quick check against cache — no DB/store hit if key is fresh
+    const cached = keys.find((k) => k.status === "active");
+    if (cached) {
+      const age = Date.now() - cached.createdAt.getTime();
+      if (age < rotationThresholdMs) return;
+    }
+
+    // Cache says stale (or empty) — take a lock via transaction to check + rotate atomically
+    await store.transaction(async () => {
+      const stored = await store.loadKeys();
+      const active = stored.find((k) => k.status === "active");
+
+      if (active) {
+        const age = Date.now() - active.createdAt.getTime();
+        if (age < rotationThresholdMs) {
+          // Another instance already rotated — just update our cache
+          keys = await Promise.all(stored.map(toCachedKey));
+          return;
+        }
+      }
+
+      // Still stale (or no active key) — rotate within this tx
+      const newKey = await generateNewKey(keyLifetimeMs);
+      await store.deprecateAllActive();
+      await store.insertKey(newKey);
+      await store.cleanupExpired();
+
+      // Refresh cache inside the tx for a consistent read
+      const updated = await store.loadKeys();
+      keys = await Promise.all(updated.map(toCachedKey));
+    });
+  }
+
+  // Initial load + ensure we have at least one key
+  // Initial load from store
+  const stored = await store.loadKeys();
+  keys = await Promise.all(stored.map(toCachedKey));
+  if (!keys.some((k) => k.status === "active")) {
+    if (enableRotation) {
+      await rotate();
+    } else {
+      // Read-only mode: generate a key in memory only (no store writes)
+      // This ensures signJwt works even without rotation enabled
+      const newKey = await generateNewKey(keyLifetimeMs);
+      keys.push(await toCachedKey(newKey));
+    }
+  } else if (enableRotation) {
+    await checkAndRotate();
+  }
+
+  // Periodic background check (only if rotation enabled)
+  const interval = enableRotation
+    ? setInterval(async () => {
+        try {
+          await checkAndRotate();
+        } catch (err) {
+          console.error("[key-manager] Check/rotation failed:", err);
+        }
+      }, checkIntervalMs)
+    : null;
+
+  function getActiveKey(): CachedKey {
+    const active = keys.find((k) => k.status === "active");
+    if (!active) throw new Error("[key-manager] No active signing key");
+    return active;
+  }
+
+  return {
+    getJwks(): { keys: JWK[] } {
+      return {
+        keys: keys
+          .filter((k) => k.status !== "revoked")
+          .map((k) => k.publicJwk),
+      };
+    },
+
+    async signJwt(claims: Record<string, unknown>): Promise<string> {
+      const key = getActiveKey();
+      return new SignJWT({ ...claims } as any)
+        .setProtectedHeader({ alg: ALG, kid: key.kid })
+        .setIssuer(issuer)
+        .setIssuedAt()
+        .setExpirationTime(`${tokenTtlSeconds}s`)
+        .sign(key.privateKey);
+    },
+
+    async rotate(): Promise<void> {
+      await rotate();
+    },
+
+    stop(): void {
+      if (interval) clearInterval(interval);
+    },
+  };
+}

package/src/server.ts

@@ -100,6 +100,8 @@ export interface AgentServerOptions {
   signingKey?: SigningKey;
   /** OAuth identity provider for cross-registry user linking */
   oauthIdentityProvider?: OAuthIdentityProvider;
+  /** Key store for managed key rotation (if provided, uses createKeyManager instead of simple key gen) */
+  keyStore?: import("./key-manager.js").KeyStore;
 }
 
 export interface AgentServer {
@@ -588,6 +590,14 @@ export function createAgentServer(
   // OAuth2 token handler
   // ──────────────────────────────────────────
 
+  // Resolve public-facing base URL, respecting reverse proxy headers
+  const resolveBaseUrl = (r: Request): string => {
+    const fwdProto = r.headers.get("x-forwarded-proto");
+    const fwdHost = r.headers.get("x-forwarded-host");
+    if (fwdProto && fwdHost) return `${fwdProto}://${fwdHost}`;
+    return new URL(r.url).origin;
+  };
+
   async function handleOAuthToken(req: Request): Promise<Response> {
     if (!authConfig) {
       return jsonResponse({ error: "auth_not_configured" }, 404);
@@ -638,27 +648,26 @@ export function createAgentServer(
         try {
           const assertionParts = assertion.split(".");
           if (assertionParts.length === 3) {
-            const assertionPayload = JSON.parse(atob(assertionParts[1].replace(/-/g, "+").replace(/_/g, "/"))) as any;
+            const assertionPayload = JSON.parse(Buffer.from(assertionParts[1], "base64url").toString()) as any;
             if (assertionPayload.type === "agent-registry" && assertionPayload.iss) {
-              // Find or create @remote-registry agent and store the reverse connection
-              const rrAgent = registry.get("@remote-registry") ?? registry.get("/agents/@remote-registry");
-              if (rrAgent) {
-                const setupTool = (rrAgent as any).tools?.find((t: any) => t.name === "setup_integration");
-                if (setupTool?.execute) {
-                  try {
-                    await setupTool.execute(
-                      { url: assertionPayload.iss, name: assertionPayload.name ?? "remote-registry" },
-                      { callerId: "system", callerType: "system", tenantId: "default", agentPath: "@remote-registry" },
-                    );
-                    console.error(`[jwt_exchange] Reverse connection stored for ${assertionPayload.iss}`);
-                  } catch (setupErr) {
-                    console.error(`[jwt_exchange] Reverse registration setup failed:`, setupErr);
-                  }
-                } else {
-                  console.error("[jwt_exchange] @remote-registry has no setup_integration tool — reverse registration skipped");
-                }
+              // Use add_connection (direct store) instead of setup_integration (which would cause infinite loop)
+              const addResult = await registry.call({
+                action: "execute_tool",
+                path: "@remote-registry",
+                tool: "add_connection",
+                params: {
+                  id: assertionPayload.name ?? "remote-registry",
+                  name: assertionPayload.name ?? "remote-registry",
+                  url: assertionPayload.iss,
+                  remoteTenantId: assertionPayload.tenantId ?? "default",
+                },
+                callerId: "system",
+                callerType: "system",
+              });
+              if (addResult.success) {
+                console.error(`[jwt_exchange] Reverse connection stored for ${assertionPayload.iss}`);
               } else {
-                console.error("[jwt_exchange] @remote-registry agent not found — reverse registration skipped");
+                console.error(`[jwt_exchange] Reverse connection failed:`, (addResult as any).error);
               }
             }
           }
@@ -675,7 +684,7 @@ export function createAgentServer(
 
         // User not linked yet — needs OAuth identity linking
         if (exchangeResult.needsAuth) {
-          const baseUrl = new URL(req.url).origin;
+          const baseUrl = resolveBaseUrl(req);
           const authorizeUrl = new URL(`${baseUrl}${basePath}/oauth/authorize`);
           authorizeUrl.searchParams.set("token", assertion);
           if (params.redirect_uri) {
@@ -707,7 +716,7 @@ export function createAgentServer(
             },
             sigKey.privateKey,
             sigKey.kid,
-            new URL(req.url).origin,
+            resolveBaseUrl(req),
             `${authConfig.tokenTtl ?? 3600}s`,
           );
 
@@ -789,6 +798,7 @@ export function createAgentServer(
 
   async function fetch(req: Request): Promise<Response> {
     try {
+
       const url = new URL(req.url);
       const path = url.pathname.replace(basePath, "") || "/";
 
@@ -856,17 +866,23 @@ export function createAgentServer(
           return cors ? addCors(res) : res;
         }
 
-        // Verify the JWT against trusted issuers
+        // Verify the JWT against trusted issuers (from store, falling back to config)
         let claims: Record<string, unknown> | null = null;
-        const issuerUrls = configTrustedIssuers.map(i => typeof i === "string" ? i : i.issuer);
-        for (const issuerUrl of issuerUrls) {
+        const storeIssuers = authConfig?.store?.listTrustedIssuers
+          ? await authConfig.store.listTrustedIssuers()
+          : [];
+        const configIssuerUrls = configTrustedIssuers.map(i => typeof i === "string" ? i : i.issuer);
+        const allIssuerUrls = [...new Set([...storeIssuers, ...configIssuerUrls])];
+        console.log("[oauth/authorize] storeIssuers:", storeIssuers.length, "configIssuers:", configIssuerUrls.length, "total:", allIssuerUrls.length, "urls:", allIssuerUrls);
+        for (const issuerUrl of allIssuerUrls) {
           try {
             const result = await verifyJwtFromIssuer(token, issuerUrl);
+            console.log("[oauth/authorize] verify", issuerUrl, "->", result ? "OK" : "null");
             if (result) {
               claims = result as unknown as Record<string, unknown>;
               break;
             }
-          } catch { /* try next issuer */ }
+          } catch (e: any) { console.log("[oauth/authorize] verify", issuerUrl, "-> ERROR:", e.message); }
         }
         if (!claims) {
           const res = jsonResponse(
@@ -876,7 +892,7 @@ export function createAgentServer(
           return cors ? addCors(res) : res;
         }
 
-        const baseUrl = new URL(req.url).origin;
+        const baseUrl = resolveBaseUrl(req);
         const scope = url.searchParams.get("scope") ?? undefined;
         const res = await oauthIdentityProvider.authorize(req, {
           token,
@@ -897,7 +913,7 @@ export function createAgentServer(
           );
           return cors ? addCors(res) : res;
         }
-        const baseUrl = new URL(req.url).origin;
+        const baseUrl = resolveBaseUrl(req);
         const res = await oauthIdentityProvider.callback(req, {
           baseUrl: baseUrl + basePath,
         });
@@ -921,7 +937,7 @@ export function createAgentServer(
 
       // ── GET /.well-known/configuration → Server discovery ──
       if (path === "/.well-known/configuration" && req.method === "GET") {
-        const baseUrl = new URL(req.url).origin;
+        const baseUrl = resolveBaseUrl(req);
         const res = jsonResponse({
           issuer: baseUrl,
           jwks_uri: `${baseUrl}/.well-known/jwks.json`,