from mcp.server.fastmcp import FastMCP import httpx import os import logging import json import subprocess import asyncio import re # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger("mcp-analysis") mcp = FastMCP("analysis") API_BASE_URL = "https://analysis.googleapis.com/v1alpha" class AnalysisClient: def __init__(self, api_key: str): self.api_key = api_key self.base_url = API_BASE_URL self.headers = { "Content-Type": "application/json", "x-goog-api-key": self.api_key } async def _request(self, method: str, path: str, **kwargs): # Handle session_id prefixing if path.startswith("/sessions/sessions/"): path = path.replace("/sessions/sessions/", "/sessions/") url = f"{self.base_url}{path}" async with httpx.AsyncClient() as client: response = await client.request(method, url, headers=self.headers, **kwargs) if response.status_code not in (200, 204): logger.error(f"Jules API Error ({method} {path}): {response.status_code} - {response.text}") response.raise_for_status() return response.json() if response.text else {} async def list_sources(self): return await self._request("GET", "/sources") async def get_source(self, name: str): if not name.startswith("sources/"): name = f"sources/{name}" return await self._request("GET", f"/{name}") async def create_session(self, prompt, source, branch="main", title="API Session", automation="AUTO_CREATE_PR"): # Format source for API if it is just owner/repo if "/" in source and "sources/github/" not in source: source = f"sources/github/{source}" payload = { "prompt": prompt, "sourceContext": { "source": source, "githubRepoContext": {"startingBranch": branch} } } if title: payload["title"] = title if automation: payload["automationMode"] = automation logger.info(f"Creating session with payload: {json.dumps(payload)}") try: return await self._request("POST", "/sessions", json=payload) except Exception as e: err_msg = str(e).lower() if "maximum number of sessions" in err_msg or "max sessions reached" in err_msg or "400" in err_msg: logger.warning(f"Possible quota/session limit error detected: {e}. Triggering early key rotation.") force_rotate_key() return {"error": f"Session creation failed: {e}. API key has been rotated. Please retry."} logger.warning(f"REST API create_session failed: {e}. Falling back to CLI.") cli_source = source.replace("sources/github/", "") return await self.create_session_cli(prompt, cli_source, branch, title) async def create_session_cli(self, prompt, source, branch="main", title="API Session"): """Fallback to analysis CLI for session creation.""" logger.info(f"Using analysis CLI to create session for {source}") cmd = [ "analysis", "new", "--repo", source, prompt ] env = os.environ.copy() env["JULES_API_KEY"] = self.api_key npm_bin = subprocess.check_output(["npm", "config", "get", "prefix"]).decode().strip() + "/bin" env["PATH"] = f"{env.get('PATH', '')}:{npm_bin}" try: process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env ) stdout, stderr = await process.communicate() output = stdout.decode().strip() match = re.search(r'sessions/(\d+)', output) if match: s_id = match.group(0) return {"name": s_id, "id": match.group(1)} return {"output": output} except: return {"error": "CLI failed"} async def get_session(self, session_id): session_id = str(session_id).strip() if not session_id.startswith("sessions/"): session_id = f"sessions/{session_id}" return await self._request("GET", f"/{session_id}") async def list_sessions(self, page_size=30): return await self._request("GET", "/sessions", params={"pageSize": page_size}) async def send_message(self, session_id, prompt): session_id = str(session_id).strip() if not session_id.startswith("sessions/"): session_id = f"sessions/{session_id}" return await self._request("POST", f"/{session_id}:sendMessage", json={"prompt": prompt}) async def approve_plan(self, session_id): session_id = str(session_id).strip() if not session_id.startswith("sessions/"): session_id = f"sessions/{session_id}" return await self._request("POST", f"/{session_id}:approvePlan", json={}) async def get_activity(self, name: str): name = str(name).strip() return await self._request("GET", f"/{name}") async def list_activities(self, session_id, page_size=50): session_id = str(session_id).strip() if not session_id.startswith("sessions/"): session_id = f"sessions/{session_id}" return await self._request("GET", f"/{session_id}/activities", params={"pageSize": page_size}) # API Key Rotation State REQUEST_COUNT = 0 CURRENT_KEY_INDEX = 0 def force_rotate_key(): """Immediately switch to the next available API key and reset counter.""" global REQUEST_COUNT, CURRENT_KEY_INDEX keys = [os.environ.get("JULES_API_KEY"), os.environ.get("JULES_API_KEY_FALLBACK")] keys = [k for k in keys if k] if len(keys) > 1: CURRENT_KEY_INDEX = (CURRENT_KEY_INDEX + 1) % len(keys) REQUEST_COUNT = 0 logger.info(f"FORCE ROTATION: Switched to Jules API Key index {CURRENT_KEY_INDEX}") def get_rotated_api_key(): global REQUEST_COUNT, CURRENT_KEY_INDEX keys = [os.environ.get("JULES_API_KEY"), os.environ.get("JULES_API_KEY_FALLBACK")] # Filter out None or empty keys keys = [k for k in keys if k] if not keys: logger.error("No Jules API keys found in environment.") return None key = keys[CURRENT_KEY_INDEX] REQUEST_COUNT += 1 if REQUEST_COUNT >= 3: REQUEST_COUNT = 0 old_index = CURRENT_KEY_INDEX CURRENT_KEY_INDEX = (CURRENT_KEY_INDEX + 1) % len(keys) if len(keys) > 1 and old_index != CURRENT_KEY_INDEX: logger.info(f"Requests per key reached. Rotating Jules API Key from index {old_index} to {CURRENT_KEY_INDEX}") return key def get_client(): key = get_rotated_api_key() return AnalysisClient(key) @mcp.tool() async def list_sources(): """List all available sources (repositories).""" return await get_client().list_sources() @mcp.tool() async def get_source(name: str): """ Get details about a specific source. :param name: The source name (e.g., 'sources/github/owner/repo'). """ return await get_client().get_source(name) @mcp.tool() async def create_session( prompt: str, source: str, branch: str = "main", title: str = "API Session", automation_mode: str = "AUTO_CREATE_PR" ): """ Initialize a new autonomous coding session for a repository. Args: prompt: The implementation task or instructions for Jules. source: Source repository path (e.g., 'sources/github/owner/repo'). branch: The branch to use. title: Title for the session. automation_mode: Defaults to 'AUTO_CREATE_PR'. """ return await get_client().create_session(prompt, source, branch, title, automation_mode) @mcp.tool() async def get_session(session_id: str): """ Get the current status and details of a session. :param session_id: The session ID (e.g., 'sessions/123'). """ return await get_client().get_session(session_id) @mcp.tool() async def list_sessions(page_size: int = 30): """List recent Jules sessions.""" return await get_client().list_sessions(page_size) @mcp.tool() async def send_message(session_id: str, prompt: str): """ Send a message or feedback to an existing session. :param session_id: The session ID (e.g., 'sessions/123'). :param prompt: The message/instruction text. """ return await get_client().send_message(session_id, prompt) @mcp.tool() async def approve_plan(session_id: str): """ Approve the implementation plan for a session. :param session_id: The session ID. """ return await get_client().approve_plan(session_id) @mcp.tool() async def get_activity(name: str): """ Get details about a specific activity. :param name: Activity resource name (e.g., 'sessions/123/activities/456'). """ return await get_client().get_activity(name) @mcp.tool() async def list_activities(session_id: str, page_size: int = 50): """ List activities for a specific session. :param session_id: The session ID. :param page_size: Number of activities to return. """ return await get_client().list_activities(session_id, page_size) if __name__ == "__main__": mcp.run()