linear integration

Author: rushilpatel0

src/codegen/extensions/clients/linear.py

@@ -1,23 +1,26 @@
 import json
 import logging
-import requests
 
+import requests
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
 
 # --- TYPES
 
+
 class LinearUser(BaseModel):
     id: str
     name: str
 
+
 class LinearComment(BaseModel):
     id: str
     body: str
     user: LinearUser | None = None
 
+
 class LinearIssue(BaseModel):
     id: str
     title: str
@@ -49,12 +52,7 @@ def get_issue(self, issue_id: str) -> LinearIssue:
         response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables})
         data = response.json()
         issue_data = data["data"]["issue"]
-        return LinearIssue(
-            id=issue_data["id"],
-            title=issue_data["title"],
-            description=issue_data["description"]
-        )
-        
+        return LinearIssue(id=issue_data["id"], title=issue_data["title"], description=issue_data["description"])
 
     def get_issue_comments(self, issue_id: str) -> list[LinearComment]:
         query = """
@@ -69,7 +67,7 @@ def get_issue_comments(self, issue_id: str) -> list[LinearComment]:
                             name
                         }
                     }
-                        
+
                     }
                 }
             }
@@ -83,21 +81,13 @@ def get_issue_comments(self, issue_id: str) -> list[LinearComment]:
         parsed_comments = []
         for comment in comments:
             user = comment.get("user", None)
-            parsed_comment = LinearComment(
-                id=comment["id"],
-                body=comment["body"],
-                user=LinearUser(
-                    id=user.get("id"),
-                    name=user.get("name")
-                ) if user else None
-            )
+            parsed_comment = LinearComment(id=comment["id"], body=comment["body"], user=LinearUser(id=user.get("id"), name=user.get("name")) if user else None)
             parsed_comments.append(parsed_comment)
 
         # Convert raw comments to LinearComment objects
         return parsed_comments
 
-
-    def comment_on_issue(self, issue_id: str, body: str) -> dict:   
+    def comment_on_issue(self, issue_id: str, body: str) -> dict:
         """issue_id is our internal issue ID"""
         query = """mutation makeComment($issueId: String!, $body: String!) {
           commentCreate(input: {issueId: $issueId, body: $body}) {
@@ -106,7 +96,7 @@ def comment_on_issue(self, issue_id: str, body: str) -> dict:
               body
               url
               user {
-                id  
+                id
                 name
               }
             }
@@ -125,30 +115,48 @@ def comment_on_issue(self, issue_id: str, body: str) -> dict:
 
             return comment_data
         except:
-            raise Exception(f"Error creating comment\n{data}")
+            msg = f"Error creating comment\n{data}"
+            raise Exception(msg)
 
+    def unregister_webhook(self, webhook_id: str):
+        mutation = """
+            mutation deleteWebhook($id: String!) {
+                webhookDelete(id: $id) {
+                    success
+                }
+            }
+        """
+        variables = {"id": webhook_id}
+        response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables})
+        return response.json()
 
     def register_webhook(self, webhook_url: str, team_id: str, secret: str, enabled: bool, resource_types: list[str]):
         mutation = """
             mutation createWebhook($input: WebhookCreateInput!) {
                 webhookCreate(input: $input) {
                     success
                     webhook {
-                    id
-                    enabled
+                        id
+                        enabled
                     }
                 }
             }
         """
 
-        variables = {"input": {
-            "url": webhook_url,
-            "teamId": team_id,
-            "resourceTypes": resource_types,
-            "enabled": enabled,
-            "secret": secret,
-        }}
+        variables = {
+            "input": {
+                "url": webhook_url,
+                "teamId": team_id,
+                "resourceTypes": resource_types,
+                "enabled": enabled,
+                "secret": secret,
+            }
+        }
 
         response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables})
+        if response.status_code != 200:
+            return None
+
         body = response.json()
+        body = body["data"]["webhookCreate"]["webhook"]["id"]
         return body

src/codegen/extensions/events/app.py

@@ -1,99 +1,21 @@
-import functools
-import os
-from typing import Callable, Optional
+import logging
+
 import modal
-import modal.runner
 
-from codegen.extensions.clients.linear import LinearClient
+from codegen.extensions.events.linear import Linear
+
+logger = logging.getLogger(__name__)
+
+
+class CodegenApp(modal.App):
+    linear: Linear
 
-class CodegenApp:
     def __init__(self, name, modal_api_key, image: modal.Image):
         self._modal_api_key = modal_api_key
         self._image = image
         self._name = name
 
-        self.modal_app = modal.App(name=name, image=image, )
-        # Internal registry for event handlers, keyed by event name.
-        self._registry = {}
-        # Expose a linear attribute that provides the event decorator.
-        self.linear = Linear(self)
-
-    def _register_event_handler(self, event_name, handler):
-        """Registers a handler for a given event."""
-        self._registry.setdefault(event_name, []).append(handler)
-        print(f"Registered handler '{handler}' for event '{event_name}'.")
-
-    def dispatch_event(self, event_name, request):
-        """
-        Dispatches an event to all registered handlers for the event.
-        """
-        handlers = self._registry.get(event_name, [])
-        if not handlers:
-            print(f"No handlers registered for event '{event_name}'.")
-            return
-        for handler in handlers:
-            print(f"Dispatching event '{event_name}' to {handler.__name__}.")
-            handler(request)
-
-
-
-class Linear:
-    def __init__(self, app: CodegenApp):
-        self.app = app
-        self.access_token= os.environ["LINEAR_ACCESS_TOKEN"] # move to extensions config.
-        self.signing_secret = os.environ["LINEAR_SIGNING_SECRET"]
-        self.linear_team_id = os.environ["LINEAR_TEAM_ID"]
+        super().__init__(name=name, image=image)
 
-
-    def register_handler_to_webhook(self, func_name: str, modal_app: modal.App, event_name):
-        app_name = modal_app.name
-        web_url = modal.Function.from_name(app_name=app_name,name=func_name)
-        client = LinearClient(access_token=self.access_token)
-        client.register_webhook(
-            team_id=self.linear_team_id, 
-            webhook_url=web_url, 
-            enabled=True, 
-            resource_types=[event_name], 
-            secret=self.signing_secret)
-        
-    # def unregister_handler(self):
-        # TODO
-        # pass
-
-
-    def event(self, event_name, register_hook: Optional[Callable]=None):
-        """
-        Decorator for registering an event handler.
-
-        :param event_name: The name of the event to handle.
-        :param register_hook: An optional function to call during registration,
-                              e.g., to make an API call to register the webhook.
-        """
-        def decorator(func):
-            # Register the handler with the app's registry.
-
-            modal_ready_func = apply_decorators(func=func, decorators=[modal.web_endpoint(method="POST"), self.app.modal_app.function()])
-            self.app._register_event_handler(event_name, modal_ready_func)
-
-            # If a custom registration hook is provided, execute it.
-            if callable(register_hook):
-                register_hook(event_name, func, self.app)
-
-            @functools.wraps(func)
-            def wrapper(*args, **kwargs):
-                return func(*args, **kwargs)
-            
-            result = modal.runner.deploy_app(self.app.modal_app)
-            print(result)
-            return wrapper
-            
-        
-        return decorator
-    
-        
-    
-
-def apply_decorators(func, decorators):
-    for decorator in decorators:
-        func = decorator(func)
-    return func
+        # Expose a attribute that provides the event decorator for different providers.
+        self.linear = Linear(self)

src/codegen/extensions/events/interface.py

@@ -0,0 +1,14 @@
+from typing import Protocol
+
+import modal
+
+
+class EventHandlerManagerProtocol(Protocol):
+    def subscribe_handler_to_webhook(self, func_name: str, modal_app: modal.App, event_name):
+        pass
+
+    def unsubscribe_handler_to_webhook(self, func_name: str, modal_app: modal.App, event_name):
+        pass
+
+    def unsubscribe_all_handlers(self):
+        pass

src/codegen/extensions/events/linear.py

@@ -0,0 +1,76 @@
+import functools
+import logging
+import os
+from typing import Callable
+
+import modal
+from anthropic import BaseModel
+
+from codegen.extensions.clients.linear import LinearClient
+from codegen.extensions.events.interface import EventHandlerManagerProtocol
+
+logger = logging.getLogger(__name__)
+
+
+class RegisteredWebhookHandler(BaseModel):
+    webhook_id: str | None = None
+    handler_func: Callable
+
+
+class Linear(EventHandlerManagerProtocol):
+    def __init__(self, app: modal.App):
+        self.app = app
+        self.access_token = os.environ["LINEAR_ACCESS_TOKEN"]  # move to extensions config.
+        self.signing_secret = os.environ["LINEAR_SIGNING_SECRET"]
+        self.linear_team_id = os.environ["LINEAR_TEAM_ID"]
+        self.registered_handlers = {}
+
+    def subscribe_handler_to_webhook(self, web_url: str, event_name: str):
+        client = LinearClient(access_token=self.access_token)
+
+        result = client.register_webhook(team_id=self.linear_team_id, webhook_url=web_url, enabled=True, resource_types=[event_name], secret=self.signing_secret)
+        return result
+
+    def unsubscribe_handler_to_webhook(self, registered_handler: RegisteredWebhookHandler):
+        webhook_id = registered_handler.webhook_id
+
+        client = LinearClient(access_token=self.access_token)
+        if webhook_id:
+            print(f"Unsubscribing from webhook {webhook_id}")
+            result = client.unregister_webhook(webhook_id)
+            return result
+        else:
+            print("No webhook id found for handler")
+            return None
+
+    def unsubscribe_all_handlers(self):
+        for handler in self.registered_handlers:
+            self.unsubscribe_handler_to_webhook(self.registered_handlers[handler])
+
+    def event(self, event_name):
+        """Decorator for registering an event handler.
+
+        :param event_name: The name of the event to handle.
+        :param register_hook: An optional function to call during registration,
+                              e.g., to make an API call to register the webhook.
+        """
+
+        def decorator(func):
+            # Register the handler with the app's registry.
+            modal_ready_func = func
+            func_name = func.__qualname__
+            app_name = self.app.name
+            web_url = modal.Function.from_name(app_name=app_name, name=func_name).web_url
+
+            self.registered_handlers[func_name] = RegisteredWebhookHandler(handler_func=modal_ready_func)
+
+            webhook_id = self.subscribe_handler_to_webhook(web_url=web_url, event_name=event_name)
+            self.registered_handlers[func_name].webhook_id = webhook_id
+
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator