@optimizely-opal/opal-tool-ocp-sdk 0.0.0-dev.4 → 0.0.0-devmg.11

package/src/function/ToolFunction.test.ts

@@ -1,21 +1,23 @@
 import { ToolFunction } from './ToolFunction';
 import { toolsService } from '../service/Service';
-import { Response } from '@zaiusinc/app-sdk';
+import { Response, getAppContext } from '@zaiusinc/app-sdk';
+import { getTokenVerifier } from '../auth/TokenVerifier';
 
-// Mock the toolsService
+// Mock the dependencies
 jest.mock('../service/Service', () => ({
   toolsService: {
     processRequest: jest.fn(),
-    extractBearerToken: jest.fn(),
   },
 }));
 
-// Mock the Request and Response classes and Function base class
+jest.mock('../auth/TokenVerifier', () => ({
+  getTokenVerifier: jest.fn(),
+}));
+
 jest.mock('@zaiusinc/app-sdk', () => ({
   Function: class {
     protected request: any;
     public constructor(_name?: string) {
-      // Mock constructor that accepts optional name parameter
       this.request = {};
     }
   },
@@ -26,32 +28,37 @@ jest.mock('@zaiusinc/app-sdk', () => ({
     bodyJSON: data,
     bodyAsU8Array: new Uint8Array()
   })),
+  amendLogContext: jest.fn(),
+  getAppContext: jest.fn(),
+  logger: {
+    info: jest.fn(),
+    error: jest.fn(),
+    warn: jest.fn(),
+    debug: jest.fn(),
+  },
 }));
 
 // Create a concrete implementation for testing
 class TestToolFunction extends ToolFunction {
-  private mockValidateBearerToken: jest.MockedFunction<(token: string) => boolean>;
+  private mockReady: jest.MockedFunction<() => Promise<boolean>>;
 
   public constructor(request?: any) {
-    super(request || {}); // Pass the request parameter properly
-    // Set the request directly without defaulting to empty object
+    super(request || {});
     (this as any).request = request;
-
-    this.mockValidateBearerToken = jest.fn().mockReturnValue(true);
+    this.mockReady = jest.fn().mockResolvedValue(true);
   }
 
-  // Override the concrete method with mock implementation for testing
-  protected validateBearerToken(bearerToken: string): boolean {
-    return this.mockValidateBearerToken(bearerToken);
+  // Override the ready method with mock implementation for testing
+  protected ready(): Promise<boolean> {
+    return this.mockReady();
   }
 
-  // Expose request and validation mock for testing
   public getRequest() {
     return (this as any).request;
   }
 
-  public getMockValidateBearerToken() {
-    return this.mockValidateBearerToken;
+  public getMockReady() {
+    return this.mockReady;
   }
 }
 
@@ -60,180 +67,296 @@ describe('ToolFunction', () => {
   let mockResponse: Response;
   let toolFunction: TestToolFunction;
   let mockProcessRequest: jest.MockedFunction<typeof toolsService.processRequest>;
-  let mockExtractBearerToken: jest.MockedFunction<typeof toolsService.extractBearerToken>;
+  let mockGetTokenVerifier: jest.MockedFunction<typeof getTokenVerifier>;
+  let mockGetAppContext: jest.MockedFunction<typeof getAppContext>;
+  let mockTokenVerifier: jest.Mocked<{
+    verify: (token: string) => Promise<any>;
+  }>;
 
   beforeEach(() => {
     jest.clearAllMocks();
 
-    // Create mock instances
-    mockRequest = {
-      headers: new Map(),
-      method: 'POST',
-      path: '/test'
+    // Create mock token verifier
+    mockTokenVerifier = {
+      verify: jest.fn(),
     };
-    mockResponse = {} as Response;
 
     // Setup the mocks
     mockProcessRequest = jest.mocked(toolsService.processRequest);
-    mockExtractBearerToken = jest.mocked(toolsService.extractBearerToken);
+    mockGetTokenVerifier = jest.mocked(getTokenVerifier);
+    mockGetAppContext = jest.mocked(getAppContext);
 
-    // Create test instance
+    mockGetTokenVerifier.mockResolvedValue(mockTokenVerifier as any);
+    mockGetAppContext.mockReturnValue({
+      account: {
+        organizationId: 'app-org-123'
+      }
+    } as any);
+
+    // Create mock request with bodyJSON structure
+    mockRequest = {
+      headers: new Map(),
+      method: 'POST',
+      path: '/test',
+      bodyJSON: {
+        parameters: {
+          task_id: 'task-123',
+          content_id: 'content-456'
+        },
+        auth: {
+          provider: 'OptiID',
+          credentials: {
+            token_type: 'Bearer',
+            access_token: 'valid-access-token',
+            org_sso_id: 'org-sso-123',
+            user_id: 'user-456',
+            instance_id: 'instance-789',
+            customer_id: 'app-org-123',
+            product_sku: 'OPAL'
+          }
+        },
+        environment: {
+          execution_mode: 'headless'
+        }
+      }
+    };
+
+    mockResponse = {} as Response;
     toolFunction = new TestToolFunction(mockRequest);
   });
 
-  describe('bearer token validation', () => {
-    it('should extract bearer token from headers and validate it', async () => {
-      // Arrange
-      const bearerToken = 'valid-token-123';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      toolFunction.getMockValidateBearerToken().mockReturnValue(true);
-      mockProcessRequest.mockResolvedValue(mockResponse);
+  // Helper function to create a ready request with valid auth
+  const createReadyRequestWithAuth = () => ({
+    headers: new Map(),
+    method: 'GET',
+    path: '/ready',
+    bodyJSON: {
+      auth: {
+        provider: 'OptiID',
+        credentials: {
+          access_token: 'valid-token',
+          customer_id: 'app-org-123'
+        }
+      }
+    }
+  });
 
-      // Act
-      const result = await toolFunction.perform();
+  // Helper function to setup authorization mocks to pass
+  const setupAuthMocks = () => {
+    mockTokenVerifier.verify.mockResolvedValue(true);
+    mockGetAppContext.mockReturnValue({
+      account: {
+        organizationId: 'app-org-123'
+      }
+    } as any);
+  };
 
-      // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledWith(bearerToken);
-      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, toolFunction);
-      expect(result).toBe(mockResponse);
+  describe('/ready endpoint', () => {
+    beforeEach(() => {
+      setupAuthMocks();
     });
 
-    it('should return 403 Forbidden when bearer token validation fails', async () => {
+    it('should return ready: true when ready method returns true', async () => {
       // Arrange
-      const bearerToken = 'invalid-token-456';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      toolFunction.getMockValidateBearerToken().mockReturnValue(false);
+      const readyRequest = createReadyRequestWithAuth();
+
+      toolFunction = new TestToolFunction(readyRequest);
+      toolFunction.getMockReady().mockResolvedValue(true);
 
       // Act
       const result = await toolFunction.perform();
 
       // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledWith(bearerToken);
-      expect(mockProcessRequest).not.toHaveBeenCalled();
-      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(toolFunction.getMockReady()).toHaveBeenCalledTimes(1);
+      expect(result).toEqual(new Response(200, { ready: true }));
+      expect(mockProcessRequest).not.toHaveBeenCalled(); // Should not call service
     });
 
-    it('should handle complex bearer token validation scenarios', async () => {
-      // Arrange - simulate a token that should be valid based on some complex logic
-      const complexToken = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9l' +
-        'IiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c';
-      mockExtractBearerToken.mockReturnValue(complexToken);
-      toolFunction.getMockValidateBearerToken().mockReturnValue(true);
-      mockProcessRequest.mockResolvedValue(mockResponse);
+    it('should return ready: false when ready method returns false', async () => {
+      // Arrange
+      const readyRequest = createReadyRequestWithAuth();
+
+      toolFunction = new TestToolFunction(readyRequest);
+      toolFunction.getMockReady().mockResolvedValue(false);
 
       // Act
       const result = await toolFunction.perform();
 
       // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledWith(complexToken);
-      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, toolFunction);
-      expect(result).toBe(mockResponse);
+      expect(toolFunction.getMockReady()).toHaveBeenCalledTimes(1);
+      expect(result).toEqual(new Response(200, { ready: false }));
+      expect(mockProcessRequest).not.toHaveBeenCalled(); // Should not call service
     });
 
-    it('should handle bearer token validation throwing an error', async () => {
+    it('should handle ready method throwing an error', async () => {
       // Arrange
-      const bearerToken = 'malformed-token';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      toolFunction.getMockValidateBearerToken().mockImplementation(() => {
-        throw new Error('Token validation error');
-      });
+      const readyRequest = createReadyRequestWithAuth();
+
+      toolFunction = new TestToolFunction(readyRequest);
+      toolFunction.getMockReady().mockRejectedValue(new Error('Ready check failed'));
 
       // Act & Assert
-      await expect(toolFunction.perform()).rejects.toThrow('Token validation error');
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledWith(bearerToken);
-      expect(mockProcessRequest).not.toHaveBeenCalled();
+      await expect(toolFunction.perform()).rejects.toThrow('Ready check failed');
+      expect(toolFunction.getMockReady()).toHaveBeenCalledTimes(1);
+      expect(mockProcessRequest).not.toHaveBeenCalled(); // Should not call service
     });
 
-    it('should call validateBearerToken only once per request', async () => {
-      // Arrange
-      const bearerToken = 'test-token';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      toolFunction.getMockValidateBearerToken().mockReturnValue(true);
-      mockProcessRequest.mockResolvedValue(mockResponse);
-
-      // Act
-      await toolFunction.perform();
+    it('should use default ready implementation when not overridden', async () => {
+      // Create a class that doesn't override ready method
+      class DefaultReadyToolFunction extends ToolFunction {
+        public constructor(request?: any) {
+          super(request || {});
+          (this as any).request = request;
+        }
 
-      // Assert
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledTimes(1);
-      expect(toolFunction.getMockValidateBearerToken()).toHaveBeenCalledWith(bearerToken);
-    });
+        public getRequest() {
+          return (this as any).request;
+        }
+      }
 
-    it('should extract bearer token only once per request', async () => {
       // Arrange
-      const bearerToken = 'test-token';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      toolFunction.getMockValidateBearerToken().mockReturnValue(true);
-      mockProcessRequest.mockResolvedValue(mockResponse);
+      const readyRequest = createReadyRequestWithAuth();
+      const defaultToolFunction = new DefaultReadyToolFunction(readyRequest);
 
       // Act
-      await toolFunction.perform();
+      const result = await defaultToolFunction.perform();
 
-      // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledTimes(1);
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
+      // Assert - Default implementation should return true
+      expect(result).toEqual(new Response(200, { ready: true }));
+      expect(mockProcessRequest).not.toHaveBeenCalled(); // Should not call service
     });
+  });
 
-    it('should skip validation when bearer token is undefined', async () => {
-      // Arrange
-      mockExtractBearerToken.mockReturnValue(undefined);
+  describe('perform', () => {
+    it('should execute successfully with valid token and matching organization', async () => {
+      // Setup mock token verifier to return true for valid token
+      mockTokenVerifier.verify.mockResolvedValue(true);
       mockProcessRequest.mockResolvedValue(mockResponse);
 
-      // Act
       const result = await toolFunction.perform();
 
-      // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).not.toHaveBeenCalled();
-      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, toolFunction);
       expect(result).toBe(mockResponse);
+      expect(mockGetTokenVerifier).toHaveBeenCalled();
+      expect(mockTokenVerifier.verify).toHaveBeenCalledWith('valid-access-token');
+      expect(mockGetAppContext).toHaveBeenCalled();
+      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, toolFunction);
     });
 
-    it('should skip validation when bearer token is empty string', async () => {
-      // Arrange
-      mockExtractBearerToken.mockReturnValue('');
-      mockProcessRequest.mockResolvedValue(mockResponse);
+    it('should return 403 response with invalid token', async () => {
+      // Setup mock token verifier to return false
+      mockTokenVerifier.verify.mockResolvedValue(false);
 
-      // Act
       const result = await toolFunction.perform();
 
-      // Assert
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(toolFunction.getMockValidateBearerToken()).not.toHaveBeenCalled();
-      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, toolFunction);
-      expect(result).toBe(mockResponse);
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetTokenVerifier).toHaveBeenCalled();
+      expect(mockTokenVerifier.verify).toHaveBeenCalledWith('valid-access-token');
+      expect(mockProcessRequest).not.toHaveBeenCalled();
     });
 
-    it('should use default validateBearerToken implementation when not overridden', async () => {
-      // Create a class that doesn't override validateBearerToken
-      class DefaultToolFunction extends ToolFunction {
-        public constructor(request?: any) {
-          super(request || {});
-          (this as any).request = request;
+    it('should return 403 response when organization ID does not match', async () => {
+      // Update mock request with different customer_id
+      const requestWithDifferentOrgId = {
+        ...mockRequest,
+        bodyJSON: {
+          ...mockRequest.bodyJSON,
+          auth: {
+            ...mockRequest.bodyJSON.auth,
+            credentials: {
+              ...mockRequest.bodyJSON.auth.credentials,
+              customer_id: 'different-org-123'
+            }
+          }
         }
+      };
 
-        public getRequest() {
-          return (this as any).request;
+      const toolFunctionWithDifferentOrgId = new TestToolFunction(requestWithDifferentOrgId);
+
+      const result = await toolFunctionWithDifferentOrgId.perform();
+
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetAppContext).toHaveBeenCalled();
+      expect(mockProcessRequest).not.toHaveBeenCalled();
+    });
+
+    it('should return 403 response when access token is missing', async () => {
+      // Create request without access token
+      const requestWithoutToken = {
+        ...mockRequest,
+        bodyJSON: {
+          ...mockRequest.bodyJSON,
+          auth: {
+            ...mockRequest.bodyJSON.auth,
+            credentials: {
+              ...mockRequest.bodyJSON.auth.credentials,
+              access_token: undefined
+            }
+          }
         }
-      }
+      };
 
-      // Arrange
-      const defaultToolFunction = new DefaultToolFunction(mockRequest);
-      const bearerToken = 'any-token';
-      mockExtractBearerToken.mockReturnValue(bearerToken);
-      mockProcessRequest.mockResolvedValue(mockResponse);
+      const toolFunctionWithoutToken = new TestToolFunction(requestWithoutToken);
 
-      // Act
-      const result = await defaultToolFunction.perform();
+      const result = await toolFunctionWithoutToken.perform();
 
-      // Assert - Default implementation should return true and allow request to proceed
-      expect(mockExtractBearerToken).toHaveBeenCalledWith(mockRequest.headers);
-      expect(mockProcessRequest).toHaveBeenCalledWith(mockRequest, defaultToolFunction);
-      expect(result).toBe(mockResponse);
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetTokenVerifier).not.toHaveBeenCalled();
+      expect(mockProcessRequest).not.toHaveBeenCalled();
+    });
+
+    it('should return 403 response when organisation id is missing', async () => {
+      // Create request without customer_id
+      const requestWithoutCustomerId = {
+        ...mockRequest,
+        bodyJSON: {
+          ...mockRequest.bodyJSON,
+          auth: {
+            ...mockRequest.bodyJSON.auth,
+            credentials: {
+              ...mockRequest.bodyJSON.auth.credentials,
+              customer_id: undefined
+            }
+          }
+        }
+      };
+
+      const toolFunctionWithoutCustomerId = new TestToolFunction(requestWithoutCustomerId);
+
+      const result = await toolFunctionWithoutCustomerId.perform();
+
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetTokenVerifier).not.toHaveBeenCalled();
+      expect(mockProcessRequest).not.toHaveBeenCalled();
+    });
+
+    it('should return 403 response when auth structure is missing', async () => {
+      // Create request without auth structure
+      const requestWithoutAuth = {
+        ...mockRequest,
+        bodyJSON: {
+          parameters: mockRequest.bodyJSON.parameters,
+          environment: mockRequest.bodyJSON.environment
+        }
+      };
+
+      const toolFunctionWithoutAuth = new TestToolFunction(requestWithoutAuth);
+
+      const result = await toolFunctionWithoutAuth.perform();
+
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetTokenVerifier).not.toHaveBeenCalled();
+      expect(mockProcessRequest).not.toHaveBeenCalled();
+    });
+
+    it('should return 403 response when token verifier initialization fails', async () => {
+      // Setup mock to fail during token verifier initialization
+      mockGetTokenVerifier.mockRejectedValue(new Error('Failed to initialize token verifier'));
+
+      const result = await toolFunction.perform();
+
+      expect(result).toEqual(new Response(403, { error: 'Forbidden' }));
+      expect(mockGetTokenVerifier).toHaveBeenCalled();
+      expect(mockProcessRequest).not.toHaveBeenCalled();
     });
   });
 

package/src/function/ToolFunction.ts

@@ -1,5 +1,7 @@
-import { Function, Response } from '@zaiusinc/app-sdk';
+import { Function, Response, amendLogContext, getAppContext, logger } from '@zaiusinc/app-sdk';
 import { toolsService } from '../service/Service';
+import { getTokenVerifier } from '../auth/TokenVerifier';
+import { OptiIdAuthData } from '../types/Models';
 
 /**
  * Abstract base class for tool-based function execution
@@ -7,31 +9,78 @@ import { toolsService } from '../service/Service';
  */
 export abstract class ToolFunction extends Function {
 
+  /**
+   * Override this method to implement any required credentials and/or other configuration
+   * exist and are valid. Reasonable caching should be utilized to prevent excessive requests to external resources.
+   * @async
+   * @returns true if the opal function is ready to use
+   */
+  protected ready(): Promise<boolean> {
+    return Promise.resolve(true);
+  }
+
   /**
    * Process the incoming request using the tools service
    *
    * @returns Response as the HTTP response
    */
   public async perform(): Promise<Response> {
-    const bearerToken = toolsService.extractBearerToken(this.request.headers);
-    if (bearerToken && !this.validateBearerToken(bearerToken)) {
+    amendLogContext({ opalThreadId: this.request.headers.get('x-opal-thread-id') || '' });
+    if (!(await this.authorizeRequest())) {
       return new Response(403, { error: 'Forbidden' });
     }
 
+    if (this.request.path === '/ready') {
+      const isReady = await this.ready();
+      return new Response(200, { ready: isReady });
+    }
     // Pass 'this' as context so decorated methods can use the existing instance
     return toolsService.processRequest(this.request, this);
   }
 
   /**
-   * Validates the bearer token for authorization.
+   * Authenticate the incoming request by validating the OptiID token and organization ID
    *
-   * This method provides a default implementation that accepts all tokens.
-   * Subclasses can override this method to implement custom bearer token validation logic.
-   *
-   * @param _bearerToken - The bearer token extracted from the Authorization header
-   * @returns true if the token is valid and the request should proceed, false to return 403 Forbidden
+   * @throws true if authentication succeeds
    */
-  protected validateBearerToken(_bearerToken: string): boolean {
-    return true;
+  private async authorizeRequest(): Promise<boolean> {
+    logger.info('Authorizing request:', this.request.bodyJSON);
+    if (this.request.path === '/discovery' || this.request.path === '/ready') {
+      return true;
+    }
+    const authData = this.request.bodyJSON?.auth as OptiIdAuthData;
+    const accessToken = authData?.credentials?.access_token;
+    if (!accessToken || authData?.provider?.toLowerCase() !== 'optiid') {
+      logger.error('OptiID token is required but not provided');
+      return false;
+    }
+
+    const customerId = authData.credentials?.customer_id;
+    if (!customerId) {
+      logger.error('Organisation ID is required but not provided');
+      return false;
+    }
+
+    const appOrganisationId = getAppContext().account?.organizationId;
+    if (customerId !== appOrganisationId) {
+      logger.error(`Invalid organisation ID: expected ${appOrganisationId}, received ${customerId}`);
+      return false;
+    }
+
+    return await this.validateAccessToken(accessToken);
   }
+
+  private async validateAccessToken(accessToken: string | undefined): Promise<boolean> {
+    try {
+      if (!accessToken) {
+        return false;
+      }
+      const tokenVerifier = await getTokenVerifier();
+      return await tokenVerifier.verify(accessToken);
+    } catch (error) {
+      logger.error('OptiID token validation failed:', error);
+      return false;
+    }
+  }
+
 }

package/src/index.ts

@@ -1,4 +1,5 @@
 export * from './function/ToolFunction';
 export * from './types/Models';
 export * from './decorator/Decorator';
+export * from './auth/TokenVerifier';
 export { Tool, Interaction, InteractionResult } from './service/Service';