@slashfi/agents-sdk 0.12.0 → 0.13.0

package/src/agent-definitions/auth.ts

@@ -25,7 +25,7 @@
  */
 
 import { defineAgent, defineTool } from "../define.js";
-import { signJwt, generateSigningKey, exportSigningKey, type ExportedKeyPair } from "../jwt.js";
+import { signJwt, generateSigningKey, exportSigningKey, verifyJwtFromIssuer, type ExportedKeyPair } from "../jwt.js";
 import type { AgentDefinition, ToolContext, ToolDefinition } from "../types.js";
 
 // ============================================
@@ -173,6 +173,30 @@ export interface AuthStore {
     userId: string;
     clientId: string;
   } | null>;
+
+  // --- Tenant Identity ---
+
+  /** Store a tenant identity mapping (foreign issuer + ID -> local tenant). */
+  storeTenantIdentity?(tenantId: string, provider: string, providerTenantId: string): Promise<void>;
+
+  /** Resolve a local tenant ID from a foreign identity. */
+  resolveTenantByIdentity?(provider: string, providerTenantId: string): Promise<string | null>;
+
+  // --- User Identity ---
+
+  /** Store a user identity mapping (foreign issuer + ID -> local user). */
+  storeUserIdentity?(userId: string, provider: string, providerUserId: string): Promise<void>;
+
+  /** Resolve a local user ID from a foreign identity. */
+  resolveUserByIdentity?(provider: string, providerUserId: string): Promise<string | null>;
+
+  // --- Transaction ---
+
+  /**
+   * Run operations atomically. The store scopes its own methods to the
+   * transaction context. For stores without real tx support, run fn sequentially.
+   */
+  transaction<T>(fn: () => Promise<T>): Promise<T>;
 }
 
 // ============================================
@@ -218,6 +242,8 @@ export function createMemoryAuthStore(): AuthStore {
   const tokens = new Map<string, AuthToken>();
   const signingKeys = new Map<string, ExportedKeyPair>();
   const trustedIssuers = new Set<string>();
+  const tenantIdentities = new Map<string, string>(); // "provider:providerTenantId" -> tenantId
+  const userIdentities = new Map<string, string>(); // "provider:providerUserId" -> userId
 
   return {
     async createTenant(name, _externalRef) {
@@ -346,6 +372,27 @@ export function createMemoryAuthStore(): AuthStore {
     async listTrustedIssuers() {
       return Array.from(trustedIssuers);
     },
+
+    async storeTenantIdentity(tenantId, provider, providerTenantId) {
+      tenantIdentities.set(`${provider}:${providerTenantId}`, tenantId);
+    },
+
+    async resolveTenantByIdentity(provider, providerTenantId) {
+      return tenantIdentities.get(`${provider}:${providerTenantId}`) ?? null;
+    },
+
+    async storeUserIdentity(userId, provider, providerUserId) {
+      userIdentities.set(`${provider}:${providerUserId}`, userId);
+    },
+
+    async resolveUserByIdentity(provider, providerUserId) {
+      return userIdentities.get(`${provider}:${providerUserId}`) ?? null;
+    },
+
+    async transaction<T>(fn: () => Promise<T>): Promise<T> {
+      // In-memory: just run sequentially (single-threaded, no concurrency issues)
+      return fn();
+    },
   };
 }
 
@@ -783,28 +830,88 @@ export function createAuthAgent(
           type: "string" as const,
           description: "JWT signed by a trusted issuer",
         },
-        connectBaseUrl: {
-          type: "string" as const,
-          description: "Base URL for the OAuth connect flow (returned in needsAuth response)",
-        },
+
       },
       required: ["token"],
     },
     execute: async (
-      _input: { token: string; connectBaseUrl?: string },
+      input: { token: string },
     ) => {
-      // This tool is a stub — the actual implementation needs:
-      // 1. JWT verification (via verifyJwtFromIssuer)
-      // 2. Tenant resolution (via tenant_identity table)
-      // 3. User resolution (via user_identity table)
-      // These depend on the store having identity lookup methods.
-      //
-      // For now, return the structure so the flow can be wired.
-      // The atlas-environments CockroachDB implementation overrides this.
-      return {
-        error: "exchange_token requires a store with identity resolution support",
-        hint: "Override this tool in your environment implementation",
-      };
+      if (!store.resolveTenantByIdentity || !store.resolveUserByIdentity) {
+        return {
+          error: "exchange_token requires a store with identity resolution support",
+          hint: "Implement storeTenantIdentity/resolveUserByIdentity on your AuthStore",
+        };
+      }
+
+      // 1. Decode JWT to read iss claim (no verification yet)
+      const parts = input.token.split(".");
+      if (parts.length !== 3) {
+        return { success: false, error: "Invalid JWT format" };
+      }
+      let decoded: any;
+      try {
+        decoded = JSON.parse(Buffer.from(parts[1], "base64url").toString());
+      } catch {
+        return { success: false, error: "Failed to decode JWT payload" };
+      }
+
+      const issuer = decoded.iss;
+      const sub = decoded.sub;
+      const foreignTenantId = decoded.tenantId;
+      if (!issuer || !sub) {
+        return { success: false, error: "JWT missing iss or sub claims" };
+      }
+
+      // 2. Match issuer against trusted issuers
+      const trustedIssuers = store.listTrustedIssuers ? await store.listTrustedIssuers() : [];
+      if (!trustedIssuers.includes(issuer)) {
+        return { success: false, error: `Issuer ${issuer} is not trusted` };
+      }
+
+      // 3. Verify JWT against the matched issuer's JWKS
+      let payload: any;
+      try {
+        payload = await verifyJwtFromIssuer(input.token, issuer);
+      } catch {
+        return { success: false, error: "JWT verification failed" };
+      }
+      if (!payload) {
+        return { success: false, error: "JWT verification returned empty payload" };
+      }
+
+      // 4. Resolve tenant + user inside a transaction for consistency
+      return store.transaction(async () => {
+        const localTenantId = await (async () => {
+          if (!foreignTenantId) return null;
+          const existing = await store.resolveTenantByIdentity!(issuer, foreignTenantId);
+          if (existing) return existing;
+          // Auto-create tenant identity link on first encounter
+          if (store.storeTenantIdentity) {
+            await store.storeTenantIdentity(foreignTenantId, issuer, foreignTenantId);
+          }
+          return foreignTenantId;
+        })();
+
+        // 5. Resolve user
+        const localUserId = await store.resolveUserByIdentity!(issuer, sub);
+        if (localUserId) {
+          return {
+            success: true,
+            tenantId: localTenantId ?? undefined,
+            userId: localUserId,
+          };
+        }
+
+        // 4. User not linked — caller decides how to handle (e.g. OIDC flow)
+        return {
+          success: false,
+          needsAuth: true,
+          tenantId: localTenantId ?? undefined,
+          issuer,
+          sub,
+        };
+      });
     },
   });
 

package/src/agent-definitions/integrations.ts

@@ -342,7 +342,10 @@ export async function exchangeCodeForToken(
     throw new Error(`Token exchange failed (${response.status}): ${text}`);
   }
 
-  const data = (await response.json()) as Record<string, unknown>;
+  const responseText = await response.text();
+  console.log("[token-exchange] Slack response:", responseText.substring(0, 500));
+  let data: Record<string, unknown>;
+  try { data = JSON.parse(responseText); } catch (e) { throw new Error(`Failed to parse JSON: ${responseText.substring(0, 200)}`); }
 
   return {
     accessToken: String(data[oauth.accessTokenField ?? "access_token"] ?? ""),
@@ -408,7 +411,10 @@ export async function refreshAccessToken(
     throw new Error(`Token refresh failed (${response.status}): ${text}`);
   }
 
-  const data = (await response.json()) as Record<string, unknown>;
+  const responseText = await response.text();
+  console.log("[token-exchange] Slack response:", responseText.substring(0, 500));
+  let data: Record<string, unknown>;
+  try { data = JSON.parse(responseText); } catch (e) { throw new Error(`Failed to parse JSON: ${responseText.substring(0, 200)}`); }
 
   return {
     accessToken: String(data[oauth.accessTokenField ?? "access_token"] ?? ""),

package/src/agent-definitions/remote-registry.ts

@@ -177,6 +177,7 @@ export function createRemoteRegistryAgent(
 
   // Extract setup/connect as standalone functions to avoid circular reference
   const setupFn = async (params: Record<string, unknown>, _ctx: IntegrationMethodContext): Promise<IntegrationMethodResult> => {
+    console.log("[remote-registry] setupFn called with:", JSON.stringify(params));
     const url = params.url as string;
     const name = (params.name as string) ?? "registry";
     const oidcUserId = params.oidcUserId as string | undefined;
@@ -189,7 +190,7 @@ export function createRemoteRegistryAgent(
         const jwt = await signJwt({ sub: oidcUserId, action: "setup", type: "agent-registry" });
         const tokenRes = await globalThis.fetch(baseUrl + "/oauth/token", {
           method: "POST", headers: { "Content-Type": "application/json" },
-          body: JSON.stringify({ grant_type: "jwt_exchange", assertion: jwt, scope: "setup" }),
+          body: JSON.stringify({ grant_type: "jwt_exchange", assertion: jwt, scope: "setup", redirect_uri: params.redirect_uri ?? "" }),
         });
         const tokenData = await tokenRes.json() as any;
         if (!tokenData.access_token && !tokenData.tenant_id) {
@@ -203,22 +204,30 @@ export function createRemoteRegistryAgent(
 
       // Phase 1: Discover JWKS, establish trust, then request OIDC
       const configUrl = baseUrl + "/.well-known/configuration";
+      console.log("[setupFn] fetching config:", configUrl);
       const configRes = await globalThis.fetch(configUrl);
+      console.log("[setupFn] config status:", configRes.status);
       if (!configRes.ok) return { success: false, error: "Failed to discover registry at " + configUrl };
       const remoteConfig = await configRes.json() as any;
       if (remoteConfig.jwks_uri) {
+        console.log("[setupFn] fetching JWKS:", remoteConfig.jwks_uri);
         const jwksRes = await globalThis.fetch(remoteConfig.jwks_uri);
+        console.log("[setupFn] JWKS status:", jwksRes.status);
         if (!jwksRes.ok) return { success: false, error: "JWKS not reachable" };
       }
-      if (addTrustedIssuer) await addTrustedIssuer(baseUrl);
+      if (addTrustedIssuer) { console.log("[setupFn] adding trusted issuer:", baseUrl); await addTrustedIssuer(baseUrl); console.log("[setupFn] added trusted issuer"); }
 
       // Request identity — atlas will return authorize URL for Slack OIDC
+      console.log("[setupFn] Phase 1: requesting identity via jwt_exchange");
       const jwt = await signJwt({ action: "setup", type: "agent-registry", targetUrl: url });
+      console.log("[setupFn] POSTing to:", baseUrl + "/oauth/token");
       const tokenRes = await globalThis.fetch(baseUrl + "/oauth/token", {
         method: "POST", headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({ grant_type: "jwt_exchange", assertion: jwt, scope: "setup" }),
+        body: JSON.stringify({ grant_type: "jwt_exchange", assertion: jwt, scope: "setup", redirect_uri: params.redirect_uri ?? "" }),
       });
+      console.log("[setupFn] token status:", tokenRes.status);
       const tokenData = await tokenRes.json() as any;
+      console.log("[setupFn] tokenData:", JSON.stringify(tokenData).substring(0, 300));
 
       // If already set up (user linked), store connection directly
       if (tokenData.access_token) {

package/src/index.ts

@@ -194,3 +194,4 @@ export type {
 export * from "./integrations-store.js";
 export * from "./integration-interface.js";
 export type { ContextFactory } from "./registry.js";
+export { createKeyManager, type KeyManager, type KeyStore, type KeyManagerOptions, type StoredKey, type KeyStatus } from "./key-manager.js";

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();
+  });
+});