From 8c9bc67335ae9f913d03b39b9865c6316fe74fad Mon Sep 17 00:00:00 2001 From: aran-yogesh Date: Tue, 10 Feb 2026 11:22:52 -0800 Subject: [PATCH] refactor: extract shared utils from server.py into agent/utils/ package --- apps/agent/agent/middleware/__init__.py | 3 +- apps/agent/agent/middleware/open_pr.py | 25 ++- apps/agent/agent/middleware/post_to_linear.py | 3 +- apps/agent/agent/server.py | 182 +----------------- apps/agent/agent/utils/__init__.py | 0 apps/agent/agent/utils/github.py | 133 +++++++++++++ apps/agent/agent/utils/linear.py | 58 ++++++ apps/agent/agent/utils/sandbox_state.py | 8 + 8 files changed, 214 insertions(+), 198 deletions(-) create mode 100644 apps/agent/agent/utils/__init__.py create mode 100644 apps/agent/agent/utils/github.py create mode 100644 apps/agent/agent/utils/linear.py create mode 100644 apps/agent/agent/utils/sandbox_state.py diff --git a/apps/agent/agent/middleware/__init__.py b/apps/agent/agent/middleware/__init__.py index 0ce6ccde..250091d4 100644 --- a/apps/agent/agent/middleware/__init__.py +++ b/apps/agent/agent/middleware/__init__.py @@ -1,10 +1,9 @@ -from .check_message_queue import LinearNotifyState, check_message_queue_before_model +from .check_message_queue import check_message_queue_before_model from .open_pr import open_pr_if_needed from .post_to_linear import post_to_linear_after_model from .tool_error_handler import ToolErrorMiddleware __all__ = [ - "LinearNotifyState", "ToolErrorMiddleware", "check_message_queue_before_model", "open_pr_if_needed", diff --git a/apps/agent/agent/middleware/open_pr.py b/apps/agent/agent/middleware/open_pr.py index 93d40e80..0dd29489 100644 --- a/apps/agent/agent/middleware/open_pr.py +++ b/apps/agent/agent/middleware/open_pr.py @@ -17,6 +17,11 @@ from langchain.agents.middleware import AgentState, after_agent from langgraph.config import get_config from langgraph.runtime import Runtime +from ..encryption import decrypt_token +from ..utils.github import create_github_pr, get_github_default_branch +from ..utils.linear import comment_on_linear_issue +from ..utils.sandbox_state import SANDBOX_BACKENDS + logger = logging.getLogger(__name__) @@ -41,19 +46,11 @@ def _extract_pr_params_from_messages(messages: list) -> dict[str, str] | None: @after_agent -async def open_pr_if_needed( # noqa: PLR0912, PLR0915 +async def open_pr_if_needed( state: AgentState, - runtime: Runtime, # noqa: ARG001 + runtime: Runtime, ) -> dict[str, Any] | None: """Middleware that commits/pushes changes and comments on Linear after agent runs.""" - from ..encryption import decrypt_token - from ..server import ( - _SANDBOX_BACKENDS, - comment_on_linear_issue, - create_github_pr, - get_github_default_branch, - ) - logger.info("After-agent middleware started") pr_url = None pr_number = None @@ -103,7 +100,7 @@ async def open_pr_if_needed( # noqa: PLR0912, PLR0915 repo_owner = repo_config.get("owner") repo_name = repo_config.get("name") - sandbox_backend = _SANDBOX_BACKENDS.get(thread_id) + sandbox_backend = SANDBOX_BACKENDS.get(thread_id) repo_dir = f"/workspace/{repo_name}" @@ -214,7 +211,7 @@ async def open_pr_if_needed( # noqa: PLR0912, PLR0915 if linear_issue_id and last_message_content: if pr_url: - comment = f"""✅ **Pull Request Created** + comment = f"""**Pull Request Created** I've created a pull request to address this issue: @@ -226,7 +223,7 @@ I've created a pull request to address this issue: {last_message_content}""" else: - comment = f"""🤖 **Agent Response** + comment = f""" **Agent Response** {last_message_content}""" await comment_on_linear_issue(linear_issue_id, comment) @@ -241,7 +238,7 @@ I've created a pull request to address this issue: linear_issue = configurable.get("linear_issue", {}) linear_issue_id = linear_issue.get("id") if linear_issue_id: - error_comment = f"""❌ **Agent Error** + error_comment = f""" **Agent Error** An error occurred while processing this issue: diff --git a/apps/agent/agent/middleware/post_to_linear.py b/apps/agent/agent/middleware/post_to_linear.py index 52326ac9..c8a1c569 100644 --- a/apps/agent/agent/middleware/post_to_linear.py +++ b/apps/agent/agent/middleware/post_to_linear.py @@ -13,6 +13,7 @@ from langchain.agents.middleware import after_model from langgraph.config import get_config from langgraph.runtime import Runtime +from ..utils.linear import comment_on_linear_issue from .check_message_queue import LinearNotifyState logger = logging.getLogger(__name__) @@ -34,8 +35,6 @@ async def post_to_linear_after_model( # noqa: PLR0911, PLR0912 - The AI response has text content (not just tool calls) - The message hasn't already been sent (tracked via linear_messages_sent_count) """ - from ..server import comment_on_linear_issue - try: config = get_config() configurable = config.get("configurable", {}) diff --git a/apps/agent/agent/server.py b/apps/agent/agent/server.py index 1623020c..31dd5fe1 100644 --- a/apps/agent/agent/server.py +++ b/apps/agent/agent/server.py @@ -6,7 +6,6 @@ import logging import os import warnings -from typing import Any logger = logging.getLogger(__name__) @@ -95,184 +94,7 @@ SANDBOX_CREATING = "__creating__" SANDBOX_CREATION_TIMEOUT = 180 SANDBOX_POLL_INTERVAL = 1.0 -# HTTP status codes -HTTP_CREATED = 201 -HTTP_UNPROCESSABLE_ENTITY = 422 - -# Message count thresholds -_SANDBOX_BACKENDS: dict[str, Any] = {} - -import httpx - -LINEAR_API_KEY = os.environ.get("LINEAR_API_KEY", "") - - -async def create_github_pr( - repo_owner: str, - repo_name: str, - github_token: str, - title: str, - head_branch: str, - base_branch: str, - body: str, -) -> tuple[str | None, int | None]: - """Create a GitHub pull request via the API. - - Args: - repo_owner: Repository owner (e.g., "langchain-ai") - repo_name: Repository name (e.g., "deepagents") - github_token: GitHub access token - title: PR title - head_branch: Source branch name - base_branch: Target branch name - body: PR description - - Returns: - Tuple of (pr_url, pr_number) if successful, (None, None) otherwise - """ - pr_payload = { - "title": title, - "head": head_branch, - "base": base_branch, - "body": body, - } - - logger.info( - "Creating PR: head=%s, base=%s, repo=%s/%s", - head_branch, - base_branch, - repo_owner, - repo_name, - ) - - try: - async with httpx.AsyncClient() as http_client: - pr_response = await http_client.post( - f"https://api.github.com/repos/{repo_owner}/{repo_name}/pulls", - headers={ - "Authorization": f"Bearer {github_token}", - "Accept": "application/vnd.github+json", - "X-GitHub-Api-Version": "2022-11-28", - }, - json=pr_payload, - ) - - pr_data = pr_response.json() - - if pr_response.status_code == HTTP_CREATED: - pr_url = pr_data.get("html_url") - pr_number = pr_data.get("number") - logger.info("PR created successfully: %s", pr_url) - return pr_url, pr_number - - if pr_response.status_code == HTTP_UNPROCESSABLE_ENTITY: - logger.error("GitHub API validation error (422): %s", pr_data.get("message")) - else: - logger.error( - "GitHub API error (%s): %s", - pr_response.status_code, - pr_data.get("message"), - ) - - if "errors" in pr_data: - logger.error("GitHub API errors detail: %s", pr_data.get("errors")) - - return None, None - - except httpx.HTTPError: - logger.exception("Failed to create PR via GitHub API") - return None, None - - -async def get_github_default_branch( - repo_owner: str, - repo_name: str, - github_token: str, -) -> str: - """Get the default branch of a GitHub repository via the API. - - Args: - repo_owner: Repository owner (e.g., "langchain-ai") - repo_name: Repository name (e.g., "deepagents") - github_token: GitHub access token - - Returns: - The default branch name (e.g., "main" or "master") - """ - try: - async with httpx.AsyncClient() as http_client: - response = await http_client.get( - f"https://api.github.com/repos/{repo_owner}/{repo_name}", - headers={ - "Authorization": f"Bearer {github_token}", - "Accept": "application/vnd.github+json", - "X-GitHub-Api-Version": "2022-11-28", - }, - ) - - if response.status_code == 200: # noqa: PLR2004 - repo_data = response.json() - default_branch = repo_data.get("default_branch", "main") - logger.debug("Got default branch from GitHub API: %s", default_branch) - return default_branch - - logger.warning( - "Failed to get repo info from GitHub API (%s), falling back to 'main'", - response.status_code, - ) - return "main" - - except httpx.HTTPError: - logger.exception("Failed to get default branch from GitHub API, falling back to 'main'") - return "main" - - -async def comment_on_linear_issue(issue_id: str, comment_body: str) -> bool: - """Add a comment to a Linear issue. - - Args: - issue_id: The Linear issue ID - comment_body: The comment text - - Returns: - True if successful, False otherwise - """ - if not LINEAR_API_KEY: - return False - - import httpx - - url = "https://api.linear.app/graphql" - - mutation = """ - mutation CommentCreate($issueId: String!, $body: String!) { - commentCreate(input: { issueId: $issueId, body: $body }) { - success - comment { - id - } - } - } - """ - - async with httpx.AsyncClient() as http_client: - try: - response = await http_client.post( - url, - headers={ - "Authorization": LINEAR_API_KEY, - "Content-Type": "application/json", - }, - json={ - "query": mutation, - "variables": {"issueId": issue_id, "body": comment_body}, - }, - ) - response.raise_for_status() - result = response.json() - return bool(result.get("data", {}).get("commentCreate", {}).get("success")) - except Exception: # noqa: BLE001 - return False +from .utils.sandbox_state import SANDBOX_BACKENDS async def _clone_or_pull_repo_in_sandbox( # noqa: PLR0915 @@ -532,7 +354,7 @@ async def get_agent(config: RunnableConfig) -> Pregel: # noqa: PLR0915 logger.exception("Failed to pull repo in existing sandbox") raise - _SANDBOX_BACKENDS[thread_id] = sandbox_backend + SANDBOX_BACKENDS[thread_id] = sandbox_backend linear_issue = config["configurable"].get("linear_issue", {}) linear_project_id = linear_issue.get("linear_project_id", "") diff --git a/apps/agent/agent/utils/__init__.py b/apps/agent/agent/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/apps/agent/agent/utils/github.py b/apps/agent/agent/utils/github.py new file mode 100644 index 00000000..216cf581 --- /dev/null +++ b/apps/agent/agent/utils/github.py @@ -0,0 +1,133 @@ +"""GitHub API utilities.""" + +from __future__ import annotations + +import logging + +import httpx + +logger = logging.getLogger(__name__) + +# HTTP status codes +HTTP_CREATED = 201 +HTTP_UNPROCESSABLE_ENTITY = 422 + + +async def create_github_pr( + repo_owner: str, + repo_name: str, + github_token: str, + title: str, + head_branch: str, + base_branch: str, + body: str, +) -> tuple[str | None, int | None]: + """Create a GitHub pull request via the API. + + Args: + repo_owner: Repository owner (e.g., "langchain-ai") + repo_name: Repository name (e.g., "deepagents") + github_token: GitHub access token + title: PR title + head_branch: Source branch name + base_branch: Target branch name + body: PR description + + Returns: + Tuple of (pr_url, pr_number) if successful, (None, None) otherwise + """ + pr_payload = { + "title": title, + "head": head_branch, + "base": base_branch, + "body": body, + } + + logger.info( + "Creating PR: head=%s, base=%s, repo=%s/%s", + head_branch, + base_branch, + repo_owner, + repo_name, + ) + + try: + async with httpx.AsyncClient() as http_client: + pr_response = await http_client.post( + f"https://api.github.com/repos/{repo_owner}/{repo_name}/pulls", + headers={ + "Authorization": f"Bearer {github_token}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }, + json=pr_payload, + ) + + pr_data = pr_response.json() + + if pr_response.status_code == HTTP_CREATED: + pr_url = pr_data.get("html_url") + pr_number = pr_data.get("number") + logger.info("PR created successfully: %s", pr_url) + return pr_url, pr_number + + if pr_response.status_code == HTTP_UNPROCESSABLE_ENTITY: + logger.error("GitHub API validation error (422): %s", pr_data.get("message")) + else: + logger.error( + "GitHub API error (%s): %s", + pr_response.status_code, + pr_data.get("message"), + ) + + if "errors" in pr_data: + logger.error("GitHub API errors detail: %s", pr_data.get("errors")) + + return None, None + + except httpx.HTTPError: + logger.exception("Failed to create PR via GitHub API") + return None, None + + +async def get_github_default_branch( + repo_owner: str, + repo_name: str, + github_token: str, +) -> str: + """Get the default branch of a GitHub repository via the API. + + Args: + repo_owner: Repository owner (e.g., "langchain-ai") + repo_name: Repository name (e.g., "deepagents") + github_token: GitHub access token + + Returns: + The default branch name (e.g., "main" or "master") + """ + try: + async with httpx.AsyncClient() as http_client: + response = await http_client.get( + f"https://api.github.com/repos/{repo_owner}/{repo_name}", + headers={ + "Authorization": f"Bearer {github_token}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + + if response.status_code == 200: # noqa: PLR2004 + repo_data = response.json() + default_branch = repo_data.get("default_branch", "main") + logger.debug("Got default branch from GitHub API: %s", default_branch) + return default_branch + + logger.warning( + "Failed to get repo info from GitHub API (%s), falling back to 'main'", + response.status_code, + ) + return "main" + + except httpx.HTTPError: + logger.exception("Failed to get default branch from GitHub API, falling back to 'main'") + return "main" diff --git a/apps/agent/agent/utils/linear.py b/apps/agent/agent/utils/linear.py new file mode 100644 index 00000000..96bb9cd6 --- /dev/null +++ b/apps/agent/agent/utils/linear.py @@ -0,0 +1,58 @@ +"""Linear API utilities.""" + +from __future__ import annotations + +import logging +import os + +import httpx + +logger = logging.getLogger(__name__) + +LINEAR_API_KEY = os.environ.get("LINEAR_API_KEY", "") + + +async def comment_on_linear_issue(issue_id: str, comment_body: str) -> bool: + """Add a comment to a Linear issue. + + Args: + issue_id: The Linear issue ID + comment_body: The comment text + + Returns: + True if successful, False otherwise + """ + if not LINEAR_API_KEY: + return False + + url = "https://api.linear.app/graphql" + + mutation = """ + mutation CommentCreate($issueId: String!, $body: String!) { + commentCreate(input: { issueId: $issueId, body: $body }) { + success + comment { + id + } + } + } + """ + + async with httpx.AsyncClient() as http_client: + try: + response = await http_client.post( + url, + headers={ + "Authorization": LINEAR_API_KEY, + "Content-Type": "application/json", + }, + json={ + "query": mutation, + "variables": {"issueId": issue_id, "body": comment_body}, + }, + ) + response.raise_for_status() + result = response.json() + return bool(result.get("data", {}).get("commentCreate", {}).get("success")) + except Exception: # noqa: BLE001 + return False diff --git a/apps/agent/agent/utils/sandbox_state.py b/apps/agent/agent/utils/sandbox_state.py new file mode 100644 index 00000000..5d8d2b9c --- /dev/null +++ b/apps/agent/agent/utils/sandbox_state.py @@ -0,0 +1,8 @@ +"""Shared sandbox state used by server and middleware.""" + +from __future__ import annotations + +from typing import Any + +# Thread ID -> SandboxBackend mapping, shared between server.py and middleware +SANDBOX_BACKENDS: dict[str, Any] = {}