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

package/src/__tests__/runtime/route-interception.test.ts

@@ -0,0 +1,686 @@
+/**
+ * Route Interception Tests
+ *
+ * Tests for route interception, resume tokens, and tool call resumption.
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { MCPIRuntimeBase } from '../../runtime/base';
+import { ProviderRuntimeConfig } from '../../config';
+import { createMockProviders, MockClockProvider } from '../utils/mock-providers';
+import { DelegationRequiredError } from '../../types/tool-protection';
+
+describe('MCPIRuntimeBase - Route Interception', () => {
+  let runtime: MCPIRuntimeBase;
+  let config: ProviderRuntimeConfig;
+  let mockProviders: ReturnType<typeof createMockProviders>;
+  let mockToolProtectionService: {
+    checkToolProtection: ReturnType<typeof vi.fn>;
+  };
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    mockProviders = createMockProviders();
+    
+    mockToolProtectionService = {
+      checkToolProtection: vi.fn()
+    };
+
+    config = {
+      ...mockProviders,
+      environment: 'development',
+      session: {
+        timestampSkewSeconds: 120,
+        ttlMinutes: 30
+      },
+      audit: {
+        enabled: false
+      },
+      toolProtectionService: mockToolProtectionService as any
+    };
+    runtime = new MCPIRuntimeBase(config);
+    await runtime.initialize();
+  });
+
+  describe('generateResumeToken', () => {
+    it('should generate unique resume tokens for different calls', async () => {
+      const call1 = {
+        toolName: 'tool1',
+        args: { arg: 'value1' },
+        sessionId: 'session1',
+        timestamp: Date.now(),
+        expiresAt: Date.now() + 1800000
+      };
+
+      const call2 = {
+        toolName: 'tool2',
+        args: { arg: 'value2' },
+        sessionId: 'session2',
+        timestamp: Date.now(),
+        expiresAt: Date.now() + 1800000
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      const token1 = generateToken(call1);
+      const token2 = generateToken(call2);
+
+      expect(token1).toBeDefined();
+      expect(token2).toBeDefined();
+      expect(token1).not.toBe(token2);
+      expect(token1).toMatch(/^resume_/);
+      expect(token2).toMatch(/^resume_/);
+    });
+
+    it('should generate different tokens due to timestamp component', async () => {
+      const call = {
+        toolName: 'tool1',
+        args: { arg: 'value' },
+        sessionId: 'session1',
+        timestamp: 1000,
+        expiresAt: 1801000
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      
+      // Generate multiple tokens to verify they differ when generated at different times
+      // Note: generateResumeToken uses Date.now() internally, so tokens will differ
+      // unless generated in the exact same millisecond (very rare)
+      const tokens = new Set<string>();
+      for (let i = 0; i < 10; i++) {
+        const token = generateToken(call);
+        tokens.add(token);
+        // Small delay to ensure different timestamps
+        await new Promise(resolve => setTimeout(resolve, 1));
+      }
+
+      // Verify token format: resume_{hash}_{timestamp}
+      const firstToken = Array.from(tokens)[0];
+      expect(firstToken).toMatch(/^resume_[a-z0-9]+_[a-z0-9]+$/);
+      
+      const parts = firstToken.split('_');
+      expect(parts.length).toBe(3);
+      expect(parts[0]).toBe('resume');
+      
+      // Verify hash part is consistent for same call context
+      const hashPart = parts[1];
+      for (const token of tokens) {
+        const tokenParts = token.split('_');
+        expect(tokenParts[1]).toBe(hashPart); // Same hash for same call
+      }
+      
+      // Most tokens should be unique (allowing for edge case of same millisecond)
+      // If all 10 tokens are unique, we're confident the timestamp component works
+      expect(tokens.size).toBeGreaterThan(1); // At least some should differ
+    });
+
+    it('should include call context in token generation', async () => {
+      const call1 = {
+        toolName: 'tool1',
+        args: { arg: 'value1' },
+        sessionId: 'session1',
+        timestamp: 1000,
+        expiresAt: 1801000
+      };
+
+      const call2 = {
+        toolName: 'tool1',
+        args: { arg: 'value2' }, // Different args
+        sessionId: 'session1',
+        timestamp: 1000,
+        expiresAt: 1801000
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      const token1 = generateToken(call1);
+      const token2 = generateToken(call2);
+
+      // Different args should produce different tokens
+      expect(token1).not.toBe(token2);
+    });
+  });
+
+  describe('buildConsentUrl', () => {
+    it('should build consent URL with tool and scopes', () => {
+      const session = {
+        id: 'session123',
+        agentDid: 'did:key:zagent123',
+        audience: 'https://client.example.com'
+      };
+
+      const buildUrl = (runtime as any).buildConsentUrl.bind(runtime);
+      const url = buildUrl(
+        'protectedTool',
+        ['files:write', 'files:read'],
+        session,
+        'resume_token_123'
+      );
+
+      expect(url).toContain('https://kya.vouched.id/bouncer/consent');
+      expect(url).toContain('tool=protectedTool');
+      expect(url).toContain('scopes=files%3Awrite%2Cfiles%3Aread');
+      expect(url).toContain('session_id=session123');
+      expect(url).toContain('agent_did=did%3Akey%3Azagent123');
+      expect(url).toContain('resume_token=resume_token_123');
+    });
+
+    it('should build consent URL without resume token if not provided', () => {
+      const session = {
+        id: 'session123',
+        agentDid: 'did:key:zagent123',
+        audience: 'https://client.example.com'
+      };
+
+      const buildUrl = (runtime as any).buildConsentUrl.bind(runtime);
+      const url = buildUrl(
+        'protectedTool',
+        ['files:write'],
+        session
+      );
+
+      expect(url).toContain('https://kya.vouched.id/bouncer/consent');
+      expect(url).not.toContain('resume_token=');
+    });
+
+    it('should handle missing session fields gracefully', () => {
+      const buildUrl = (runtime as any).buildConsentUrl.bind(runtime);
+      const url = buildUrl(
+        'protectedTool',
+        ['files:write'],
+        {},
+        'resume_token_123'
+      );
+
+      expect(url).toContain('https://kya.vouched.id/bouncer/consent');
+      expect(url).toContain('tool=protectedTool');
+      expect(url).toContain('session_id=');
+      expect(url).toContain('agent_did=');
+    });
+
+    it('should URL encode special characters in scopes', () => {
+      const session = {
+        id: 'session123',
+        agentDid: 'did:key:zagent123'
+      };
+
+      const buildUrl = (runtime as any).buildConsentUrl.bind(runtime);
+      const url = buildUrl(
+        'protectedTool',
+        ['files:write', 'scope with spaces', 'scope:with:colons'],
+        session
+      );
+
+      expect(url).toContain('scopes=files%3Awrite%2Cscope+with+spaces%2Cscope%3Awith%3Acolons');
+    });
+  });
+
+  describe('intercepted calls storage', () => {
+    const mockHandler = vi.fn().mockResolvedValue({ result: 'success' });
+    const session = {
+      id: 'session123',
+      audience: 'https://client.example.com',
+      nonce: 'test-nonce'
+    };
+
+    it('should store intercepted call when delegation required', async () => {
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          { arg: 'value' },
+          mockHandler,
+          session
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        expect(error.resumeToken).toBeDefined();
+        
+        // Check that intercepted call is stored
+        const interceptedCalls = (runtime as any).interceptedCalls;
+        const storedCall = interceptedCalls.get(error.resumeToken);
+        expect(storedCall).toBeDefined();
+        expect(storedCall.toolName).toBe('protectedTool');
+        expect(storedCall.args).toEqual({ arg: 'value' });
+        expect(storedCall.sessionId).toBe('session123');
+      }
+    });
+
+    it('should set expiry time for intercepted calls', async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          { arg: 'value' },
+          mockHandler,
+          session
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        const interceptedCalls = (runtime as any).interceptedCalls;
+        const storedCall = interceptedCalls.get(error.resumeToken);
+        expect(storedCall.expiresAt).toBe(currentTime + 1800000); // 30 minutes
+      }
+    });
+  });
+
+  describe('resumeToolCall', () => {
+    const mockHandler = vi.fn().mockResolvedValue({ result: 'success' });
+    let session: any;
+    let resumeToken: string;
+
+    beforeEach(async () => {
+      // Create a handshake session first
+      const handshakeResponse = await runtime.handleHandshake({
+        clientDid: 'did:key:zclient123',
+        audience: 'https://client.example.com'
+      });
+
+      session = {
+        id: handshakeResponse.sessionId, // Use actual session ID from handshake
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      // Intercept a tool call
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          { arg: 'value' },
+          mockHandler,
+          session
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        resumeToken = error.resumeToken;
+      }
+    });
+
+    it('should resume tool call with valid resume token', async () => {
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      const result = await runtime.resumeToolCall(
+        resumeToken,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      expect(mockHandler).toHaveBeenCalledWith({ arg: 'value' });
+      expect(result).toEqual({ result: 'success' });
+    });
+
+    it('should throw error for invalid resume token', async () => {
+      await expect(
+        runtime.resumeToolCall(
+          'invalid-resume-token',
+          mockHandler,
+          'delegation-token-123'
+        )
+      ).rejects.toThrow('Invalid or expired resume token');
+    });
+
+    it('should throw error for expired resume token', async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      
+      // Advance time past expiry (30 minutes + 1 second)
+      clock.setTime(Date.now() + (30 * 60 * 1000) + 1000);
+
+      await expect(
+        runtime.resumeToolCall(
+          resumeToken,
+          mockHandler,
+          'delegation-token-123'
+        )
+      ).rejects.toThrow('Resume token expired');
+    });
+
+    it('should throw error if session not found', async () => {
+      // Create intercepted call with non-existent session ID
+      const call = {
+        toolName: 'protectedTool',
+        args: { arg: 'value' },
+        sessionId: 'non-existent-session',
+        timestamp: Date.now(),
+        expiresAt: Date.now() + 1800000
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      const invalidToken = generateToken(call);
+      
+      // Manually store the intercepted call
+      (runtime as any).interceptedCalls.set(invalidToken, call);
+
+      await expect(
+        runtime.resumeToolCall(
+          invalidToken,
+          mockHandler,
+          'delegation-token-123'
+        )
+      ).rejects.toThrow('Session not found');
+    });
+
+    it('should add delegation token to session when resuming', async () => {
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      await runtime.resumeToolCall(
+        resumeToken,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      // Verify that processToolCall was called with session containing delegationToken
+      expect(mockHandler).toHaveBeenCalled();
+    });
+
+    it('should clean up intercepted call after successful resumption', async () => {
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      await runtime.resumeToolCall(
+        resumeToken,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      // Call should be removed from storage
+      const interceptedCalls = (runtime as any).interceptedCalls;
+      expect(interceptedCalls.has(resumeToken)).toBe(false);
+    });
+
+    it('should not clean up intercepted call if resumption fails', async () => {
+      const errorHandler = vi.fn().mockRejectedValue(new Error('Handler failed'));
+
+      try {
+        await runtime.resumeToolCall(
+          resumeToken,
+          errorHandler,
+          'delegation-token-123'
+        );
+        expect.fail('Should have thrown error');
+      } catch (error) {
+        // Expected
+      }
+
+      // Call should still be in storage (can retry)
+      const interceptedCalls = (runtime as any).interceptedCalls;
+      expect(interceptedCalls.has(resumeToken)).toBe(true);
+    });
+  });
+
+  describe('cleanupExpiredInterceptedCalls', () => {
+    it('should remove expired intercepted calls', async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      const expiredCall = {
+        toolName: 'expiredTool',
+        args: { arg: 'value' },
+        sessionId: 'session123',
+        timestamp: currentTime - 1000,
+        expiresAt: currentTime - 500 // Already expired
+      };
+
+      const validCall = {
+        toolName: 'validTool',
+        args: { arg: 'value' },
+        sessionId: 'session123',
+        timestamp: currentTime,
+        expiresAt: currentTime + 1800000 // Not expired
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      const expiredToken = generateToken(expiredCall);
+      const validToken = generateToken(validCall);
+
+      const interceptedCalls = (runtime as any).interceptedCalls;
+      interceptedCalls.set(expiredToken, expiredCall);
+      interceptedCalls.set(validToken, validCall);
+
+      // Trigger cleanup
+      const cleanup = (runtime as any).cleanupExpiredInterceptedCalls.bind(runtime);
+      cleanup();
+
+      expect(interceptedCalls.has(expiredToken)).toBe(false);
+      expect(interceptedCalls.has(validToken)).toBe(true);
+    });
+
+    it('should be called automatically when intercepting calls', async () => {
+      const clock = mockProviders.clockProvider as MockClockProvider;
+      const currentTime = Date.now();
+      clock.setTime(currentTime);
+
+      // Add an expired call
+      const expiredCall = {
+        toolName: 'expiredTool',
+        args: { arg: 'value' },
+        sessionId: 'session123',
+        timestamp: currentTime - 1000,
+        expiresAt: currentTime - 500
+      };
+
+      const generateToken = (runtime as any).generateResumeToken.bind(runtime);
+      const expiredToken = generateToken(expiredCall);
+
+      const interceptedCalls = (runtime as any).interceptedCalls;
+      interceptedCalls.set(expiredToken, expiredCall);
+
+      // Intercept a new call (should trigger cleanup)
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      try {
+        await runtime.processToolCall(
+          'newTool',
+          { arg: 'value' },
+          vi.fn(),
+          { id: 'session123' }
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error) {
+        // Expected
+      }
+
+      // Expired call should be cleaned up
+      expect(interceptedCalls.has(expiredToken)).toBe(false);
+    });
+  });
+
+  describe('integration: full interception flow', () => {
+    it('should complete full flow: intercept -> authorize -> resume', async () => {
+      const mockHandler = vi.fn().mockResolvedValue({ result: 'success' });
+      const session = {
+        id: 'session123',
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      // Step 1: Create session
+      const handshakeResponse = await runtime.handleHandshake({
+        clientDid: 'did:key:zclient123',
+        audience: 'https://client.example.com'
+      });
+
+      const sessionWithId = {
+        id: handshakeResponse.sessionId,
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      // Step 2: Intercept tool call
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      let resumeToken: string;
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          { arg: 'value' },
+          mockHandler,
+          sessionWithId
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        resumeToken = error.resumeToken;
+        expect(error).toBeInstanceOf(DelegationRequiredError);
+        expect(error.consentUrl).toBeDefined();
+      }
+
+      // Step 3: User authorizes (simulated by having delegation token)
+      // Step 4: Resume tool call
+      const result = await runtime.resumeToolCall(
+        resumeToken!,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      // Step 5: Verify successful execution
+      expect(mockHandler).toHaveBeenCalledWith({ arg: 'value' });
+      expect(result).toEqual({ result: 'success' });
+
+      // Step 6: Verify cleanup
+      const interceptedCalls = (runtime as any).interceptedCalls;
+      expect(interceptedCalls.has(resumeToken!)).toBe(false);
+    });
+  });
+
+  describe('edge cases', () => {
+    it('should handle resume token reuse attempt', async () => {
+      const mockHandler = vi.fn().mockResolvedValue({ result: 'success' });
+      
+      // Create session first
+      const handshakeResponse = await runtime.handleHandshake({
+        clientDid: 'did:key:zclient123',
+        audience: 'https://client.example.com'
+      });
+
+      const sessionWithId = {
+        id: handshakeResponse.sessionId,
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      let resumeToken: string;
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          { arg: 'value' },
+          mockHandler,
+          sessionWithId
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        resumeToken = error.resumeToken;
+      }
+
+      // First resume should succeed
+      await runtime.resumeToolCall(
+        resumeToken!,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      // Second resume should fail (token already consumed)
+      await expect(
+        runtime.resumeToolCall(
+          resumeToken!,
+          mockHandler,
+          'delegation-token-123'
+        )
+      ).rejects.toThrow('Invalid or expired resume token');
+    });
+
+    it('should handle complex nested args in intercepted calls', async () => {
+      const mockHandler = vi.fn().mockResolvedValue({ result: 'success' });
+      const session = {
+        id: 'session123',
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      const handshakeResponse = await runtime.handleHandshake({
+        clientDid: 'did:key:zclient123',
+        audience: 'https://client.example.com'
+      });
+
+      const sessionWithId = {
+        id: handshakeResponse.sessionId,
+        audience: 'https://client.example.com',
+        nonce: 'test-nonce'
+      };
+
+      mockToolProtectionService.checkToolProtection.mockResolvedValue({
+        requiresDelegation: true,
+        requiredScopes: ['files:write']
+      });
+
+      const complexArgs = {
+        nested: {
+          deeply: {
+            nested: {
+              value: 'test',
+              array: [1, 2, 3],
+              object: { key: 'value' }
+            }
+          }
+        },
+        array: [{ id: 1 }, { id: 2 }]
+      };
+
+      let resumeToken: string;
+      try {
+        await runtime.processToolCall(
+          'protectedTool',
+          complexArgs,
+          mockHandler,
+          sessionWithId
+        );
+        expect.fail('Should have thrown DelegationRequiredError');
+      } catch (error: any) {
+        resumeToken = error.resumeToken;
+      }
+
+      const result = await runtime.resumeToolCall(
+        resumeToken!,
+        mockHandler,
+        'delegation-token-123'
+      );
+
+      expect(mockHandler).toHaveBeenCalledWith(complexArgs);
+      expect(result).toEqual({ result: 'success' });
+    });
+  });
+});
+