@kya-os/mcp-i-core 1.2.3-canary.7 → 1.3.0

package/src/__tests__/runtime/base.test.ts

@@ -0,0 +1,869 @@
+/**
+ * Tests for MCPIRuntimeBase
+ */
+
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+import { MCPIRuntimeBase } from "../../runtime/base";
+import { ProviderRuntimeConfig } from "../../config";
+import {
+  createMockProviders,
+  MockClockProvider,
+  MockFetchProvider,
+  MockIdentityProvider,
+  MockNonceCacheProvider,
+} from "../utils/mock-providers";
+
+describe("MCPIRuntimeBase", () => {
+  let runtime: MCPIRuntimeBase;
+  let config: ProviderRuntimeConfig;
+  let mockProviders: ReturnType<typeof createMockProviders>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockProviders = createMockProviders();
+    config = {
+      ...mockProviders,
+      environment: "development",
+      session: {
+        timestampSkewSeconds: 120,
+        ttlMinutes: 30,
+      },
+      audit: {
+        enabled: true,
+        logFunction: vi.fn(),
+        includePayloads: true,
+      },
+      wellKnown: {
+        enabled: true,
+        serviceName: "Test Service",
+        serviceEndpoint: "https://test.example.com",
+      },
+      showVerifyLink: true,
+      identityBadge: true,
+    };
+    runtime = new MCPIRuntimeBase(config);
+  });
+
+  describe("constructor", () => {
+    it("should initialize with providers", () => {
+      expect(runtime).toBeDefined();
+      expect((runtime as any).crypto).toBe(mockProviders.cryptoProvider);
+      expect((runtime as any).clock).toBe(mockProviders.clockProvider);
+      expect((runtime as any).fetch).toBe(mockProviders.fetchProvider);
+      expect((runtime as any).storage).toBe(mockProviders.storageProvider);
+      expect((runtime as any).nonceCache).toBe(
+        mockProviders.nonceCacheProvider
+      );
+      expect((runtime as any).identity).toBe(mockProviders.identityProvider);
+    });
+
+    it("should store config", () => {
+      expect((runtime as any).config).toEqual(config);
+    });
+  });
+
+  describe("initialize", () => {
+    it("should load identity on initialization", async () => {
+      await runtime.initialize();
+      const identity = await runtime.getIdentity();
+      expect(identity.did).toBe("did:key:zmock123");
+    });
+
+    it("should call nonce cache initialize if available", async () => {
+      const initializeFn = vi.fn();
+      (mockProviders.nonceCacheProvider as any).initialize = initializeFn;
+
+      await runtime.initialize();
+      expect(initializeFn).toHaveBeenCalled();
+    });
+
+    it("should log audit event on initialization", async () => {
+      await runtime.initialize();
+
+      expect(config.audit!.logFunction).toHaveBeenCalledWith(
+        expect.stringContaining("runtime_initialized")
+      );
+
+      const logCall = (config.audit!.logFunction as any).mock.calls[0][0];
+      const logData = JSON.parse(logCall);
+      expect(logData.event).toBe("runtime_initialized");
+      expect(logData.data.did).toBe("did:key:zmock123");
+      expect(logData.data.environment).toBe("development");
+    });
+
+    it("should not log audit if disabled", async () => {
+      config.audit!.enabled = false;
+      await runtime.initialize();
+      expect(config.audit!.logFunction).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("getIdentity", () => {
+    it("should return cached identity if available", async () => {
+      await runtime.initialize();
+      const identity1 = await runtime.getIdentity();
+      const identity2 = await runtime.getIdentity();
+
+      expect(identity1).toBe(identity2);
+    });
+
+    it("should load identity if not cached", async () => {
+      const identity = await runtime.getIdentity();
+      expect(identity.did).toBe("did:key:zmock123");
+    });
+  });
+
+  describe("handleHandshake", () => {
+    const handshakeRequest = {
+      nonce: "test-nonce-123",
+      audience: "https://client.example.com",
+      timestamp: Date.now(),
+      clientDid: "did:key:zclient123",
+    };
+
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should create session and return handshake response", async () => {
+      const response = await runtime.handleHandshake(handshakeRequest);
+
+      expect(response.sessionId).toBeDefined();
+      expect(response.agentDid).toBe("did:key:zmock123");
+      expect(response.timestamp).toBe(
+        (mockProviders.clockProvider as MockClockProvider).now()
+      );
+      expect(response.capabilities).toEqual(["identity", "proof", "audit"]);
+      expect(response.signature).toBeDefined();
+    });
+
+    it("should store session internally", async () => {
+      const response = await runtime.handleHandshake(handshakeRequest);
+      const session = await runtime.getCurrentSession();
+
+      expect(session).toBeDefined();
+      expect(session.id).toBe(response.sessionId);
+      expect(session.clientDid).toBe(handshakeRequest.clientDid);
+      expect(session.agentDid).toBeUndefined(); // ✅ FIXED: agentDid should be undefined when not provided (no fallback)
+      expect(session.serverDid).toBe("did:key:zmock123"); // ✅ NEW: serverDid should be set to identity.did
+      expect(session.audience).toBe(handshakeRequest.audience);
+    });
+
+    it("should set session expiry", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      await runtime.handleHandshake(handshakeRequest);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.expiresAt).toBe(currentTime + 30 * 60 * 1000); // 30 minutes
+    });
+
+    it("should sign handshake response", async () => {
+      const response = await runtime.handleHandshake(handshakeRequest);
+      expect(response.signature).toBe(
+        Buffer.from(new Uint8Array([1, 2, 3, 4, 5])).toString("base64")
+      );
+    });
+
+    it("should use agentDid when provided in handshake request", async () => {
+      const handshakeWithAgentDid = {
+        ...handshakeRequest,
+        agentDid: "did:key:zagent456",
+      };
+
+      const response = await runtime.handleHandshake(handshakeWithAgentDid);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.agentDid).toBe("did:key:zagent456"); // ✅ Should use agentDid, not clientDid
+      expect(session.clientDid).toBe(handshakeRequest.clientDid); // clientDid should remain unchanged
+      expect(session.serverDid).toBe("did:key:zmock123"); // ✅ serverDid should be set to identity.did
+    });
+
+    it("should extract and store clientInfo when provided", async () => {
+      const handshakeWithClientInfo = {
+        ...handshakeRequest,
+        clientInfo: {
+          name: "Claude Desktop",
+          title: "Claude Desktop Preview",
+          version: "1.0.0",
+          platform: "desktop",
+          vendor: "Anthropic",
+          persistentId: "persistent-client-id",
+        },
+        clientProtocolVersion: "2025-06-18",
+        clientCapabilities: {
+          tools: { listChanged: true },
+        },
+      };
+
+      await runtime.handleHandshake(handshakeWithClientInfo);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.clientInfo).toBeDefined();
+      expect(session.clientInfo?.name).toBe("Claude Desktop");
+      expect(session.clientInfo?.title).toBe("Claude Desktop Preview");
+      expect(session.clientInfo?.version).toBe("1.0.0");
+      expect(session.clientInfo?.platform).toBe("desktop");
+      expect(session.clientInfo?.clientId).toBeDefined();
+      expect(typeof session.clientInfo?.clientId).toBe("string");
+      expect(session.clientInfo?.vendor).toBe("Anthropic");
+      expect(session.clientInfo?.persistentId).toBe("persistent-client-id");
+      expect(session.clientInfo?.protocolVersion).toBe("2025-06-18");
+      expect(session.clientInfo?.capabilities).toEqual({
+        tools: { listChanged: true },
+      });
+    });
+
+    it("should handle handshake without clientInfo gracefully", async () => {
+      await runtime.handleHandshake(handshakeRequest);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.clientInfo).toBeUndefined();
+    });
+
+    it("should preserve clientId when provided by adapter", async () => {
+      const handshakeWithClientId = {
+        ...handshakeRequest,
+        clientInfo: {
+          name: "Claude Desktop",
+          clientId: "server-specified-client-id",
+        },
+      };
+
+      await runtime.handleHandshake(handshakeWithClientId);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.clientInfo?.clientId).toBe("server-specified-client-id");
+    });
+
+    it("should default clientInfo name to 'unknown' when name is missing", async () => {
+      const handshakeWithPartialClientInfo = {
+        ...handshakeRequest,
+        clientInfo: {
+          version: "1.0.0",
+          platform: "desktop",
+        },
+      };
+
+      await runtime.handleHandshake(handshakeWithPartialClientInfo);
+      const session = await runtime.getCurrentSession();
+
+      expect(session.clientInfo).toBeDefined();
+      expect(session.clientInfo?.name).toBe("unknown");
+      expect(session.clientInfo?.version).toBe("1.0.0");
+      expect(session.clientInfo?.platform).toBe("desktop");
+    });
+  });
+
+  describe("processToolCall", () => {
+    const mockHandler = vi.fn().mockResolvedValue({ result: "success" });
+    const session = {
+      id: "session123",
+      audience: "https://client.example.com",
+      nonce: "test-nonce",
+    };
+
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should execute tool handler and return clean result", async () => {
+      const result = await runtime.processToolCall(
+        "testTool",
+        { arg: "value" },
+        mockHandler,
+        session
+      );
+
+      expect(mockHandler).toHaveBeenCalledWith({ arg: "value" });
+      expect(result).toEqual({ result: "success" });
+    });
+
+    it("should create and store proof", async () => {
+      await runtime.processToolCall(
+        "testTool",
+        { arg: "value" },
+        mockHandler,
+        session
+      );
+
+      const proof = runtime.getLastProof();
+      expect(proof).toBeDefined();
+      expect(proof.did).toBe("did:key:zmock123");
+      expect(proof.signature).toBeDefined();
+      expect(proof.nonce).toBe("test-nonce");
+      expect(proof.sessionId).toBe("session123");
+      expect(proof.audience).toBe("https://client.example.com");
+    });
+
+    it("should log audit event", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      await runtime.processToolCall(
+        "testTool",
+        { arg: "value" },
+        mockHandler,
+        session
+      );
+
+      expect(config.audit!.logFunction).toHaveBeenCalled();
+
+      // Find the tool_executed log entry (might not be the first one)
+      const logCalls = (config.audit!.logFunction as any).mock.calls;
+      const toolLogCall = logCalls.find((call: any[]) => {
+        const log = JSON.parse(call[0]);
+        return log.event === "tool_executed";
+      });
+
+      expect(toolLogCall).toBeDefined();
+      const logData = JSON.parse(toolLogCall[0]);
+      expect(logData.event).toBe("tool_executed");
+      expect(logData.data.tool).toBe("testTool");
+      expect(logData.data.sessionId).toBe("session123");
+      expect(logData.data.timestamp).toBe(currentTime);
+    });
+
+    it("should handle tool handler errors", async () => {
+      const errorHandler = vi.fn().mockRejectedValue(new Error("Tool failed"));
+
+      await expect(
+        runtime.processToolCall("errorTool", {}, errorHandler)
+      ).rejects.toThrow("Tool failed");
+    });
+  });
+
+  describe("issueNonce", () => {
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should generate and cache nonce", async () => {
+      const nonce = await runtime.issueNonce("session123");
+
+      expect(nonce).toBeDefined();
+      expect(nonce).toHaveLength(44); // Base64 encoded 32 bytes
+
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      const identity = await runtime.getIdentity();
+      expect(await nonceCache.has(nonce, identity.did)).toBe(true);
+    });
+
+    it("should set 5 minute expiry", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      const nonce = await runtime.issueNonce("session123");
+
+      // Check that nonce expires after 5 minutes
+      clock.setTime(currentTime + 5 * 60 * 1000 + 1);
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      const identity = await runtime.getIdentity();
+      expect(await nonceCache.has(nonce, identity.did)).toBe(false);
+    });
+  });
+
+  describe("createProof", () => {
+    const testData = { foo: "bar" };
+
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should create proof with all fields", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      const proof = await runtime.createProof(testData);
+
+      expect(proof.timestamp).toBe(currentTime);
+      expect(proof.nonce).toBeDefined();
+      expect(proof.did).toBe("did:key:zmock123");
+      expect(proof.signature).toBeDefined();
+      expect(proof.algorithm).toBe("Ed25519");
+      expect(proof.sessionId).toBeUndefined();
+      expect(proof.audience).toBeUndefined();
+    });
+
+    it("should use session nonce if provided", async () => {
+      const session = {
+        id: "session123",
+        nonce: "session-nonce",
+        audience: "https://client.example.com",
+      };
+
+      const proof = await runtime.createProof(testData, session);
+
+      expect(proof.nonce).toBe("session-nonce");
+      expect(proof.sessionId).toBe("session123");
+      expect(proof.audience).toBe("https://client.example.com");
+    });
+
+    it("should generate and cache new nonce if not in session", async () => {
+      const proof = await runtime.createProof(testData);
+
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      const identity = await runtime.getIdentity();
+      expect(await nonceCache.has(proof.nonce, identity.did)).toBe(true);
+    });
+  });
+
+  describe("verifyProof", () => {
+    const testData = { foo: "bar" };
+    let validProof: any;
+
+    beforeEach(async () => {
+      await runtime.initialize();
+
+      // Set up DID document for verification
+      const fetchProvider = mockProviders.fetchProvider as MockFetchProvider;
+      fetchProvider.setDIDDocument("did:key:zmock123", {
+        verificationMethod: [
+          {
+            publicKeyBase64: "mock-public-key",
+          },
+        ],
+      });
+
+      // Create a valid proof
+      validProof = await runtime.createProof(testData);
+    });
+
+    it("should verify valid proof", async () => {
+      // Remove the nonce from cache first (simulating a fresh verification)
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      nonceCache.clear();
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(true);
+    });
+
+    it("should reject if nonce already used", async () => {
+      // Pre-add nonce to cache with agent-scoped key
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      await nonceCache.add(
+        validProof.nonce,
+        Date.now() + 10000,
+        validProof.did
+      );
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(false);
+    });
+
+    it("should use agent-scoped nonces in legacy verification (prevents cross-agent replay)", async () => {
+      const agentDid1 = "did:key:agent1";
+      const agentDid2 = "did:key:agent2";
+      const sameNonce = "same-nonce-value";
+
+      // Set up DID documents for both agents
+      const fetchProvider = mockProviders.fetchProvider as MockFetchProvider;
+      fetchProvider.setDIDDocument(agentDid1, {
+        verificationMethod: [
+          {
+            publicKeyBase64: "mock-public-key-agent1",
+          },
+        ],
+      });
+      fetchProvider.setDIDDocument(agentDid2, {
+        verificationMethod: [
+          {
+            publicKeyBase64: "mock-public-key-agent2",
+          },
+        ],
+      });
+
+      // Create proof for agent 1
+      const proof1 = {
+        ...validProof,
+        did: agentDid1,
+        nonce: sameNonce,
+      };
+
+      // Create proof for agent 2 with same nonce
+      const proof2 = {
+        ...validProof,
+        did: agentDid2,
+        nonce: sameNonce,
+      };
+
+      const nonceCache =
+        mockProviders.nonceCacheProvider as MockNonceCacheProvider;
+      nonceCache.clear();
+
+      // Verify proof for agent 1 (should succeed)
+      const isValid1 = await runtime.verifyProof(testData, proof1);
+      expect(isValid1).toBe(true);
+
+      // Verify that nonce was added with agent-scoped key
+      // The nonce should be scoped to agentDid1, so agentDid2 should be able to use the same nonce
+      const hasNonceAgent1 = await nonceCache.has(sameNonce, agentDid1);
+      expect(hasNonceAgent1).toBe(true);
+
+      // Verify proof for agent 2 with same nonce (should succeed because nonces are agent-scoped)
+      const isValid2 = await runtime.verifyProof(testData, proof2);
+      expect(isValid2).toBe(true);
+
+      // Verify that both agents' nonces are cached separately
+      const hasNonceAgent2 = await nonceCache.has(sameNonce, agentDid2);
+      expect(hasNonceAgent2).toBe(true);
+    });
+
+    it("should reject if timestamp outside skew", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime + 121 * 1000); // 121 seconds later (skew is 120)
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(false);
+    });
+
+    it("should reject if DID cannot be resolved", async () => {
+      const fetchProvider = mockProviders.fetchProvider as MockFetchProvider;
+      fetchProvider.setDIDDocument("did:key:zmock123", null as any);
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(false);
+    });
+
+    it("should reject if signature invalid", async () => {
+      validProof.signature = "invalid-signature";
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(false);
+    });
+
+    it("should handle verification errors gracefully", async () => {
+      const fetchProvider = mockProviders.fetchProvider as MockFetchProvider;
+      (fetchProvider.resolveDID as any) = vi
+        .fn()
+        .mockRejectedValue(new Error("Network error"));
+
+      const isValid = await runtime.verifyProof(testData, validProof);
+      expect(isValid).toBe(false);
+    });
+  });
+
+  describe("getCurrentSession", () => {
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should return active session", async () => {
+      await runtime.handleHandshake({
+        clientDid: "did:key:zclient123",
+        audience: "https://client.example.com",
+      });
+
+      const session = await runtime.getCurrentSession();
+      expect(session).toBeDefined();
+      expect(session.clientDid).toBe("did:key:zclient123");
+    });
+
+    it("should return null if no sessions", async () => {
+      const session = await runtime.getCurrentSession();
+      expect(session).toBeNull();
+    });
+
+    it("should return null if session expired", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      await runtime.handleHandshake({
+        clientDid: "did:key:zclient123",
+        audience: "https://client.example.com",
+      });
+
+      // Advance time past session expiry
+      clock.setTime(currentTime + 31 * 60 * 1000); // 31 minutes
+
+      const session = await runtime.getCurrentSession();
+      expect(session).toBeNull();
+    });
+  });
+
+  describe("getLastProof", () => {
+    it("should return undefined initially", () => {
+      const proof = runtime.getLastProof();
+      expect(proof).toBeUndefined();
+    });
+
+    it("should return last generated proof", async () => {
+      await runtime.initialize();
+      await runtime.processToolCall(
+        "testTool",
+        {},
+        vi.fn().mockResolvedValue({ result: "ok" })
+      );
+
+      const proof = runtime.getLastProof();
+      expect(proof).toBeDefined();
+      expect(proof.did).toBe("did:key:zmock123");
+    });
+  });
+
+  describe("createWellKnownHandler", () => {
+    let handler: any;
+
+    beforeEach(async () => {
+      await runtime.initialize();
+      handler = runtime.createWellKnownHandler({
+        serviceName: "Test Service",
+        serviceEndpoint: "https://test.example.com",
+      });
+    });
+
+    it("should handle /.well-known/did.json", async () => {
+      const response = await handler("/.well-known/did.json");
+
+      expect(response.status).toBe(200);
+      expect(response.headers["Content-Type"]).toBe("application/did+json");
+
+      const result = JSON.parse(response.body);
+      expect(result).toEqual({
+        "@context": ["https://www.w3.org/ns/did/v1"],
+        id: "did:key:zmock123",
+        verificationMethod: [
+          {
+            id: "did:key:zmock123#key-1",
+            type: "Ed25519VerificationKey2020",
+            controller: "did:key:zmock123",
+            publicKeyBase64: "mock-public-key",
+          },
+        ],
+        authentication: ["did:key:zmock123#key-1"],
+        assertionMethod: ["did:key:zmock123#key-1"],
+      });
+    });
+
+    it("should handle /.well-known/mcp-identity", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      const result = await handler("/.well-known/mcp-identity");
+
+      expect(result).toEqual({
+        did: "did:key:zmock123",
+        publicKey: "mock-public-key",
+        serviceName: "Test Service",
+        serviceEndpoint: "https://test.example.com",
+        timestamp: currentTime,
+      });
+    });
+
+    it("should return null for unknown paths", async () => {
+      const result = await handler("/unknown/path");
+      expect(result).toBeNull();
+    });
+
+    it("should use default values if config not provided", async () => {
+      const defaultHandler = runtime.createWellKnownHandler();
+      const result = await defaultHandler("/.well-known/mcp-identity");
+
+      expect(result.serviceName).toBe("MCP-I Service");
+      expect(result.serviceEndpoint).toBe("https://example.com");
+    });
+  });
+
+  describe("createDebugEndpoint", () => {
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should return debug info in development", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      const endpoint = runtime.createDebugEndpoint();
+      expect(endpoint).toBeDefined();
+
+      const result = await endpoint();
+      expect(result).toEqual({
+        identity: {
+          did: "did:key:zmock123",
+          publicKey: "mock-public-key",
+        },
+        session: null,
+        config: {
+          environment: "development",
+          timestampSkewSeconds: 120,
+          sessionTtlMinutes: 30,
+        },
+        timestamp: currentTime,
+      });
+    });
+
+    it("should include session if active", async () => {
+      await runtime.handleHandshake({
+        clientDid: "did:key:zclient123",
+        audience: "https://client.example.com",
+      });
+
+      const endpoint = runtime.createDebugEndpoint();
+      const result = await endpoint();
+
+      expect(result.session).toBeDefined();
+      expect(result.session.clientDid).toBe("did:key:zclient123");
+    });
+
+    it("should return null in production", () => {
+      config.environment = "production";
+      const prodRuntime = new MCPIRuntimeBase(config);
+
+      const endpoint = prodRuntime.createDebugEndpoint();
+      expect(endpoint).toBeNull();
+    });
+  });
+
+  describe("getAuditLogger", () => {
+    it("should return audit logger", () => {
+      const logger = runtime.getAuditLogger();
+      expect(logger).toBeDefined();
+      expect(logger.log).toBeDefined();
+    });
+
+    it("should log through runtime audit", async () => {
+      await runtime.initialize();
+      const logger = runtime.getAuditLogger();
+
+      logger.log("test_event", { test: "data" });
+
+      expect(config.audit!.logFunction).toHaveBeenCalledWith(
+        expect.stringContaining("test_event")
+      );
+    });
+  });
+
+  describe("rotateKeys", () => {
+    beforeEach(async () => {
+      await runtime.initialize();
+    });
+
+    it("should rotate keys and update identity", async () => {
+      const oldIdentity = await runtime.getIdentity();
+      const newIdentity = await runtime.rotateKeys();
+
+      expect(newIdentity.did).toBe("did:key:zmock456-1");
+      expect(newIdentity.privateKey).toBe("mock-private-key-rotated-1");
+      expect(newIdentity.publicKey).toBe("mock-public-key-rotated-1");
+
+      const currentIdentity = await runtime.getIdentity();
+      expect(currentIdentity).toEqual(newIdentity);
+    });
+
+    it("should log audit event", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      // Get the old identity first
+      const oldIdentity = await runtime.getIdentity();
+
+      await runtime.rotateKeys();
+
+      expect(config.audit!.logFunction).toHaveBeenCalled();
+
+      // Find the keys_rotated log entry
+      const logCalls = (config.audit!.logFunction as any).mock.calls;
+      const rotateLogCall = logCalls.find((call: any[]) => {
+        const log = JSON.parse(call[0]);
+        return log.event === "keys_rotated";
+      });
+
+      expect(rotateLogCall).toBeDefined();
+      const logData = JSON.parse(rotateLogCall[0]);
+      expect(logData.event).toBe("keys_rotated");
+      expect(logData.data.oldDid).toBe(oldIdentity.did);
+      expect(logData.data.newDid).toBe("did:key:zmock456-1");
+      expect(logData.data.timestamp).toBe(currentTime);
+    });
+  });
+
+  describe("audit logging", () => {
+    it("should use custom log function if provided", async () => {
+      await runtime.initialize();
+      expect(config.audit!.logFunction).toHaveBeenCalled();
+    });
+
+    it("should console.log if no custom function", async () => {
+      const consoleLogSpy = vi
+        .spyOn(console, "log")
+        .mockImplementation(() => {});
+
+      config.audit!.logFunction = undefined;
+      const runtimeNoLogFn = new MCPIRuntimeBase(config);
+      await runtimeNoLogFn.initialize();
+
+      expect(consoleLogSpy).toHaveBeenCalledWith("[AUDIT]", expect.any(String));
+      consoleLogSpy.mockRestore();
+    });
+
+    it("should exclude payloads if configured", async () => {
+      config.audit!.includePayloads = false;
+      const runtimeNoPayloads = new MCPIRuntimeBase(config);
+      await runtimeNoPayloads.initialize();
+
+      const logCall = (config.audit!.logFunction as any).mock.calls[0][0];
+      const logData = JSON.parse(logCall);
+      expect(logData.data).toBeUndefined();
+    });
+
+    it("should include formatted timestamp", async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      await runtime.initialize();
+
+      const logCall = (config.audit!.logFunction as any).mock.calls[0][0];
+      const logData = JSON.parse(logCall);
+      expect(logData.timestampFormatted).toBe(
+        new Date(currentTime).toISOString()
+      );
+    });
+  });
+
+  describe("helper methods", () => {
+    it("should handle DID document with publicKeyMultibase", async () => {
+      await runtime.initialize();
+
+      const fetchProvider = mockProviders.fetchProvider as MockFetchProvider;
+      fetchProvider.setDIDDocument("did:key:ztest", {
+        verificationMethod: [
+          {
+            publicKeyMultibase: "zBase58EncodedKey",
+          },
+        ],
+      });
+
+      // This would be called internally during verification
+      const didDoc = await fetchProvider.resolveDID("did:key:ztest");
+      const extractPublicKey = (runtime as any).extractPublicKey.bind(runtime);
+      const publicKey = extractPublicKey(didDoc);
+
+      expect(publicKey).toBe("zBase58EncodedKey");
+    });
+
+    it("should throw if no public key in DID document", async () => {
+      await runtime.initialize();
+
+      const didDoc = { verificationMethod: [{}] };
+      const extractPublicKey = (runtime as any).extractPublicKey.bind(runtime);
+
+      expect(() => extractPublicKey(didDoc)).toThrow("Public key not found");
+    });
+  });
+});