From ac864113b7b333faa07fc0b7cf15d841a7312368 Mon Sep 17 00:00:00 2001 From: tomcodgen Date: Wed, 26 Feb 2025 00:46:54 +0100 Subject: [PATCH] [CG-10837] feat: Linear tools error handling, extra test, request retry (#653) # Motivation # Content # Testing # Please check the following before marking your PR as ready for review - [ ] I have added tests for my changes - [ ] I have updated the documentation or added new documentation as needed --------- Co-authored-by: tomcodgen <191515280+tomcodgen@users.noreply.github.com> --- pyproject.toml | 1 + .../extensions/linear/linear_client.py | 43 +++-- src/codegen/extensions/tools/linear/linear.py | 179 ++++++++++++++++++ tests/integration/extension/test_linear.py | 16 ++ uv.lock | 4 + 5 files changed, 229 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8a3c851a0..5f1fbe6b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ dependencies = [ "lox>=0.12.0", "httpx>=0.28.1", "docker>=6.1.3", + "urllib3>=2.0.0", ] license = { text = "Apache-2.0" } diff --git a/src/codegen/extensions/linear/linear_client.py b/src/codegen/extensions/linear/linear_client.py index 37291b806..f9c2bef05 100644 --- a/src/codegen/extensions/linear/linear_client.py +++ b/src/codegen/extensions/linear/linear_client.py @@ -1,9 +1,10 @@ -import json import logging import os from typing import Optional import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry from codegen.extensions.linear.types import LinearComment, LinearIssue, LinearTeam, LinearUser @@ -14,7 +15,7 @@ class LinearClient: api_headers: dict api_endpoint = "https://api.linear.app/graphql" - def __init__(self, access_token: Optional[str] = None, team_id: Optional[str] = None): + def __init__(self, access_token: Optional[str] = None, team_id: Optional[str] = None, max_retries: int = 3, backoff_factor: float = 0.5): if not access_token: access_token = os.getenv("LINEAR_ACCESS_TOKEN") if not access_token: @@ -31,6 +32,18 @@ def __init__(self, access_token: Optional[str] = None, team_id: Optional[str] = "Authorization": self.access_token, } + # Set up a session with retry logic + self.session = requests.Session() + retry_strategy = Retry( + total=max_retries, + backoff_factor=backoff_factor, + status_forcelist=[429, 500, 502, 503, 504], + allowed_methods=["POST", "GET"], # POST is important for GraphQL + ) + adapter = HTTPAdapter(max_retries=retry_strategy) + self.session.mount("https://", adapter) + self.session.mount("http://", adapter) + def get_issue(self, issue_id: str) -> LinearIssue: query = """ query getIssue($issueId: String!) { @@ -42,7 +55,7 @@ def get_issue(self, issue_id: str) -> LinearIssue: } """ variables = {"issueId": issue_id} - response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables}) + response = self.session.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"]) @@ -66,7 +79,7 @@ def get_issue_comments(self, issue_id: str) -> list[LinearComment]: } """ variables = {"issueId": issue_id} - response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables}) + response = self.session.post(self.api_endpoint, headers=self.api_headers, json={"query": query, "variables": variables}) data = response.json() comments = data["data"]["issue"]["comments"]["nodes"] @@ -80,8 +93,8 @@ def get_issue_comments(self, issue_id: str) -> list[LinearComment]: # Convert raw comments to LinearComment objects return parsed_comments - def comment_on_issue(self, issue_id: str, body: str) -> dict: - """issue_id is our internal issue ID""" + def comment_on_issue(self, issue_id: str, body: str) -> LinearComment: + """Add a comment to an issue.""" query = """mutation makeComment($issueId: String!, $body: String!) { commentCreate(input: {issueId: $issueId, body: $body}) { comment { @@ -97,19 +110,21 @@ def comment_on_issue(self, issue_id: str, body: str) -> dict: } """ variables = {"issueId": issue_id, "body": body} - response = requests.post( + response = self.session.post( self.api_endpoint, headers=self.api_headers, - data=json.dumps({"query": query, "variables": variables}), + json={"query": query, "variables": variables}, ) data = response.json() try: comment_data = data["data"]["commentCreate"]["comment"] + user_data = comment_data.get("user", None) + user = LinearUser(id=user_data["id"], name=user_data["name"]) if user_data else None - return comment_data + return LinearComment(id=comment_data["id"], body=comment_data["body"], user=user) except Exception as e: msg = f"Error creating comment\n{data}\n{e}" - raise Exception(msg) + raise ValueError(msg) def register_webhook(self, webhook_url: str, team_id: str, secret: str, enabled: bool, resource_types: list[str]): mutation = """ @@ -134,7 +149,7 @@ def register_webhook(self, webhook_url: str, team_id: str, secret: str, enabled: } } - response = requests.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables}) + response = self.session.post(self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables}) body = response.json() return body @@ -160,7 +175,7 @@ def search_issues(self, query: str, limit: int = 10) -> list[LinearIssue]: } """ variables = {"query": query, "limit": limit} - response = requests.post( + response = self.session.post( self.api_endpoint, headers=self.api_headers, json={"query": graphql_query, "variables": variables}, @@ -222,7 +237,7 @@ def create_issue(self, title: str, description: str | None = None, team_id: str } } - response = requests.post( + response = self.session.post( self.api_endpoint, headers=self.api_headers, json={"query": mutation, "variables": variables}, @@ -258,7 +273,7 @@ def get_teams(self) -> list[LinearTeam]: } """ - response = requests.post( + response = self.session.post( self.api_endpoint, headers=self.api_headers, json={"query": query}, diff --git a/src/codegen/extensions/tools/linear/linear.py b/src/codegen/extensions/tools/linear/linear.py index d079f45d7..31c534bd7 100644 --- a/src/codegen/extensions/tools/linear/linear.py +++ b/src/codegen/extensions/tools/linear/linear.py @@ -2,6 +2,7 @@ from typing import ClassVar +import requests from pydantic import Field from codegen.extensions.linear.linear_client import LinearClient @@ -97,7 +98,32 @@ def linear_get_issue_tool(client: LinearClient, issue_id: str) -> LinearIssueObs issue_id=issue_id, issue_data=issue.dict(), ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearIssueObservation( + status="error", + error=f"Network error when fetching issue: {e!s}", + issue_id=issue_id, + issue_data={}, + ) + except ValueError as e: + # Input validation errors + return LinearIssueObservation( + status="error", + error=f"Invalid input: {e!s}", + issue_id=issue_id, + issue_data={}, + ) + except KeyError as e: + # Missing data in response + return LinearIssueObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + issue_id=issue_id, + issue_data={}, + ) except Exception as e: + # Catch-all for other errors return LinearIssueObservation( status="error", error=f"Failed to get issue: {e!s}", @@ -115,7 +141,32 @@ def linear_get_issue_comments_tool(client: LinearClient, issue_id: str) -> Linea issue_id=issue_id, comments=[comment.dict() for comment in comments], ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearCommentsObservation( + status="error", + error=f"Network error when fetching comments: {e!s}", + issue_id=issue_id, + comments=[], + ) + except ValueError as e: + # Input validation errors + return LinearCommentsObservation( + status="error", + error=f"Invalid input: {e!s}", + issue_id=issue_id, + comments=[], + ) + except KeyError as e: + # Missing data in response + return LinearCommentsObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + issue_id=issue_id, + comments=[], + ) except Exception as e: + # Catch-all for other errors return LinearCommentsObservation( status="error", error=f"Failed to get issue comments: {e!s}", @@ -133,7 +184,32 @@ def linear_comment_on_issue_tool(client: LinearClient, issue_id: str, body: str) issue_id=issue_id, comment=comment, ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearCommentObservation( + status="error", + error=f"Network error when adding comment: {e!s}", + issue_id=issue_id, + comment={}, + ) + except ValueError as e: + # Input validation errors + return LinearCommentObservation( + status="error", + error=f"Invalid input: {e!s}", + issue_id=issue_id, + comment={}, + ) + except KeyError as e: + # Missing data in response + return LinearCommentObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + issue_id=issue_id, + comment={}, + ) except Exception as e: + # Catch-all for other errors return LinearCommentObservation( status="error", error=f"Failed to comment on issue: {e!s}", @@ -159,7 +235,35 @@ def linear_register_webhook_tool( team_id=team_id, response=response, ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearWebhookObservation( + status="error", + error=f"Network error when registering webhook: {e!s}", + webhook_url=webhook_url, + team_id=team_id, + response={}, + ) + except ValueError as e: + # Input validation errors + return LinearWebhookObservation( + status="error", + error=f"Invalid input: {e!s}", + webhook_url=webhook_url, + team_id=team_id, + response={}, + ) + except KeyError as e: + # Missing data in response + return LinearWebhookObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + webhook_url=webhook_url, + team_id=team_id, + response={}, + ) except Exception as e: + # Catch-all for other errors return LinearWebhookObservation( status="error", error=f"Failed to register webhook: {e!s}", @@ -178,7 +282,32 @@ def linear_search_issues_tool(client: LinearClient, query: str, limit: int = 10) query=query, issues=[issue.dict() for issue in issues], ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearSearchObservation( + status="error", + error=f"Network error when searching issues: {e!s}", + query=query, + issues=[], + ) + except ValueError as e: + # Input validation errors + return LinearSearchObservation( + status="error", + error=f"Invalid input: {e!s}", + query=query, + issues=[], + ) + except KeyError as e: + # Missing data in response + return LinearSearchObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + query=query, + issues=[], + ) except Exception as e: + # Catch-all for other errors return LinearSearchObservation( status="error", error=f"Failed to search issues: {e!s}", @@ -197,7 +326,35 @@ def linear_create_issue_tool(client: LinearClient, title: str, description: str team_id=team_id, issue_data=issue.dict(), ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearCreateIssueObservation( + status="error", + error=f"Network error when creating issue: {e!s}", + title=title, + team_id=team_id, + issue_data={}, + ) + except ValueError as e: + # Input validation errors + return LinearCreateIssueObservation( + status="error", + error=f"Invalid input: {e!s}", + title=title, + team_id=team_id, + issue_data={}, + ) + except KeyError as e: + # Missing data in response + return LinearCreateIssueObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + title=title, + team_id=team_id, + issue_data={}, + ) except Exception as e: + # Catch-all for other errors return LinearCreateIssueObservation( status="error", error=f"Failed to create issue: {e!s}", @@ -215,7 +372,29 @@ def linear_get_teams_tool(client: LinearClient) -> LinearTeamsObservation: status="success", teams=[team.dict() for team in teams], ) + except requests.exceptions.RequestException as e: + # Network-related errors + return LinearTeamsObservation( + status="error", + error=f"Network error when fetching teams: {e!s}", + teams=[], + ) + except ValueError as e: + # Input validation errors + return LinearTeamsObservation( + status="error", + error=f"Invalid input: {e!s}", + teams=[], + ) + except KeyError as e: + # Missing data in response + return LinearTeamsObservation( + status="error", + error=f"Unexpected API response format: {e!s}", + teams=[], + ) except Exception as e: + # Catch-all for other errors return LinearTeamsObservation( status="error", error=f"Failed to get teams: {e!s}", diff --git a/tests/integration/extension/test_linear.py b/tests/integration/extension/test_linear.py index ece4ab677..2ad1ce043 100644 --- a/tests/integration/extension/test_linear.py +++ b/tests/integration/extension/test_linear.py @@ -3,6 +3,7 @@ import os import pytest +import requests from codegen.extensions.linear.linear_client import LinearClient from codegen.extensions.tools.linear.linear import ( @@ -98,3 +99,18 @@ def test_get_teams(client: LinearClient) -> None: assert "id" in team assert "name" in team assert "key" in team + + +def test_linear_get_issue_network_error(client): + """Test handling of network errors.""" + # Setup mock to raise network error + client.get_issue.side_effect = requests.exceptions.ConnectionError("Network error") + + # Call function + result = linear_get_issue_tool(client, "TEST-123") + + # Verify + assert result.status == "error" + assert "Network error" in result.error + assert result.issue_id == "TEST-123" + assert result.issue_data == {} diff --git a/uv.lock b/uv.lock index 6e6175e1f..39b85167b 100644 --- a/uv.lock +++ b/uv.lock @@ -1,4 +1,5 @@ version = 1 +revision = 1 requires-python = ">=3.12, <3.14" resolution-markers = [ "python_full_version >= '3.12.4'", @@ -602,6 +603,7 @@ dependencies = [ { name = "tree-sitter-typescript" }, { name = "typing-extensions" }, { name = "unidiff" }, + { name = "urllib3" }, { name = "uvicorn", extra = ["standard"] }, { name = "watchfiles" }, { name = "wrapt" }, @@ -731,11 +733,13 @@ requires-dist = [ { name = "types-toml", marker = "extra == 'types'", specifier = ">=0.10.8.20240310" }, { name = "typing-extensions", specifier = ">=4.12.2" }, { name = "unidiff", specifier = ">=0.7.5" }, + { name = "urllib3", specifier = ">=2.0.0" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.30.0" }, { name = "watchfiles", specifier = ">=1.0.0,<1.1.0" }, { name = "wrapt", specifier = ">=1.16.0,<2.0.0" }, { name = "xmltodict", specifier = ">=0.13.0,<1.0.0" }, ] +provides-extras = ["lsp", "types"] [package.metadata.requires-dev] dev = [