@slashfi/agents-sdk 0.11.2 → 0.13.0

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.test.ts

@@ -0,0 +1,284 @@
+import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
+import {
+  createAgentServer,
+  createAgentRegistry,
+  detectAuth,
+  resolveAuth,
+  canSeeAgent,
+} from './index';
+import type { AgentDefinition, TrustedIssuer, AgentServer } from './index';
+import { generateKeyPair, exportJWK, SignJWT } from 'jose';
+
+// ─── Helpers ─────────────────────────────────────────────────────
+
+function makeAgent(
+  path: string,
+  opts: Partial<AgentDefinition> = {},
+): AgentDefinition {
+  return {
+    path,
+    entrypoint: 'test',
+    tools: [],
+    visibility: 'internal',
+    config: { name: path.split('/').pop(), supportedActions: ['load'] },
+    ...opts,
+  } as AgentDefinition;
+}
+
+async function mcpCall(
+  port: number,
+  toolName: string,
+  args: Record<string, unknown>,
+  authToken?: string,
+) {
+  const headers: Record<string, string> = { 'Content-Type': 'application/json' };
+  if (authToken) headers.Authorization = `Bearer ${authToken}`;
+
+  const res = await fetch(`http://localhost:${port}`, {
+    method: 'POST',
+    headers,
+    body: JSON.stringify({
+      jsonrpc: '2.0',
+      id: Date.now(),
+      method: 'tools/call',
+      params: { name: toolName, arguments: args },
+    }),
+  });
+  return res.json() as Promise<any>;
+}
+
+function parseResult(rpc: any): any {
+  const text = rpc.result?.content?.[0]?.text;
+  return text ? JSON.parse(text) : null;
+}
+
+// ─── E2E: Full server auth flow ──────────────────────────────────
+//
+// These tests spin up a real createAgentServer, send actual HTTP
+// requests, and verify the complete path:
+//   HTTP request → auth resolution → handleToolCall → registry.call → access check
+//
+// This is what actually broke today: authConfig was null → resolveAuth
+// was skipped → trusted issuer tokens were ignored.
+
+describe('E2E: createAgentServer with trusted issuers', () => {
+  let privateKey: CryptoKey;
+  let publicJwk: any;
+  let jwksHttpServer: ReturnType<typeof Bun.serve>;
+  let server: AgentServer;
+  const JWKS_PORT = 19880;
+  const SDK_PORT = 19881;
+  const ISSUER_URL = `http://localhost:${JWKS_PORT}`;
+  const KID = 'test-e2e-key';
+
+  beforeAll(async () => {
+    // 1. Generate ES256 keypair and serve JWKS
+    const keyPair = await generateKeyPair('ES256', { extractable: true });
+    privateKey = keyPair.privateKey;
+    publicJwk = await exportJWK(keyPair.publicKey);
+    publicJwk.kid = KID;
+    publicJwk.alg = 'ES256';
+    publicJwk.use = 'sig';
+
+    jwksHttpServer = Bun.serve({
+      port: JWKS_PORT,
+      fetch(req) {
+        const url = new URL(req.url);
+        if (url.pathname === '/.well-known/jwks.json') {
+          return new Response(JSON.stringify({ keys: [publicJwk] }), {
+            headers: { 'Content-Type': 'application/json' },
+          });
+        }
+        return new Response('Not found', { status: 404 });
+      },
+    });
+
+    // 2. Create registry with internal + public agents
+    const registry = createAgentRegistry();
+    registry.register(makeAgent('/agents/@clock', { visibility: 'internal' }));
+    registry.register(makeAgent('/agents/public-bot', { visibility: 'public' }));
+
+    // 3. Create server with trusted issuer — NO @auth agent registered
+    //    This is the exact scenario that was broken.
+    server = createAgentServer(registry, {
+      port: SDK_PORT,
+      trustedIssuers: [{
+        issuer: ISSUER_URL,
+        scopes: ['agents:admin'],
+      }],
+    });
+    await server.start();
+  });
+
+  afterAll(() => {
+    server?.stop?.();
+    jwksHttpServer?.stop();
+  });
+
+  async function signToken(claims: Record<string, unknown> = {}): Promise<string> {
+    return new SignJWT({ sub: 'atlas-api', ...claims } as any)
+      .setProtectedHeader({ alg: 'ES256', kid: KID })
+      .setIssuer(ISSUER_URL)
+      .setIssuedAt()
+      .setExpirationTime('5m')
+      .sign(privateKey);
+  }
+
+  // ─── Core auth flow tests ───────────────────────────────────
+
+  test('system token → can load internal agent', async () => {
+    const token = await signToken();
+    const rpc = await mcpCall(SDK_PORT, 'call_agent', {
+      request: { action: 'load', path: '/agents/@clock' },
+    }, token);
+
+    const result = parseResult(rpc);
+    expect(result.success).toBe(true);
+  });
+
+  test('no token → access denied for internal agent', async () => {
+    const rpc = await mcpCall(SDK_PORT, 'call_agent', {
+      request: { action: 'load', path: '/agents/@clock' },
+    });
+
+    const result = parseResult(rpc);
+    expect(result.success).toBe(false);
+    expect(result.code).toBe('ACCESS_DENIED');
+  });
+
+  test('no token → public agent still accessible', async () => {
+    const rpc = await mcpCall(SDK_PORT, 'call_agent', {
+      request: { action: 'load', path: '/agents/public-bot' },
+    });
+
+    const result = parseResult(rpc);
+    expect(result.success).toBe(true);
+  });
+
+  test('garbage token → access denied', async () => {
+    const rpc = await mcpCall(SDK_PORT, 'call_agent', {
+      request: { action: 'load', path: '/agents/@clock' },
+    }, 'not.a.valid.jwt');
+
+    const result = parseResult(rpc);
+    expect(result.success).toBe(false);
+    expect(result.code).toBe('ACCESS_DENIED');
+  });
+
+  test('token with wrong issuer → access denied', async () => {
+    // Sign with correct key but wrong iss claim
+    const token = await new SignJWT({ sub: 'evil' } as any)
+      .setProtectedHeader({ alg: 'ES256', kid: KID })
+      .setIssuer('http://evil:9999') // not in trustedIssuers
+      .setIssuedAt()
+      .setExpirationTime('5m')
+      .sign(privateKey);
+
+    const rpc = await mcpCall(SDK_PORT, 'call_agent', {
+      request: { action: 'load', path: '/agents/@clock' },
+    }, token);
+
+    const result = parseResult(rpc);
+    expect(result.success).toBe(false);
+    expect(result.code).toBe('ACCESS_DENIED');
+  });
+
+  // ─── Visibility in list_agents ─────────────────────────────
+
+  test('list_agents without token → only public agents', async () => {
+    const rpc = await mcpCall(SDK_PORT, 'list_agents', {});
+    const result = parseResult(rpc);
+    expect(result.success).toBe(true);
+
+    const paths = result.agents.map((a: any) => a.path);
+    expect(paths).toContain('/agents/public-bot');
+    expect(paths).not.toContain('/agents/@clock');
+  });
+
+  test('list_agents with system token → all agents visible', async () => {
+    const token = await signToken();
+    const rpc = await mcpCall(SDK_PORT, 'list_agents', {}, token);
+    const result = parseResult(rpc);
+    expect(result.success).toBe(true);
+
+    const paths = result.agents.map((a: any) => a.path);
+    expect(paths).toContain('/agents/public-bot');
+    expect(paths).toContain('/agents/@clock');
+  });
+
+  // ─── Scopes: limited issuer ────────────────────────────────
+
+  test('issuer with limited scopes → resolves as agent, not system', async () => {
+    // Create a separate server with limited-scope issuer
+    const limitedRegistry = createAgentRegistry();
+    limitedRegistry.register(makeAgent('/agents/@private-agent', { visibility: 'private' }));
+    limitedRegistry.register(makeAgent('/agents/@internal-agent', { visibility: 'internal' }));
+
+    const limitedServer = createAgentServer(limitedRegistry, {
+      port: 19882,
+      trustedIssuers: [{
+        issuer: ISSUER_URL,
+        scopes: ['agents:read'], // NOT agents:admin or *
+      }],
+    });
+    await limitedServer.start();
+
+    try {
+      const token = await signToken();
+
+      // agents:read grants agent-level access (not system)
+      // Internal agents are accessible to authenticated agents
+      const internalRpc = await mcpCall(19882, 'call_agent', {
+        request: { action: 'load', path: '/agents/@internal-agent' },
+      }, token);
+      expect(parseResult(internalRpc).success).toBe(true);
+
+      // Private agents should be denied (only self can access)
+      const privateRpc = await mcpCall(19882, 'call_agent', {
+        request: { action: 'load', path: '/agents/@private-agent' },
+      }, token);
+      expect(parseResult(privateRpc).success).toBe(false);
+      expect(parseResult(privateRpc).code).toBe('ACCESS_DENIED');
+    } finally {
+      limitedServer?.stop?.();
+    }
+  });
+});
+
+// ─── Unit: detectAuth ────────────────────────────────────────────
+
+describe('detectAuth', () => {
+  test('returns non-null even without @auth agent', () => {
+    const registry = createAgentRegistry();
+    const config = detectAuth(registry);
+    expect(config).toBeDefined();
+    expect(config).not.toBeNull();
+  });
+
+  test('returns empty config (no store, no rootKey) without @auth agent', () => {
+    const registry = createAgentRegistry();
+    const config = detectAuth(registry);
+    expect(config.store).toBeUndefined();
+    expect(config.rootKey).toBeUndefined();
+  });
+});
+
+// ─── Unit: canSeeAgent ───────────────────────────────────────────
+
+describe('canSeeAgent', () => {
+  test('system auth can see internal agents', () => {
+    const agent = makeAgent('/agents/@clock', { visibility: 'internal' });
+    const auth = { callerId: 'api', callerType: 'system' as const, scopes: ['*'], isRoot: true };
+    expect(canSeeAgent(agent, auth)).toBe(true);
+  });
+
+  test('null auth cannot see internal agents', () => {
+    const agent = makeAgent('/agents/@clock', { visibility: 'internal' });
+    expect(canSeeAgent(agent, null)).toBe(false);
+  });
+
+  test('null auth can see public agents', () => {
+    const agent = makeAgent('/agents/public', { visibility: 'public' });
+    expect(canSeeAgent(agent, null)).toBe(true);
+  });
+});