diff --git a/roboco/mcp/journal_server.py b/roboco/mcp/journal_server.py index 447f0629..81a8c61a 100644 --- a/roboco/mcp/journal_server.py +++ b/roboco/mcp/journal_server.py @@ -16,12 +16,8 @@ Tools: from typing import Any -import httpx -from fastapi import status from mcp.server.fastmcp import FastMCP -from roboco.agents_config import get_agent_role, get_agent_team -from roboco.config import settings from roboco.llm import ToonAdapter from roboco.mcp.schemas import ( DecisionLogInput, @@ -30,74 +26,15 @@ from roboco.mcp.schemas import ( StruggleInput, TaskReflectionInput, ) +from roboco.mcp.utils import ApiClient, format_error_response # Global TOON adapter for encoding journal data _toon = ToonAdapter() - -# ============================================================================= -# HELPER FUNCTIONS -# ============================================================================= - - -def _format_error_response( - error_code: str, - message: str, - details: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Format a standardized error response.""" - return { - "error": { - "code": error_code, - "message": message, - "details": details or {}, - } - } - - -def _get_agent_headers(agent_id: str) -> dict[str, str]: - """Get standard headers for API calls.""" - headers = { - "X-Agent-ID": agent_id, - "X-Agent-Role": get_agent_role(agent_id), - } - team = get_agent_team(agent_id) - if team: - headers["X-Agent-Team"] = team - return headers - - -async def _post_journal_entry( - endpoint: str, - payload: dict[str, Any], - agent_id: str, -) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: - """Post to a journal endpoint. Returns (data, error).""" - async with httpx.AsyncClient(timeout=30.0) as client: - try: - resp = await client.post( - f"{settings.internal_api_url}/journals/me/{endpoint}", - json=payload, - headers=_get_agent_headers(agent_id), - ) - except httpx.TimeoutException: - return None, _format_error_response( - "TIMEOUT", - f"Request to create {endpoint.rstrip('s')} timed out", - ) - except httpx.RequestError as e: - return None, _format_error_response( - "CONNECTION_ERROR", - f"Failed to connect to API: {type(e).__name__}", - ) - - if resp.status_code not in [200, 201]: - return None, _format_error_response( - "CREATE_FAILED", - f"Failed to create {endpoint.rstrip('s')}", - {"status_code": resp.status_code, "detail": resp.text}, - ) - return resp.json(), None +# Valid entry types +VALID_ENTRY_TYPES = frozenset( + ["general", "task_reflection", "decision_log", "learning", "struggle"] +) # ============================================================================= @@ -106,14 +43,13 @@ async def _post_journal_entry( async def _handle_journal_entry( - data: JournalEntryInput, agent_id: str + data: JournalEntryInput, client: ApiClient ) -> dict[str, Any]: """Handle journal entry creation.""" - valid_types = ["general", "task_reflection", "decision_log", "learning", "struggle"] - if data.entry_type not in valid_types: - return _format_error_response( + if data.entry_type not in VALID_ENTRY_TYPES: + return format_error_response( "INVALID_TYPE", - f"Invalid entry type. Must be one of: {valid_types}", + f"Invalid entry type. Must be one of: {list(VALID_ENTRY_TYPES)}", ) payload = { @@ -125,9 +61,14 @@ async def _handle_journal_entry( "is_private": data.is_private, } - entry, error = await _post_journal_entry("entries", payload, agent_id) + entry, error = await client.post_or_error( + "/journals/me/entries", + json=payload, + error_code="CREATE_FAILED", + error_message="Failed to create entry", + ) if error or entry is None: - return error or _format_error_response("ERROR", "Failed to create entry") + return error or format_error_response("ERROR", "Failed to create entry") return { "status": "created", @@ -139,7 +80,9 @@ async def _handle_journal_entry( } -async def _handle_reflect(data: TaskReflectionInput, agent_id: str) -> dict[str, Any]: +async def _handle_reflect( + data: TaskReflectionInput, client: ApiClient +) -> dict[str, Any]: """Handle task reflection creation.""" payload = { "task_id": data.task_id, @@ -151,7 +94,12 @@ async def _handle_reflect(data: TaskReflectionInput, agent_id: str) -> dict[str, "tags": data.tags, } - entry, error = await _post_journal_entry("reflections", payload, agent_id) + entry, error = await client.post_or_error( + "/journals/me/reflections", + json=payload, + error_code="CREATE_FAILED", + error_message="Failed to create reflection", + ) if error: return error @@ -165,7 +113,7 @@ async def _handle_reflect(data: TaskReflectionInput, agent_id: str) -> dict[str, } -async def _handle_decision(data: DecisionLogInput, agent_id: str) -> dict[str, Any]: +async def _handle_decision(data: DecisionLogInput, client: ApiClient) -> dict[str, Any]: """Handle decision log creation.""" payload = { "title": data.title, @@ -178,7 +126,12 @@ async def _handle_decision(data: DecisionLogInput, agent_id: str) -> dict[str, A "tags": data.tags, } - entry, error = await _post_journal_entry("decisions", payload, agent_id) + entry, error = await client.post_or_error( + "/journals/me/decisions", + json=payload, + error_code="CREATE_FAILED", + error_message="Failed to create decision log", + ) if error: return error @@ -192,7 +145,7 @@ async def _handle_decision(data: DecisionLogInput, agent_id: str) -> dict[str, A } -async def _handle_learning(data: LearningInput, agent_id: str) -> dict[str, Any]: +async def _handle_learning(data: LearningInput, client: ApiClient) -> dict[str, Any]: """Handle learning entry creation.""" payload = { "title": data.title, @@ -203,7 +156,12 @@ async def _handle_learning(data: LearningInput, agent_id: str) -> dict[str, Any] "tags": data.tags, } - entry, error = await _post_journal_entry("learnings", payload, agent_id) + entry, error = await client.post_or_error( + "/journals/me/learnings", + json=payload, + error_code="CREATE_FAILED", + error_message="Failed to create learning entry", + ) if error: return error @@ -214,7 +172,7 @@ async def _handle_learning(data: LearningInput, agent_id: str) -> dict[str, Any] } -async def _handle_struggle(data: StruggleInput, agent_id: str) -> dict[str, Any]: +async def _handle_struggle(data: StruggleInput, client: ApiClient) -> dict[str, Any]: """Handle struggle entry creation.""" payload = { "title": data.title, @@ -226,7 +184,12 @@ async def _handle_struggle(data: StruggleInput, agent_id: str) -> dict[str, Any] "tags": data.tags, } - entry, error = await _post_journal_entry("struggles", payload, agent_id) + entry, error = await client.post_or_error( + "/journals/me/struggles", + json=payload, + error_code="CREATE_FAILED", + error_message="Failed to create struggle entry", + ) if error: return error @@ -239,34 +202,19 @@ async def _handle_struggle(data: StruggleInput, agent_id: str) -> dict[str, Any] return {"status": "created", "entry": entry, "guidance": guidance} -async def _handle_search(query: str, top_k: int, agent_id: str) -> dict[str, Any]: +async def _handle_search(query: str, top_k: int, client: ApiClient) -> dict[str, Any]: """Handle journal search.""" - async with httpx.AsyncClient(timeout=30.0) as client: - payload = {"query": query, "top_k": min(top_k, 20)} - try: - resp = await client.post( - f"{settings.internal_api_url}/journals/me/search", - json=payload, - headers=_get_agent_headers(agent_id), - ) - except httpx.TimeoutException: - return _format_error_response( - "TIMEOUT", "Search request timed out" - ) - except httpx.RequestError as e: - return _format_error_response( - "CONNECTION_ERROR", - f"Failed to connect to API: {type(e).__name__}", - ) + max_results = 20 + payload = {"query": query, "top_k": min(top_k, max_results)} - if resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "SEARCH_FAILED", - "Failed to search journal", - {"status_code": resp.status_code, "detail": resp.text}, - ) - - entries = resp.json() + entries, error = await client.post_or_error( + "/journals/me/search", + json=payload, + error_code="SEARCH_FAILED", + error_message="Failed to search journal", + ) + if error: + return error if not entries: return { @@ -281,33 +229,14 @@ async def _handle_search(query: str, top_k: int, agent_id: str) -> dict[str, Any } -async def _handle_stats(agent_id: str) -> dict[str, Any]: +async def _handle_stats(client: ApiClient) -> dict[str, Any]: """Handle journal stats retrieval.""" - async with httpx.AsyncClient(timeout=30.0) as client: - try: - stats_resp = await client.get( - f"{settings.internal_api_url}/journals/me/stats", - headers=_get_agent_headers(agent_id), - ) - growth_resp = await client.get( - f"{settings.internal_api_url}/journals/me/growth", - headers=_get_agent_headers(agent_id), - ) - except httpx.RequestError: - # Return empty stats on connection error - stats_resp = None - growth_resp = None + # Fetch stats and growth in parallel would be better but keep simple for now + stats_resp = await client.get("/journals/me/stats") + growth_resp = await client.get("/journals/me/growth") - stats = ( - stats_resp.json() - if stats_resp and stats_resp.status_code == status.HTTP_200_OK - else {} - ) - growth = ( - growth_resp.json() - if growth_resp and growth_resp.status_code == status.HTTP_200_OK - else {} - ) + stats = stats_resp.json() if stats_resp.ok else {} + growth = growth_resp.json() if growth_resp.ok else {} return { "total_entries": stats.get("total_entries", 0), @@ -332,42 +261,26 @@ async def _handle_recent( entry_type: str | None, task_id: str | None, limit: int, - agent_id: str, + client: ApiClient, ) -> dict[str, Any]: """Handle recent entries retrieval.""" - async with httpx.AsyncClient(timeout=30.0) as client: - params: dict[str, Any] = {"limit": min(limit, 50)} - if entry_type: - params["entry_type"] = entry_type - if task_id: - params["task_id"] = task_id + max_limit = 50 + params: dict[str, Any] = {"limit": min(limit, max_limit)} + if entry_type: + params["entry_type"] = entry_type + if task_id: + params["task_id"] = task_id - try: - resp = await client.get( - f"{settings.internal_api_url}/journals/me/entries", - params=params, - headers=_get_agent_headers(agent_id), - ) - except httpx.TimeoutException: - return _format_error_response( - "TIMEOUT", "Request to list entries timed out" - ) - except httpx.RequestError as e: - return _format_error_response( - "CONNECTION_ERROR", - f"Failed to connect to API: {type(e).__name__}", - ) + entries, error = await client.get_or_error( + "/journals/me/entries", + params=params, + error_code="LIST_FAILED", + error_message="Failed to list entries", + ) + if error: + return error - if resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "LIST_FAILED", - "Failed to list entries", - {"status_code": resp.status_code, "detail": resp.text}, - ) - - entries = resp.json() - - return {"entries": entries, "count": len(entries)} + return {"entries": entries, "count": len(entries) if entries else 0} # ============================================================================= @@ -387,6 +300,9 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: """ mcp = FastMCP(f"roboco-journal-{agent_id}", json_response=True) + # Create shared API client for this agent + client = ApiClient(agent_id) + @mcp.tool() async def roboco_journal_entry(data: JournalEntryInput) -> dict[str, Any]: """ @@ -395,7 +311,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Your journal is personal - use it to track thoughts, progress, and document your journey on tasks. """ - return await _handle_journal_entry(data, agent_id) + return await _handle_journal_entry(data, client) @mcp.tool() async def roboco_journal_reflect(data: TaskReflectionInput) -> dict[str, Any]: @@ -405,7 +321,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: IMPORTANT: Call this when completing a task. Reflections help build institutional memory and track your growth. """ - return await _handle_reflect(data, agent_id) + return await _handle_reflect(data, client) @mcp.tool() async def roboco_journal_decision(data: DecisionLogInput) -> dict[str, Any]: @@ -415,7 +331,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Use when choosing between approaches. Creates a record of WHY you made the decision for future context. """ - return await _handle_decision(data, agent_id) + return await _handle_decision(data, client) @mcp.tool() async def roboco_journal_learning(data: LearningInput) -> dict[str, Any]: @@ -424,7 +340,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Track learnings to build your knowledge base and help future you. """ - return await _handle_learning(data, agent_id) + return await _handle_learning(data, client) @mcp.tool() async def roboco_journal_struggle(data: StruggleInput) -> dict[str, Any]: @@ -434,7 +350,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Recording struggles helps track problem-solving patterns and create documentation for others. """ - return await _handle_struggle(data, agent_id) + return await _handle_struggle(data, client) @mcp.tool() async def roboco_journal_search(query: str, top_k: int = 5) -> dict[str, Any]: @@ -443,7 +359,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Uses semantic search to find relevant entries based on meaning. """ - return await _handle_search(query, top_k, agent_id) + return await _handle_search(query, top_k, client) @mcp.tool() async def roboco_journal_stats() -> dict[str, Any]: @@ -452,7 +368,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Returns counts by entry type, growth metrics, and other stats. """ - return await _handle_stats(agent_id) + return await _handle_stats(client) @mcp.tool() async def roboco_journal_recent( @@ -466,7 +382,7 @@ def create_journal_mcp_server(agent_id: str) -> FastMCP: Filter by entry_type (general, task_reflection, decision_log, learning, struggle) or by task_id. """ - return await _handle_recent(entry_type, task_id, limit, agent_id) + return await _handle_recent(entry_type, task_id, limit, client) return mcp diff --git a/roboco/mcp/message_server.py b/roboco/mcp/message_server.py index 917d05dc..2391ddef 100644 --- a/roboco/mcp/message_server.py +++ b/roboco/mcp/message_server.py @@ -12,73 +12,36 @@ Tools: - roboco_channel_history: Get channel message history """ -from collections.abc import Awaitable, Callable from datetime import UTC, datetime, timedelta from typing import Any -import httpx from fastapi import status from mcp.server.fastmcp import FastMCP -from roboco.agents_config import CHANNEL_ACCESS, get_agent_role, get_agent_team -from roboco.config import settings +from roboco.agents_config import CHANNEL_ACCESS from roboco.llm import ToonAdapter from roboco.mcp.schemas import ( AskQuestionInput, ReportBlockerInput, SendMessageInput, ) +from roboco.mcp.utils import ApiClient, format_error_response, resolve_agent_uuid # Global TOON adapter for encoding message data _toon = ToonAdapter() - -def _get_agent_headers(agent_id: str) -> dict[str, str]: - """Get standard headers for API calls.""" - headers = { - "X-Agent-ID": agent_id, - "X-Agent-Role": get_agent_role(agent_id), - } - team = get_agent_team(agent_id) - if team: - headers["X-Agent-Team"] = team - return headers - - # Cache for agent slug -> UUID resolution _agent_uuid_cache: dict[str, str] = {} -async def _resolve_agent_uuid(agent_id: str, headers: dict[str, str]) -> str | None: - """Resolve agent slug to UUID. Returns None if not found.""" - from uuid import UUID - - # Check if already a valid UUID - try: - UUID(agent_id) - return agent_id # Already a UUID - except ValueError: - pass - - # Check cache +async def _resolve_agent_uuid_cached(agent_id: str, client: ApiClient) -> str | None: + """Resolve agent slug to UUID with caching. Returns None if not found.""" if agent_id in _agent_uuid_cache: return _agent_uuid_cache[agent_id] - - # Query API to resolve slug to UUID - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/agents", - params={"slug": agent_id}, - headers=headers, - ) - if resp.status_code == status.HTTP_200_OK: - agents = resp.json() - if agents: - uuid: str | None = agents[0].get("id") - if uuid: - _agent_uuid_cache[agent_id] = uuid - return uuid - return None + result = await resolve_agent_uuid(agent_id, client._get_headers()) + if result: + _agent_uuid_cache[agent_id] = result + return result # ============================================================================= @@ -100,21 +63,6 @@ def _check_channel_access(agent_id: str, channel_slug: str, action: str) -> bool return bool(action == "read" and agent_id in channel.get("silent", [])) -def _format_error_response( - error_code: str, - message: str, - details: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Format a standardized error response.""" - return { - "error": { - "code": error_code, - "message": message, - "details": details or {}, - } - } - - def _validate_message_send( agent_id: str, channel_slug: str, @@ -131,7 +79,7 @@ def _validate_message_send( "technical", ] if message_type not in valid_types: - return _format_error_response( + return format_error_response( "INVALID_TYPE", f"Invalid message type '{message_type}'. Must be one of: {valid_types}", ) @@ -140,20 +88,20 @@ def _validate_message_send( writable = [ ch for ch in CHANNEL_ACCESS if _check_channel_access(agent_id, ch, "write") ] - return _format_error_response( + return format_error_response( "ACCESS_DENIED", f"You don't have write access to #{channel_slug}", {"your_writable_channels": writable}, ) if agent_id in CHANNEL_ACCESS.get(channel_slug, {}).get("silent", []): - return _format_error_response( + return format_error_response( "SILENT_OBSERVER", "You are a silent observer on this channel and cannot post messages.", ) if not content or not content.strip(): - return _format_error_response( + return format_error_response( "EMPTY_CONTENT", "Message content cannot be empty." ) @@ -161,26 +109,22 @@ def _validate_message_send( async def _get_default_group( - client: httpx.AsyncClient, + client: ApiClient, channel_id: str, - headers: dict[str, str], ) -> str | dict[str, Any]: """Get the default (first) group for a channel. Returns group_id or error dict.""" - groups_resp = await client.get( - f"{settings.internal_api_url}/channels/{channel_id}/groups", - headers=headers, - ) + resp = await client.get(f"/channels/{channel_id}/groups") - if groups_resp.status_code != status.HTTP_200_OK: - return _format_error_response( + if not resp.ok: + return format_error_response( "GROUPS_ERROR", "Failed to get channel groups", - {"status": groups_resp.status_code}, + {"status": resp.status_code}, ) - groups = groups_resp.json() + groups = resp.json() if not groups: - return _format_error_response("NO_GROUPS", "Channel has no groups") + return format_error_response("NO_GROUPS", "Channel has no groups") # Return first active group, or first group if none are active for group in groups: @@ -190,26 +134,21 @@ async def _get_default_group( async def _get_or_create_session( - client: httpx.AsyncClient, + client: ApiClient, channel_id: str, - headers: dict[str, str], ) -> str | dict[str, Any]: """Get or create session for channel. Returns session_id or error dict.""" # First get the default group for this channel - group_result = await _get_default_group(client, channel_id, headers) + group_result = await _get_default_group(client, channel_id) if isinstance(group_result, dict): return group_result # Error response group_id = group_result # Check if group has an active session - sessions_resp = await client.get( - f"{settings.internal_api_url}/sessions", - params={"group_id": group_id, "limit": 1}, - headers=headers, - ) + resp = await client.get("/sessions", params={"group_id": group_id, "limit": 1}) - if sessions_resp.status_code == status.HTTP_200_OK: - data = sessions_resp.json() + if resp.ok: + data = resp.json() items = data.get("items", []) # Find an active session for session in items: @@ -217,15 +156,11 @@ async def _get_or_create_session( return str(session["id"]) # Create new session - create_resp = await client.post( - f"{settings.internal_api_url}/sessions", - json={"group_id": group_id}, - headers=headers, - ) - if create_resp.status_code in [status.HTTP_200_OK, status.HTTP_201_CREATED]: + create_resp = await client.post("/sessions", json={"group_id": group_id}) + if create_resp.ok: return str(create_resp.json()["id"]) - return _format_error_response( + return format_error_response( "SESSION_ERROR", "Failed to create session", {"api_error": create_resp.text}, @@ -263,68 +198,56 @@ async def _handle_channel_list(agent_id: str) -> dict[str, Any]: async def _get_channel_by_slug( - client: httpx.AsyncClient, + client: ApiClient, channel_slug: str, - headers: dict[str, str], ) -> str | dict[str, Any]: """Get channel ID by slug. Returns channel_id or error dict.""" - resp = await client.get( - f"{settings.internal_api_url}/channels", - params={"slug": channel_slug}, - headers=headers, - ) + resp = await client.get("/channels", params={"slug": channel_slug}) - if resp.status_code != status.HTTP_200_OK: - return _format_error_response("API_ERROR", "Failed to fetch channels") + if not resp.ok: + return format_error_response("API_ERROR", "Failed to fetch channels") data = resp.json() items = data.get("items", data) if not items: - return _format_error_response("NOT_FOUND", f"Channel #{channel_slug} not found") + return format_error_response("NOT_FOUND", f"Channel #{channel_slug} not found") channel = items[0] if isinstance(items, list) else items return str(channel["id"]) async def _get_sessions_for_group( - client: httpx.AsyncClient, + client: ApiClient, group_id: str, - headers: dict[str, str], ) -> list | dict[str, Any]: """Get sessions for a group. Returns session list or error dict.""" - resp = await client.get( - f"{settings.internal_api_url}/sessions", - params={"group_id": group_id, "limit": 5}, - headers=headers, - ) + resp = await client.get("/sessions", params={"group_id": group_id, "limit": 5}) - if resp.status_code != status.HTTP_200_OK: - return _format_error_response("API_ERROR", "Failed to fetch sessions") + if not resp.ok: + return format_error_response("API_ERROR", "Failed to fetch sessions") items: list = resp.json().get("items", []) return items async def _fetch_messages_from_sessions( - client: httpx.AsyncClient, + client: ApiClient, sessions: list, since: datetime, limit: int, - headers: dict[str, str], ) -> list: """Fetch messages from multiple sessions.""" all_messages: list = [] for session in sessions: resp = await client.get( - f"{settings.internal_api_url}/messages", + "/messages", params={ "session_id": session["id"], "after": since.isoformat(), "limit": limit, }, - headers=headers, ) - if resp.status_code == status.HTTP_200_OK: + if resp.ok: all_messages.extend(resp.json().get("items", [])) if len(all_messages) >= limit: break @@ -334,6 +257,7 @@ async def _fetch_messages_from_sessions( async def _handle_channel_history( + client: ApiClient, agent_id: str, channel_slug: str, limit: int, @@ -341,47 +265,44 @@ async def _handle_channel_history( ) -> dict[str, Any]: """Handle channel history retrieval.""" if not _check_channel_access(agent_id, channel_slug, "read"): - return _format_error_response( + return format_error_response( "ACCESS_DENIED", f"You don't have read access to #{channel_slug}" ) - limit = min(limit, 100) + max_limit = 100 + limit = min(limit, max_limit) since = datetime.now(UTC) - timedelta(hours=hours_back) - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - # Get channel - channel_result = await _get_channel_by_slug(client, channel_slug, headers) - if isinstance(channel_result, dict): - return channel_result - channel_id = channel_result + # Get channel + channel_result = await _get_channel_by_slug(client, channel_slug) + if isinstance(channel_result, dict): + return channel_result + channel_id = channel_result - # Get group - group_result = await _get_default_group(client, channel_id, headers) - if isinstance(group_result, dict): - return group_result - group_id = group_result + # Get group + group_result = await _get_default_group(client, channel_id) + if isinstance(group_result, dict): + return group_result + group_id = group_result - # Get sessions - sessions_result = await _get_sessions_for_group(client, group_id, headers) - if isinstance(sessions_result, dict): - return sessions_result - sessions = sessions_result + # Get sessions + sessions_result = await _get_sessions_for_group(client, group_id) + if isinstance(sessions_result, dict): + return sessions_result + sessions = sessions_result - # Early return for no sessions - if not sessions: - return { - "channel": channel_slug, - "messages": [], - "total": 0, - "has_more": False, - "since": since.isoformat(), - } + # Early return for no sessions + if not sessions: + return { + "channel": channel_slug, + "messages": [], + "total": 0, + "has_more": False, + "since": since.isoformat(), + } - # Fetch messages - messages = await _fetch_messages_from_sessions( - client, sessions, since, limit, headers - ) + # Fetch messages + messages = await _fetch_messages_from_sessions(client, sessions, since, limit) return { "channel": channel_slug, @@ -393,6 +314,7 @@ async def _handle_channel_history( async def _handle_message_send( + client: ApiClient, agent_id: str, data: SendMessageInput, ) -> dict[str, Any]: @@ -402,80 +324,68 @@ async def _handle_message_send( ): return validation_error - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - # Use the helper function that properly handles paginated responses - channel_result = await _get_channel_by_slug(client, data.channel_slug, headers) - if isinstance(channel_result, dict): - return channel_result # Error response - channel_id = channel_result + # Get channel by slug + channel_result = await _get_channel_by_slug(client, data.channel_slug) + if isinstance(channel_result, dict): + return channel_result # Error response + channel_id = channel_result - session_result = await _get_or_create_session(client, channel_id, headers) - if isinstance(session_result, dict): - return session_result - session_id = session_result + session_result = await _get_or_create_session(client, channel_id) + if isinstance(session_result, dict): + return session_result + session_id = session_result - # Resolve mentions (slugs) to UUIDs - resolved_mentions: list[str] = [] - if data.mentions: - for mention in data.mentions: - resolved = await _resolve_agent_uuid(mention, headers) - if resolved: - resolved_mentions.append(resolved) - # Skip unresolved mentions rather than failing + # Resolve mentions (slugs) to UUIDs + resolved_mentions: list[str] = [] + if data.mentions: + for mention in data.mentions: + resolved = await _resolve_agent_uuid_cached(mention, client) + if resolved: + resolved_mentions.append(resolved) + # Skip unresolved mentions rather than failing - message_data = { - "session_id": session_id, - "type": data.message_type, - "content": data.content, - "is_reply": data.reply_to is not None, - "reply_to": data.reply_to, - "mentions": resolved_mentions, - "task_id": data.task_id, - } + message_data = { + "session_id": session_id, + "type": data.message_type, + "content": data.content, + "is_reply": data.reply_to is not None, + "reply_to": data.reply_to, + "mentions": resolved_mentions, + "task_id": data.task_id, + } - send_resp = await client.post( - f"{settings.internal_api_url}/messages", - json=message_data, - headers=headers, + resp = await client.post("/messages", json=message_data) + + if not resp.ok: + return format_error_response( + "SEND_FAILED", "Failed to send message", {"api_error": resp.text} ) - if send_resp.status_code not in [status.HTTP_200_OK, status.HTTP_201_CREATED]: - return _format_error_response( - "SEND_FAILED", "Failed to send message", {"api_error": send_resp.text} - ) - - return { - "status": "sent", - "message": send_resp.json(), - "channel": data.channel_slug, - "guidance": "Message sent successfully.", - } + return { + "status": "sent", + "message": resp.json(), + "channel": data.channel_slug, + "guidance": "Message sent successfully.", + } -async def _handle_message_get(message_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_message_get(client: ApiClient, message_id: str) -> dict[str, Any]: """Handle message retrieval.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/messages/{message_id}", - headers=headers, - ) + resp = await client.get(f"/messages/{message_id}") - if resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response( - "NOT_FOUND", f"Message {message_id} not found" - ) + if resp.is_status(status.HTTP_404_NOT_FOUND): + return format_error_response("NOT_FOUND", f"Message {message_id} not found") - if resp.status_code != status.HTTP_200_OK: - return _format_error_response("API_ERROR", "Failed to fetch message") + if not resp.ok: + return format_error_response("API_ERROR", "Failed to fetch message") - return {"message": resp.json()} + return {"message": resp.json()} async def _handle_ask_question( + client: ApiClient, + agent_id: str, data: AskQuestionInput, - send_fn: Callable[[SendMessageInput], Awaitable[dict[str, Any]]], ) -> dict[str, Any]: """Handle asking a question.""" content = f"**Question**: {data.question}" @@ -488,7 +398,7 @@ async def _handle_ask_question( message_type="dialogue", task_id=data.task_id, ) - result = await send_fn(msg_data) + result = await _handle_message_send(client, agent_id, msg_data) if "error" in result: return result @@ -503,8 +413,9 @@ async def _handle_ask_question( async def _handle_report_blocker( + client: ApiClient, + agent_id: str, data: ReportBlockerInput, - send_fn: Callable[[SendMessageInput], Awaitable[dict[str, Any]]], ) -> dict[str, Any]: """Handle reporting a blocker.""" content = ( @@ -519,7 +430,7 @@ async def _handle_report_blocker( message_type="blocker", task_id=data.task_id, ) - result = await send_fn(msg_data) + result = await _handle_message_send(client, agent_id, msg_data) if "error" in result: return result @@ -550,6 +461,9 @@ def create_message_mcp_server(agent_id: str) -> FastMCP: """ mcp = FastMCP(f"roboco-message-{agent_id}", json_response=True) + # Create shared API client for this agent + client = ApiClient(agent_id) + @mcp.tool() async def roboco_channel_list() -> dict[str, Any]: """List channels you have access to.""" @@ -566,7 +480,9 @@ def create_message_mcp_server(agent_id: str) -> FastMCP: You must have read access to the channel. """ - return await _handle_channel_history(agent_id, channel_slug, limit, hours_back) + return await _handle_channel_history( + client, agent_id, channel_slug, limit, hours_back + ) @mcp.tool() async def roboco_message_send(data: SendMessageInput) -> dict[str, Any]: @@ -575,12 +491,12 @@ def create_message_mcp_server(agent_id: str) -> FastMCP: You must have write access to the channel. """ - return await _handle_message_send(agent_id, data) + return await _handle_message_send(client, agent_id, data) @mcp.tool() async def roboco_message_get(message_id: str) -> dict[str, Any]: """Get a specific message by ID.""" - return await _handle_message_get(message_id, agent_id) + return await _handle_message_get(client, message_id) @mcp.tool() async def roboco_ask_question( @@ -594,17 +510,13 @@ def create_message_mcp_server(agent_id: str) -> FastMCP: After asking, wait for an answer before proceeding. """ - - async def send_fn(d: SendMessageInput) -> dict[str, Any]: - return await _handle_message_send(agent_id, d) - data = AskQuestionInput( channel_slug=channel_slug, question=question, context=context, task_id=task_id, ) - return await _handle_ask_question(data, send_fn) + return await _handle_ask_question(client, agent_id, data) @mcp.tool() async def roboco_report_blocker( @@ -618,17 +530,13 @@ def create_message_mcp_server(agent_id: str) -> FastMCP: The PM will be notified automatically. """ - - async def send_fn(d: SendMessageInput) -> dict[str, Any]: - return await _handle_message_send(agent_id, d) - data = ReportBlockerInput( channel_slug=channel_slug, blocker_description=blocker_description, what_needed=what_needed, task_id=task_id, ) - return await _handle_report_blocker(data, send_fn) + return await _handle_report_blocker(client, agent_id, data) return mcp diff --git a/roboco/mcp/notify_server.py b/roboco/mcp/notify_server.py index b37571bb..f3107082 100644 --- a/roboco/mcp/notify_server.py +++ b/roboco/mcp/notify_server.py @@ -13,7 +13,6 @@ Tools: from typing import Any -import httpx from fastapi import status from mcp.server.fastmcp import FastMCP @@ -21,28 +20,15 @@ from roboco.agents_config import ( NOTIFICATION_PERMISSIONS, get_agent_cell, get_agent_role, - get_agent_team, ) -from roboco.config import settings from roboco.mcp.schemas import SendNotificationInput +from roboco.mcp.utils import ApiClient, format_error_response # ============================================================================= # HELPER FUNCTIONS # ============================================================================= -def _get_agent_headers(agent_id: str) -> dict[str, str]: - """Get standard headers for API calls.""" - headers = { - "X-Agent-ID": agent_id, - "X-Agent-Role": get_agent_role(agent_id), - } - team = get_agent_team(agent_id) - if team: - headers["X-Agent-Team"] = team - return headers - - def _check_cell_scope(sender_id: str) -> tuple[bool, str]: """Check if sender can notify within their cell.""" sender_cell = get_agent_cell(sender_id) @@ -75,21 +61,6 @@ def _can_send_notification(sender_id: str, recipient_id: str) -> tuple[bool, str return False, f"You cannot send notifications to {recipient_id}" -def _format_error_response( - error_code: str, - message: str, - details: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Format a standardized error response.""" - return { - "error": { - "code": error_code, - "message": message, - "details": details or {}, - } - } - - # Valid notification types and priorities VALID_NOTIFICATION_TYPES = frozenset( ["info", "alert", "task", "escalation", "approval"] @@ -101,7 +72,7 @@ def _validate_notification_type(notification_type: str) -> dict[str, Any] | None """Validate notification type. Returns error dict or None if valid.""" if notification_type not in VALID_NOTIFICATION_TYPES: valid = sorted(VALID_NOTIFICATION_TYPES) - return _format_error_response( + return format_error_response( "INVALID_TYPE", f"Invalid type. Must be one of: {valid}" ) return None @@ -110,7 +81,7 @@ def _validate_notification_type(notification_type: str) -> dict[str, Any] | None def _validate_priority(priority: str) -> dict[str, Any] | None: """Validate priority. Returns error dict or None if valid.""" if priority not in VALID_PRIORITIES: - return _format_error_response( + return format_error_response( "INVALID_PRIORITY", f"Invalid priority. Must be one of: {sorted(VALID_PRIORITIES)}", ) @@ -123,30 +94,23 @@ def _validate_priority(priority: str) -> dict[str, Any] | None: async def _handle_list( - agent_id: str, + client: ApiClient, unread_only: bool, pending_ack_only: bool, limit: int, ) -> dict[str, Any]: """Handle notification listing.""" - async with httpx.AsyncClient() as client: - params: dict[str, str | int] = { - "unread_only": str(unread_only).lower(), - "pending_ack_only": str(pending_ack_only).lower(), - "limit": limit, - } + params: dict[str, str | int] = { + "unread_only": str(unread_only).lower(), + "pending_ack_only": str(pending_ack_only).lower(), + "limit": limit, + } - resp = await client.get( - f"{settings.internal_api_url}/notifications", - params=params, - headers=_get_agent_headers(agent_id), - ) - - if resp.status_code != status.HTTP_200_OK: - return _format_error_response("API_ERROR", "Failed to fetch notifications") - - data = resp.json() + resp = await client.get("/notifications", params=params) + if not resp.ok: + return format_error_response("API_ERROR", "Failed to fetch notifications") + data = resp.json() unread = data.get("unread_count", 0) pending_ack = data.get("pending_ack_count", 0) @@ -170,27 +134,22 @@ async def _handle_list( } -async def _handle_get(agent_id: str, notification_id: str) -> dict[str, Any]: +async def _handle_get(client: ApiClient, notification_id: str) -> dict[str, Any]: """Handle getting a specific notification.""" - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/notifications/{notification_id}", - headers=_get_agent_headers(agent_id), + resp = await client.get(f"/notifications/{notification_id}") + + if resp.is_status(status.HTTP_404_NOT_FOUND): + return format_error_response("NOT_FOUND", "Notification not found") + + if resp.is_status(status.HTTP_403_FORBIDDEN): + return format_error_response( + "NOT_RECIPIENT", "You are not a recipient of this notification" ) - if resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", "Notification not found") - - if resp.status_code == status.HTTP_403_FORBIDDEN: - return _format_error_response( - "NOT_RECIPIENT", "You are not a recipient of this notification" - ) - - if resp.status_code != status.HTTP_200_OK: - return _format_error_response("API_ERROR", "Failed to fetch notification") - - notification = resp.json() + if not resp.ok: + return format_error_response("API_ERROR", "Failed to fetch notification") + notification = resp.json() guidance = "" if notification.get("requires_ack") and not notification.get("is_acknowledged"): guidance = ( @@ -201,34 +160,27 @@ async def _handle_get(agent_id: str, notification_id: str) -> dict[str, Any]: return {"notification": notification, "guidance": guidance} -async def _handle_ack(agent_id: str, notification_id: str) -> dict[str, Any]: +async def _handle_ack(client: ApiClient, notification_id: str) -> dict[str, Any]: """Handle acknowledging a notification.""" - async with httpx.AsyncClient() as client: - resp = await client.post( - f"{settings.internal_api_url}/notifications/{notification_id}/ack", - headers=_get_agent_headers(agent_id), + resp = await client.post(f"/notifications/{notification_id}/ack") + + if resp.is_status(status.HTTP_404_NOT_FOUND): + return format_error_response("NOT_FOUND", "Notification not found") + + if resp.is_status(status.HTTP_403_FORBIDDEN): + return format_error_response( + "NOT_RECIPIENT", "You are not a recipient of this notification" ) - if resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", "Notification not found") + if resp.is_status(status.HTTP_400_BAD_REQUEST): + return format_error_response( + "NO_ACK_REQUIRED", "This notification does not require acknowledgment" + ) - if resp.status_code == status.HTTP_403_FORBIDDEN: - return _format_error_response( - "NOT_RECIPIENT", "You are not a recipient of this notification" - ) - - if resp.status_code == status.HTTP_400_BAD_REQUEST: - return _format_error_response( - "NO_ACK_REQUIRED", "This notification does not require acknowledgment" - ) - - if resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "API_ERROR", "Failed to acknowledge notification" - ) - - notification = resp.json() + if not resp.ok: + return format_error_response("API_ERROR", "Failed to acknowledge notification") + notification = resp.json() return { "status": "acknowledged", "notification": notification, @@ -241,7 +193,7 @@ def _check_send_permission(agent_id: str) -> dict[str, Any] | None: role = get_agent_role(agent_id) permissions = NOTIFICATION_PERMISSIONS.get(role, {"can_send": False}) if not permissions.get("can_send", False): - return _format_error_response( + return format_error_response( "NOT_AUTHORIZED", f"Agents with role '{role}' cannot send notifications. " "Only PMs, Board members, and Auditor can send notifications.", @@ -259,7 +211,7 @@ def _check_recipients(agent_id: str, recipients: list[str]) -> dict[str, Any] | if not can_send ] if denied: - return _format_error_response( + return format_error_response( "RECIPIENT_DENIED", "Cannot send to one or more recipients", {"denied": denied}, @@ -267,7 +219,9 @@ def _check_recipients(agent_id: str, recipients: list[str]) -> dict[str, Any] | return None -async def _handle_send(agent_id: str, data: SendNotificationInput) -> dict[str, Any]: +async def _handle_send( + client: ApiClient, agent_id: str, data: SendNotificationInput +) -> dict[str, Any]: """Handle sending a notification.""" # Validate permissions and data if error := _check_send_permission(agent_id): @@ -279,30 +233,24 @@ async def _handle_send(agent_id: str, data: SendNotificationInput) -> dict[str, if error := _validate_priority(data.priority): return error - async with httpx.AsyncClient() as client: - payload = { - "type": data.notification_type, - "priority": data.priority, - "to_agents": data.recipients, - "subject": data.subject, - "body": data.body, - "requires_ack": data.requires_ack, - "related_task_id": data.related_task_id, - } + payload = { + "type": data.notification_type, + "priority": data.priority, + "to_agents": data.recipients, + "subject": data.subject, + "body": data.body, + "requires_ack": data.requires_ack, + "related_task_id": data.related_task_id, + } - resp = await client.post( - f"{settings.internal_api_url}/notifications", - json=payload, - headers=_get_agent_headers(agent_id), + resp = await client.post("/notifications", json=payload) + + if not resp.ok: + return format_error_response( + "SEND_FAILED", "Failed to send notification", {"api_error": resp.text} ) - if resp.status_code not in [status.HTTP_200_OK, status.HTTP_201_CREATED]: - return _format_error_response( - "SEND_FAILED", "Failed to send notification", {"api_error": resp.text} - ) - - notification = resp.json() - + notification = resp.json() ack_note = "Recipients must acknowledge." if data.requires_ack else "" count = len(data.recipients) @@ -331,6 +279,9 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: """ mcp = FastMCP(f"roboco-notify-{agent_id}", json_response=True) + # Create shared API client for this agent + client = ApiClient(agent_id) + @mcp.tool() async def roboco_notify_list( unread_only: bool = False, @@ -338,17 +289,17 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: limit: int = 50, ) -> dict[str, Any]: """List your notifications.""" - return await _handle_list(agent_id, unread_only, pending_ack_only, limit) + return await _handle_list(client, unread_only, pending_ack_only, limit) @mcp.tool() async def roboco_notify_get(notification_id: str) -> dict[str, Any]: """Get a specific notification. Also marks it as read.""" - return await _handle_get(agent_id, notification_id) + return await _handle_get(client, notification_id) @mcp.tool() async def roboco_notify_ack(notification_id: str) -> dict[str, Any]: """Acknowledge a notification.""" - return await _handle_ack(agent_id, notification_id) + return await _handle_ack(client, notification_id) @mcp.tool() async def roboco_notify_send(data: SendNotificationInput) -> dict[str, Any]: @@ -358,7 +309,7 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: Only PMs, Board members, and Auditor can send notifications. Cell PMs can only notify their own cell. """ - return await _handle_send(agent_id, data) + return await _handle_send(client, agent_id, data) @mcp.tool() async def roboco_escalate( @@ -374,7 +325,7 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: """ role = get_agent_role(agent_id) if role not in ["cell_pm", "main_pm"]: - return _format_error_response( + return format_error_response( "NOT_PM", "Only PMs can use the escalate function" ) @@ -387,7 +338,7 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: requires_ack=True, related_task_id=task_id, ) - return await _handle_send(agent_id, input_data) + return await _handle_send(client, agent_id, input_data) @mcp.tool() async def roboco_request_approval( @@ -401,7 +352,7 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: """ role = get_agent_role(agent_id) if role not in ["cell_pm", "main_pm", "product_owner", "head_marketing"]: - return _format_error_response( + return format_error_response( "NOT_AUTHORIZED", "Only PMs and Board can request approvals" ) @@ -414,7 +365,7 @@ def create_notify_mcp_server(agent_id: str) -> FastMCP: requires_ack=True, related_task_id=task_id, ) - return await _handle_send(agent_id, input_data) + return await _handle_send(client, agent_id, input_data) return mcp diff --git a/roboco/mcp/schemas/__init__.py b/roboco/mcp/schemas/__init__.py index 7049d127..dc63b507 100644 --- a/roboco/mcp/schemas/__init__.py +++ b/roboco/mcp/schemas/__init__.py @@ -194,3 +194,25 @@ class TaskEscalateInput(BaseModel): escalate_to: str | None = Field( default=None, description="Override default escalation target" ) + + +class TaskBlockInput(BaseModel): + """Input for blocking a task.""" + + task_id: str = Field(..., description="Task ID to block") + reason: str = Field(..., description="Why the task is blocked") + blocker_type: str = Field( + ..., description="Type: external, internal, question, dependency" + ) + what_needed: str = Field(..., description="What is needed to unblock") + + +class TaskPauseInput(BaseModel): + """Input for pausing a task.""" + + task_id: str = Field(..., description="Task ID to pause") + reason: str = Field(..., description="Why pausing") + checkpoint_summary: str = Field(..., description="Summary of current state") + remaining_work: list[str] = Field( + default_factory=list, description="List of remaining sub-tasks" + ) diff --git a/roboco/mcp/task_server.py b/roboco/mcp/task_server.py index 9fd38c66..c5abfe00 100644 --- a/roboco/mcp/task_server.py +++ b/roboco/mcp/task_server.py @@ -25,7 +25,6 @@ Tools: from typing import Any -import httpx from fastapi import status from mcp.server.fastmcp import FastMCP @@ -36,58 +35,36 @@ from roboco.agents_config import ( get_agent_team, get_escalation_target, ) -from roboco.config import settings from roboco.llm import ToonAdapter -from roboco.mcp.schemas import TaskAssignInput, TaskCreateInput, TaskEscalateInput +from roboco.mcp.schemas import ( + TaskAssignInput, + TaskBlockInput, + TaskCreateInput, + TaskEscalateInput, + TaskPauseInput, +) +from roboco.mcp.utils import ( + ApiClient, + format_error_response, + resolve_agent_uuid, +) from roboco.services.task import extract_original_developer - -def _get_agent_headers(agent_id: str) -> dict[str, str]: - """Get standard headers for API calls.""" - headers = { - "X-Agent-ID": agent_id, - "X-Agent-Role": get_agent_role(agent_id), - } - team = get_agent_team(agent_id) - if team: - headers["X-Agent-Team"] = team - return headers - +# Alias for backwards compatibility +_format_error_response = format_error_response # Cache for agent slug -> UUID resolution _agent_uuid_cache: dict[str, str] = {} -async def _resolve_agent_uuid(agent_id: str, headers: dict[str, str]) -> str | None: - """Resolve agent slug to UUID. Returns None if not found.""" - # Check if already a valid UUID - try: - from uuid import UUID - - UUID(agent_id) - return agent_id # Already a UUID - except ValueError: - pass - - # Check cache +async def _resolve_agent_uuid_cached(agent_id: str, client: ApiClient) -> str | None: + """Resolve agent slug to UUID with caching. Returns None if not found.""" if agent_id in _agent_uuid_cache: return _agent_uuid_cache[agent_id] - - # Query API to resolve slug to UUID - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/agents", - params={"slug": agent_id}, - headers=headers, - ) - if resp.status_code == status.HTTP_200_OK: - agents = resp.json() - if agents: - uuid_str = str(agents[0]["id"]) - _agent_uuid_cache[agent_id] = uuid_str - return uuid_str - - return None + result = await resolve_agent_uuid(agent_id, client._get_headers()) + if result: + _agent_uuid_cache[agent_id] = result + return result # Global TOON adapter for encoding task data @@ -124,21 +101,6 @@ def _format_task_response( return response -def _format_error_response( - error_code: str, - message: str, - details: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Format a standardized error response.""" - return { - "error": { - "code": error_code, - "message": message, - "details": details or {}, - } - } - - def _get_next_step_guidance(status: str) -> tuple[str, str]: """Get next step and guidance based on task status.""" guidance_map = { @@ -204,74 +166,55 @@ def _get_next_step_guidance(status: str) -> tuple[str, str]: # ============================================================================= -async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]: +async def _handle_task_scan( + client: ApiClient, team: str | None, agent_id: str +) -> dict[str, Any]: """Handle task scanning.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - paused_resp = await client.get( - f"{settings.internal_api_url}/tasks/my", - params={"status": "paused"}, - headers=headers, - ) - paused_tasks = ( - paused_resp.json() if paused_resp.status_code == status.HTTP_200_OK else [] - ) + paused_resp = await client.get("/tasks/my", params={"status": "paused"}) + paused_tasks = paused_resp.json() if paused_resp.ok else [] - # Get assigned tasks using /tasks/my - # Includes: PM-assigned pending tasks + tasks being actively worked on - assigned_resp = await client.get( - f"{settings.internal_api_url}/tasks/my", - headers=headers, + # Get assigned tasks using /tasks/my + # Includes: PM-assigned pending tasks + tasks being actively worked on + assigned_resp = await client.get("/tasks/my") + assigned_data = assigned_resp.json() if assigned_resp.ok else [] + # Include pending tasks (PM assigned) + active work statuses + assigned_tasks = [ + t + for t in assigned_data + if t.get("status") + in ["pending", "claimed", "in_progress", "verifying", "needs_revision"] + ] + + # Get available tasks based on agent role + # QA agents need awaiting_qa tasks, Documenters need awaiting_documentation + agent_role = get_agent_role(agent_id) + + available_tasks: list[dict[str, Any]] = [] + + if agent_role == "qa": + # QA agents look for tasks awaiting QA review + qa_resp = await client.get( + "/tasks/awaiting-qa", + params={"team": team} if team else {}, ) - assigned_data = ( - assigned_resp.json() - if assigned_resp.status_code == status.HTTP_200_OK - else [] + if qa_resp.ok: + available_tasks = qa_resp.json() + elif agent_role == "documenter": + # Documenters look for tasks awaiting documentation + doc_resp = await client.get( + "/tasks/awaiting-docs", + params={"team": team} if team else {}, ) - # Include pending tasks (PM assigned) + active work statuses - assigned_tasks = [ - t - for t in assigned_data - if t.get("status") - in ["pending", "claimed", "in_progress", "verifying", "needs_revision"] - ] - - # Get available tasks based on agent role - # QA agents need awaiting_qa tasks, Documenters need awaiting_documentation - agent_role = get_agent_role(agent_id) - - available_tasks: list[dict[str, Any]] = [] - - if agent_role == "qa": - # QA agents look for tasks awaiting QA review - qa_resp = await client.get( - f"{settings.internal_api_url}/tasks/awaiting-qa", - params={"team": team} if team else {}, - headers=headers, - ) - if qa_resp.status_code == status.HTTP_200_OK: - available_tasks = qa_resp.json() - elif agent_role == "documenter": - # Documenters look for tasks awaiting documentation - doc_resp = await client.get( - f"{settings.internal_api_url}/tasks/awaiting-docs", - params={"team": team} if team else {}, - headers=headers, - ) - if doc_resp.status_code == status.HTTP_200_OK: - available_tasks = doc_resp.json() - else: - # Developers and PMs look for pending tasks - params: dict[str, Any] = {"status": "pending"} - if team: - params["team"] = team - pending_resp = await client.get( - f"{settings.internal_api_url}/tasks", - params=params, - headers=headers, - ) - if pending_resp.status_code == status.HTTP_200_OK: - available_tasks = pending_resp.json() + if doc_resp.ok: + available_tasks = doc_resp.json() + else: + # Developers and PMs look for pending tasks + params: dict[str, Any] = {"status": "pending"} + if team: + params["team"] = team + pending_resp = await client.get("/tasks", params=params) + if pending_resp.ok: + available_tasks = pending_resp.json() # Filter out tasks already in assigned_tasks from available_tasks # (prevents PM-assigned pending tasks from appearing in both lists) @@ -308,23 +251,17 @@ async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]: } -async def _handle_task_get(task_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_task_get(client: ApiClient, task_id: str) -> dict[str, Any]: """Handle getting task details.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, + resp = await client.get(f"/tasks/{task_id}") + + if resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response( + "NOT_FOUND", + f"Task {task_id} not found", ) - if resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response( - "NOT_FOUND", - f"Task {task_id} not found", - ) - - task = resp.json() - + task = resp.json() next_step, guidance = _get_next_step_guidance(task.get("status", "")) return _format_task_response(task, next_step, guidance) @@ -379,64 +316,54 @@ def _validate_task_claimable(task: dict, agent_role: str) -> dict[str, Any] | No return None -async def _get_project_context(project_id: str, agent_id: str) -> dict[str, Any] | None: +async def _get_project_context( + client: ApiClient, project_id: str +) -> dict[str, Any] | None: """Fetch project context if available.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - resp = await client.get( - f"{settings.internal_api_url}/projects/{project_id}", - headers=headers, - ) - if resp.status_code == status.HTTP_200_OK: - result: dict[str, Any] = resp.json() - return result + resp = await client.get(f"/projects/{project_id}") + if resp.ok: + result: dict[str, Any] = resp.json() + return result return None -async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_task_claim( + client: ApiClient, task_id: str, agent_id: str +) -> dict[str, Any]: """Handle task claiming.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - active_resp = await client.get( - f"{settings.internal_api_url}/tasks/my", - headers=headers, - ) - if active_resp.status_code == status.HTTP_200_OK: - active_tasks = active_resp.json() - if error := _check_blocking_tasks(active_tasks): - return error - if error := _check_paused_tasks(active_tasks): - return error - - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - - task = task_resp.json() - agent_role = get_agent_role(agent_id) - if error := _validate_task_claimable(task, agent_role): + active_resp = await client.get("/tasks/my") + if active_resp.ok: + active_tasks = active_resp.json() + if error := _check_blocking_tasks(active_tasks): + return error + if error := _check_paused_tasks(active_tasks): return error - claim_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/claim", - json={"agent_id": agent_id}, - headers=headers, - ) - if claim_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "CLAIM_FAILED", - "Failed to claim task", - {"api_error": claim_resp.text}, - ) + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - claimed_task = claim_resp.json() + task = task_resp.json() + agent_role = get_agent_role(agent_id) + if error := _validate_task_claimable(task, agent_role): + return error + + claim_resp = await client.post( + f"/tasks/{task_id}/claim", + json={"agent_id": agent_id}, + ) + if not claim_resp.ok: + return _format_error_response( + "CLAIM_FAILED", + "Failed to claim task", + {"api_error": claim_resp.text}, + ) + + claimed_task = claim_resp.json() project = None if claimed_task.get("project_id"): - project = await _get_project_context(claimed_task["project_id"], agent_id) + project = await _get_project_context(client, claimed_task["project_id"]) return _format_task_response( claimed_task, @@ -450,7 +377,7 @@ async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]: async def _validate_task_ownership( - task: dict, agent_id: str, headers: dict[str, str] + task: dict, agent_id: str, client: ApiClient ) -> dict[str, Any] | None: """Validate agent owns the task. Returns error or None.""" assigned_to = task.get("assigned_to") @@ -461,7 +388,7 @@ async def _validate_task_ownership( ) # Resolve agent_id (which may be a slug) to UUID for comparison - agent_uuid = await _resolve_agent_uuid(agent_id, headers) + agent_uuid = await resolve_agent_uuid(agent_id, client._get_headers()) if not agent_uuid: return _format_error_response( "AGENT_NOT_FOUND", @@ -513,40 +440,32 @@ def _build_plan_data(plan_params: dict[str, Any]) -> dict[str, Any]: async def _handle_task_plan( + client: ApiClient, task_id: str, plan_params: dict[str, Any], agent_id: str, ) -> dict[str, Any]: """Handle task planning.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + + task = task_resp.json() + if error := await _validate_task_ownership(task, agent_id, client): + return error + if error := _validate_task_status_claimed(task): + return error + + plan_data = _build_plan_data(plan_params) + update_resp = await client.patch(f"/tasks/{task_id}", json={"plan": plan_data}) + if not update_resp.ok: + return _format_error_response( + "UPDATE_FAILED", + "Failed to save plan", + {"api_error": update_resp.text}, ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error - if error := _validate_task_status_claimed(task): - return error - - plan_data = _build_plan_data(plan_params) - update_resp = await client.patch( - f"{settings.internal_api_url}/tasks/{task_id}", - json={"plan": plan_data}, - headers=headers, - ) - if update_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "UPDATE_FAILED", - "Failed to save plan", - {"api_error": update_resp.text}, - ) - - updated_task = update_resp.json() + updated_task = update_resp.json() open_questions = plan_params.get("open_questions") if open_questions: @@ -566,10 +485,10 @@ async def _handle_task_plan( async def _validate_task_start( - task: dict[str, Any], agent_id: str, headers: dict[str, str] + task: dict[str, Any], agent_id: str, client: ApiClient ) -> dict[str, Any] | None: """Validate task can be started. Returns error dict or None.""" - if error := await _validate_task_ownership(task, agent_id, headers): + if error := await _validate_task_ownership(task, agent_id, client): return error task_status = task.get("status") @@ -601,92 +520,81 @@ async def _validate_task_start( return None -async def _handle_task_start(task_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_task_start( + client: ApiClient, task_id: str, agent_id: str +) -> dict[str, Any]: """Handle task start.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if validation_error := await _validate_task_start(task, agent_id, headers): - return validation_error + if validation_error := await _validate_task_start(task, agent_id, client): + return validation_error - # Start the task - start_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/start", - headers=headers, + # Start the task + start_resp = await client.post(f"/tasks/{task_id}/start") + + if not start_resp.ok: + return _format_error_response( + "START_FAILED", + "Failed to start task", + {"api_error": start_resp.text}, ) - if start_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "START_FAILED", - "Failed to start task", - {"api_error": start_resp.text}, - ) - - return _format_task_response( - start_resp.json(), - "EXECUTE", - "Task started. Work through your plan step by step:\n" - "1. Implement each sub-task\n" - "2. Commit frequently with clear messages\n" - "3. Call roboco_task_progress to update status\n" - "4. If blocked, call roboco_task_block immediately\n" - "5. When done, call roboco_task_submit_verification", - ) + return _format_task_response( + start_resp.json(), + "EXECUTE", + "Task started. Work through your plan step by step:\n" + "1. Implement each sub-task\n" + "2. Commit frequently with clear messages\n" + "3. Call roboco_task_progress to update status\n" + "4. If blocked, call roboco_task_block immediately\n" + "5. When done, call roboco_task_submit_verification", + ) async def _handle_task_progress( + client: ApiClient, task_id: str, message: str, percentage: int | None, agent_id: str, ) -> dict[str, Any]: """Handle task progress update.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error + if error := await _validate_task_ownership(task, agent_id, client): + return error - if task.get("status") != "in_progress": - return _format_error_response( - "INVALID_STATE", - "Can only update progress for in_progress tasks", - ) - - # Add progress update - progress_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/progress", - json={ - "agent_id": agent_id, - "message": message, - "percentage": percentage, - }, - headers=headers, + if task.get("status") != "in_progress": + return _format_error_response( + "INVALID_STATE", + "Can only update progress for in_progress tasks", ) - if progress_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "UPDATE_FAILED", - "Failed to update progress", - ) + # Add progress update + progress_resp = await client.post( + f"/tasks/{task_id}/progress", + json={ + "agent_id": agent_id, + "message": message, + "percentage": percentage, + }, + ) - updated_task = progress_resp.json() + if not progress_resp.ok: + return _format_error_response( + "UPDATE_FAILED", + "Failed to update progress", + ) + + updated_task = progress_resp.json() return _format_task_response( updated_task, @@ -696,74 +604,60 @@ async def _handle_task_progress( async def _handle_task_block( - task_id: str, - reason: str, - blocker_type: str, - what_needed: str, + client: ApiClient, + data: TaskBlockInput, agent_id: str, ) -> dict[str, Any]: """Handle task blocking.""" - if not reason or not what_needed: + task_resp = await client.get(f"/tasks/{data.task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {data.task_id} not found") + + task = task_resp.json() + + if error := await _validate_task_ownership(task, agent_id, client): + return error + + if task.get("status") != "in_progress": return _format_error_response( - "MISSING_DETAILS", - "Both 'reason' and 'what_needed' are required to block a task.", + "INVALID_STATE", + "Can only block in_progress tasks", ) - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + # Build blocker note for dev_notes + blocker_note = ( + f"[BLOCKED - {data.blocker_type.upper()}]\n" + f"Reason: {data.reason}\n" + f"What's needed: {data.what_needed}" + ) + existing_notes = task.get("dev_notes") or "" + if existing_notes: + updated_notes = f"{existing_notes}\n\n{blocker_note}" + else: + updated_notes = blocker_note - task = task_resp.json() + # Block the task using PATCH to update status and notes + block_resp = await client.patch( + f"/tasks/{data.task_id}", + json={ + "status": "blocked", + "dev_notes": updated_notes, + }, + ) - if error := await _validate_task_ownership(task, agent_id, headers): - return error - - if task.get("status") != "in_progress": - return _format_error_response( - "INVALID_STATE", - "Can only block in_progress tasks", - ) - - # Build blocker note for dev_notes - blocker_note = ( - f"[BLOCKED - {blocker_type.upper()}]\n" - f"Reason: {reason}\n" - f"What's needed: {what_needed}" - ) - existing_notes = task.get("dev_notes") or "" - if existing_notes: - updated_notes = f"{existing_notes}\n\n{blocker_note}" - else: - updated_notes = blocker_note - - # Block the task using PATCH to update status and notes - block_resp = await client.patch( - f"{settings.internal_api_url}/tasks/{task_id}", - json={ - "status": "blocked", - "dev_notes": updated_notes, - }, - headers=headers, + if not block_resp.ok: + return _format_error_response( + "BLOCK_FAILED", + "Failed to block task", + {"status_code": block_resp.status_code, "detail": block_resp.text}, ) - if block_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "BLOCK_FAILED", - "Failed to block task", - {"status_code": block_resp.status_code, "detail": block_resp.text}, - ) - - blocked_task = block_resp.json() + blocked_task = block_resp.json() return _format_task_response( blocked_task, "WAIT_OR_SWITCH", - f"Task blocked: {reason}\n\n" + f"Task blocked: {data.reason}\n\n" "Options:\n" "1. WAIT - If resolution expected soon, poll for updates\n" "2. SWITCH - Call roboco_task_scan to work on another task\n" @@ -773,37 +667,31 @@ async def _handle_task_block( ) -async def _handle_task_unblock(task_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_task_unblock( + client: ApiClient, task_id: str, agent_id: str +) -> dict[str, Any]: """Handle task unblocking.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error + if error := await _validate_task_ownership(task, agent_id, client): + return error - if task.get("status") != "blocked": - return _format_error_response( - "INVALID_STATE", - "Task is not blocked", - ) - - unblock_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/unblock", - headers=headers, + if task.get("status") != "blocked": + return _format_error_response( + "INVALID_STATE", + "Task is not blocked", ) - if unblock_resp.status_code != status.HTTP_200_OK: - return _format_error_response("UNBLOCK_FAILED", "Failed to unblock task") + unblock_resp = await client.post(f"/tasks/{task_id}/unblock") - unblocked_task = unblock_resp.json() + if not unblock_resp.ok: + return _format_error_response("UNBLOCK_FAILED", "Failed to unblock task") + + unblocked_task = unblock_resp.json() return _format_task_response( unblocked_task, @@ -813,114 +701,95 @@ async def _handle_task_unblock(task_id: str, agent_id: str) -> dict[str, Any]: async def _handle_task_pause( - task_id: str, - reason: str, - checkpoint_summary: str, - remaining_work: list[str], + client: ApiClient, + data: TaskPauseInput, agent_id: str, ) -> dict[str, Any]: """Handle task pausing.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{data.task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {data.task_id} not found") - task = task_resp.json() + task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error + if error := await _validate_task_ownership(task, agent_id, client): + return error - if task.get("status") != "in_progress": - return _format_error_response( - "INVALID_STATE", - "Can only pause in_progress tasks", - ) - - # Add checkpoint - await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/checkpoint", - json={ - "agent_id": agent_id, - "state_summary": checkpoint_summary, - "remaining_work": remaining_work, - "notes": reason, - }, - headers=headers, + if task.get("status") != "in_progress": + return _format_error_response( + "INVALID_STATE", + "Can only pause in_progress tasks", ) - # Pause the task - pause_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/pause", - headers=headers, - ) + # Add checkpoint + await client.post( + f"/tasks/{data.task_id}/checkpoint", + json={ + "agent_id": agent_id, + "state_summary": data.checkpoint_summary, + "remaining_work": data.remaining_work, + "notes": data.reason, + }, + ) - if pause_resp.status_code != status.HTTP_200_OK: - return _format_error_response("PAUSE_FAILED", "Failed to pause task") + # Pause the task + pause_resp = await client.post(f"/tasks/{data.task_id}/pause") - paused_task = pause_resp.json() + if not pause_resp.ok: + return _format_error_response("PAUSE_FAILED", "Failed to pause task") + + paused_task = pause_resp.json() return _format_task_response( paused_task, "SCAN_FOR_WORK", f"Task paused. Checkpoint saved.\n" - f"Reason: {reason}\n\n" + f"Reason: {data.reason}\n\n" "To resume later, call roboco_task_start with this task_id.\n" "Now call roboco_task_scan to find your next task.", ) async def _handle_task_submit_verification( - task_id: str, agent_id: str + client: ApiClient, task_id: str, agent_id: str ) -> dict[str, Any]: """Handle task verification submission.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error + if error := await _validate_task_ownership(task, agent_id, client): + return error - if task.get("status") != "in_progress": - return _format_error_response( - "INVALID_STATE", - "Can only submit in_progress tasks for verification", - ) - - # Check for evidence of work done (commits OR progress updates) - # Non-code tasks (testing, research, docs) may not have commits - has_commits = bool(task.get("commits")) - has_progress = bool(task.get("progress_updates")) - has_checkpoints = bool(task.get("checkpoints")) - - if not (has_commits or has_progress or has_checkpoints): - return _format_error_response( - "NO_WORK_EVIDENCE", - "No evidence of work found. Add commits with roboco_task_add_commit " - "or update progress with roboco_task_progress before verification.", - ) - - verify_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/verify", - headers=headers, + if task.get("status") != "in_progress": + return _format_error_response( + "INVALID_STATE", + "Can only submit in_progress tasks for verification", ) - if verify_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "VERIFY_FAILED", "Failed to submit for verification" - ) + # Check for evidence of work done (commits OR progress updates) + # Non-code tasks (testing, research, docs) may not have commits + has_commits = bool(task.get("commits")) + has_progress = bool(task.get("progress_updates")) + has_checkpoints = bool(task.get("checkpoints")) - verifying_task = verify_resp.json() + if not (has_commits or has_progress or has_checkpoints): + return _format_error_response( + "NO_WORK_EVIDENCE", + "No evidence of work found. Add commits with roboco_task_add_commit " + "or update progress with roboco_task_progress before verification.", + ) + + verify_resp = await client.post(f"/tasks/{task_id}/verify") + + if not verify_resp.ok: + return _format_error_response( + "VERIFY_FAILED", "Failed to submit for verification" + ) + + verifying_task = verify_resp.json() # Build verification checklist from acceptance criteria criteria = task.get("acceptance_criteria", []) @@ -937,6 +806,7 @@ async def _handle_task_submit_verification( async def _handle_task_submit_qa( + client: ApiClient, task_id: str, dev_notes: str, handoff_summary: str, @@ -949,45 +819,33 @@ async def _handle_task_submit_qa( "Both dev_notes and handoff_summary are required for QA submission.", ) - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if error := await _validate_task_ownership(task, agent_id, headers): - return error + if error := await _validate_task_ownership(task, agent_id, client): + return error - if task.get("status") != "verifying": - return _format_error_response( - "INVALID_STATE", - "Can only submit verified tasks for QA", - ) - - # Update with notes - combine dev_notes and handoff summary - # (handoff summary goes into dev_notes for documenter to read) - combined_notes = f"{dev_notes}\n\n---\nHandoff Summary:\n{handoff_summary}" - await client.patch( - f"{settings.internal_api_url}/tasks/{task_id}", - json={"dev_notes": combined_notes}, - headers=headers, + if task.get("status") != "verifying": + return _format_error_response( + "INVALID_STATE", + "Can only submit verified tasks for QA", ) - # Submit for QA - qa_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/submit-qa", - headers=headers, - ) + # Update with notes - combine dev_notes and handoff summary + # (handoff summary goes into dev_notes for documenter to read) + combined_notes = f"{dev_notes}\n\n---\nHandoff Summary:\n{handoff_summary}" + await client.patch(f"/tasks/{task_id}", json={"dev_notes": combined_notes}) - if qa_resp.status_code != status.HTTP_200_OK: - return _format_error_response("SUBMIT_FAILED", "Failed to submit for QA") + # Submit for QA + qa_resp = await client.post(f"/tasks/{task_id}/submit-qa") - qa_task = qa_resp.json() + if not qa_resp.ok: + return _format_error_response("SUBMIT_FAILED", "Failed to submit for QA") + + qa_task = qa_resp.json() return _format_task_response( qa_task, @@ -999,6 +857,7 @@ async def _handle_task_submit_qa( async def _handle_task_qa_pass( + client: ApiClient, task_id: str, qa_notes: str, agent_id: str, @@ -1011,50 +870,44 @@ async def _handle_task_qa_pass( "Only QA agents can pass tasks through QA review.", ) - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if task.get("status") != "awaiting_qa": - return _format_error_response( - "INVALID_STATE", - "Task is not awaiting QA", - ) - - # Check QA is not reviewing own work - # Check against original developer stored in quick_context - quick_context = task.get("quick_context") - original_dev = extract_original_developer(quick_context) - - # Resolve agent_id to UUID for proper comparison - agent_uuid = await _resolve_agent_uuid(agent_id, headers) - if original_dev and agent_uuid and agent_uuid == original_dev: - return _format_error_response( - "SELF_REVIEW", - "Cannot review your own work.", - ) - - pass_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/pass-qa", - json={"notes": qa_notes}, - headers=headers, + if task.get("status") != "awaiting_qa": + return _format_error_response( + "INVALID_STATE", + "Task is not awaiting QA", ) - if pass_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "QA_FAILED", - "Failed to pass QA", - {"status_code": pass_resp.status_code, "api_error": pass_resp.text}, - ) + # Check QA is not reviewing own work + # Check against original developer stored in quick_context + quick_context = task.get("quick_context") + original_dev = extract_original_developer(quick_context) - passed_task = pass_resp.json() + # Resolve agent_id to UUID for proper comparison + agent_uuid = await resolve_agent_uuid(agent_id, client._get_headers()) + if original_dev and agent_uuid and agent_uuid == original_dev: + return _format_error_response( + "SELF_REVIEW", + "Cannot review your own work.", + ) + + pass_resp = await client.post( + f"/tasks/{task_id}/pass-qa", + json={"notes": qa_notes}, + ) + + if not pass_resp.ok: + return _format_error_response( + "QA_FAILED", + "Failed to pass QA", + {"status_code": pass_resp.status_code, "api_error": pass_resp.text}, + ) + + passed_task = pass_resp.json() return _format_task_response( passed_task, @@ -1065,6 +918,7 @@ async def _handle_task_qa_pass( async def _handle_task_qa_fail( + client: ApiClient, task_id: str, qa_notes: str, issues: list[str], @@ -1083,39 +937,33 @@ async def _handle_task_qa_fail( "Must specify at least one issue when failing QA.", ) - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if task.get("status") != "awaiting_qa": - return _format_error_response( - "INVALID_STATE", - "Task is not awaiting QA", - ) - - full_notes = f"{qa_notes}\n\nIssues:\n" + "\n".join(f"- {i}" for i in issues) - - fail_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/fail-qa", - json={"notes": full_notes}, - headers=headers, + if task.get("status") != "awaiting_qa": + return _format_error_response( + "INVALID_STATE", + "Task is not awaiting QA", ) - if fail_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "QA_FAILED", - "Failed to fail QA", - {"status_code": fail_resp.status_code, "api_error": fail_resp.text}, - ) + full_notes = f"{qa_notes}\n\nIssues:\n" + "\n".join(f"- {i}" for i in issues) - failed_task = fail_resp.json() + fail_resp = await client.post( + f"/tasks/{task_id}/fail-qa", + json={"notes": full_notes}, + ) + + if not fail_resp.ok: + return _format_error_response( + "QA_FAILED", + "Failed to fail QA", + {"status_code": fail_resp.status_code, "api_error": fail_resp.text}, + ) + + failed_task = fail_resp.json() return _format_task_response( failed_task, @@ -1126,41 +974,33 @@ async def _handle_task_qa_fail( ) -async def _handle_task_complete(task_id: str, agent_id: str) -> dict[str, Any]: +async def _handle_task_complete(client: ApiClient, task_id: str) -> dict[str, Any]: """Handle task completion.""" - headers = _get_agent_headers(agent_id) - async with httpx.AsyncClient() as client: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response("NOT_FOUND", f"Task {task_id} not found") + task_resp = await client.get(f"/tasks/{task_id}") + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response("NOT_FOUND", f"Task {task_id} not found") - task = task_resp.json() + task = task_resp.json() - if task.get("status") != "awaiting_documentation": - return _format_error_response( - "INVALID_STATE", - "Task must be awaiting documentation to complete", - ) - - complete_resp = await client.post( - f"{settings.internal_api_url}/tasks/{task_id}/complete", - headers=headers, + if task.get("status") != "awaiting_documentation": + return _format_error_response( + "INVALID_STATE", + "Task must be awaiting documentation to complete", ) - if complete_resp.status_code != status.HTTP_200_OK: - return _format_error_response( - "COMPLETE_FAILED", - "Failed to complete task", - { - "status_code": complete_resp.status_code, - "api_error": complete_resp.text, - }, - ) + complete_resp = await client.post(f"/tasks/{task_id}/complete") - completed_task = complete_resp.json() + if not complete_resp.ok: + return _format_error_response( + "COMPLETE_FAILED", + "Failed to complete task", + { + "status_code": complete_resp.status_code, + "api_error": complete_resp.text, + }, + ) + + completed_task = complete_resp.json() return _format_task_response( completed_task, @@ -1169,81 +1009,68 @@ async def _handle_task_complete(task_id: str, agent_id: str) -> dict[str, Any]: ) -async def _handle_agent_idle(agent_id: str) -> dict[str, Any]: +async def _handle_agent_idle(client: ApiClient, agent_id: str) -> dict[str, Any]: """Handle agent going idle (no work available).""" - headers = _get_agent_headers(agent_id) + # First, check if agent has any in-progress tasks + # Use /tasks/my endpoint which properly uses authenticated agent context + try: + scan_resp = await client.get("/tasks/my", params={"status": "in_progress"}) + if scan_resp.ok: + tasks = scan_resp.json() # /tasks/my returns list directly + if tasks: + # Agent has in-progress tasks - they must handle them first + task_info = [ + {"id": t.get("id"), "title": t.get("title")} for t in tasks + ] + return _format_error_response( + "TASKS_IN_PROGRESS", + ( + "You have in-progress tasks. Handle them before going " + "idle using: roboco_task_pause (to pause), " + "roboco_task_submit_qa (if done), or " + "roboco_task_complete (if approved)." + ), + {"tasks": task_info}, + ) + except Exception: + # If check fails, continue to mark idle (fail open) + pass - async with httpx.AsyncClient(timeout=30.0) as client: - # First, check if agent has any in-progress tasks - # Use /tasks/my endpoint which properly uses authenticated agent context - try: - scan_resp = await client.get( - f"{settings.internal_api_url}/tasks/my", - params={"status": "in_progress"}, - headers=headers, - ) - if scan_resp.status_code == status.HTTP_200_OK: - tasks = scan_resp.json() # /tasks/my returns list directly - if tasks: - # Agent has in-progress tasks - they must handle them first - task_info = [ - {"id": t.get("id"), "title": t.get("title")} for t in tasks - ] - return _format_error_response( - "TASKS_IN_PROGRESS", - ( - "You have in-progress tasks. Handle them before going " - "idle using: roboco_task_pause (to pause), " - "roboco_task_submit_qa (if done), or " - "roboco_task_complete (if approved)." - ), - {"tasks": task_info}, - ) - except Exception: - # If check fails, continue to mark idle (fail open) - pass - - try: - # Signal to orchestrator that this agent is idle - resp = await client.post( - f"{settings.internal_api_url}/orchestrator/agents/{agent_id}/mark-waiting", - params={"waiting_for": "task_assignment"}, - headers=headers, - ) - except httpx.TimeoutException: - return _format_error_response( - "TIMEOUT", - "Request to mark idle timed out", - ) - except httpx.RequestError as e: - return _format_error_response( - "CONNECTION_ERROR", - f"Failed to connect to orchestrator: {type(e).__name__}", - ) - - if resp.status_code == status.HTTP_204_NO_CONTENT: - return { - "status": "idle", - "message": ( - "You are now in WAITING state. Your container will terminate " - "to save resources. You will be respawned when work is available." - ), - "action": "EXIT_GRACEFULLY", - } - - # Handle specific error codes - if resp.status_code == status.HTTP_503_SERVICE_UNAVAILABLE: - return _format_error_response( - "ORCHESTRATOR_UNAVAILABLE", - "Orchestrator is not running. Cannot mark idle state.", - {"detail": resp.text}, - ) - - return _format_error_response( - "IDLE_FAILED", - "Failed to signal idle state to orchestrator", - {"status_code": resp.status_code, "detail": resp.text}, + try: + # Signal to orchestrator that this agent is idle + resp = await client.post( + f"/orchestrator/agents/{agent_id}/mark-waiting", + params={"waiting_for": "task_assignment"}, ) + except Exception as e: + return _format_error_response( + "CONNECTION_ERROR", + f"Failed to connect to orchestrator: {type(e).__name__}", + ) + + if resp.is_status(status.HTTP_204_NO_CONTENT): + return { + "status": "idle", + "message": ( + "You are now in WAITING state. Your container will terminate " + "to save resources. You will be respawned when work is available." + ), + "action": "EXIT_GRACEFULLY", + } + + # Handle specific error codes + if resp.is_status(status.HTTP_503_SERVICE_UNAVAILABLE): + return _format_error_response( + "ORCHESTRATOR_UNAVAILABLE", + "Orchestrator is not running. Cannot mark idle state.", + {"detail": resp.text}, + ) + + return _format_error_response( + "IDLE_FAILED", + "Failed to signal idle state to orchestrator", + {"status_code": resp.status_code, "detail": resp.text}, + ) # ============================================================================= @@ -1252,11 +1079,11 @@ async def _handle_agent_idle(agent_id: str) -> dict[str, Any]: async def _handle_task_create( + client: ApiClient, input_data: TaskCreateInput, agent_id: str, ) -> dict[str, Any]: """Handle task creation by PM.""" - headers = _get_agent_headers(agent_id) agent_team = get_agent_team(agent_id) # Validate PM role @@ -1276,49 +1103,44 @@ async def _handle_task_create( {"requested_team": input_data.team, "agent_team": agent_team}, ) - async with httpx.AsyncClient(timeout=30.0) as client: - # Build task payload - payload: dict[str, Any] = { - "title": input_data.title, - "description": input_data.description, - "acceptance_criteria": input_data.acceptance_criteria, - "team": input_data.team, - "priority": input_data.priority, - "estimated_complexity": input_data.complexity, - } - if input_data.parent_task_id: - payload["parent_task_id"] = input_data.parent_task_id + # Build task payload + payload: dict[str, Any] = { + "title": input_data.title, + "description": input_data.description, + "acceptance_criteria": input_data.acceptance_criteria, + "team": input_data.team, + "priority": input_data.priority, + "estimated_complexity": input_data.complexity, + } + if input_data.parent_task_id: + payload["parent_task_id"] = input_data.parent_task_id - # Create the task - try: - create_resp = await client.post( - f"{settings.internal_api_url}/tasks", - json=payload, - headers=headers, - ) - except httpx.RequestError as e: - return _format_error_response( - "CONNECTION_ERROR", - f"Failed to connect to API: {type(e).__name__}", - ) + # Create the task + try: + create_resp = await client.post("/tasks", json=payload) + except Exception as e: + return _format_error_response( + "CONNECTION_ERROR", + f"Failed to connect to API: {type(e).__name__}", + ) - if create_resp.status_code != status.HTTP_201_CREATED: - return _format_error_response( - "CREATE_FAILED", - "Failed to create task", - {"status_code": create_resp.status_code, "detail": create_resp.text}, - ) + if not create_resp.is_status(status.HTTP_201_CREATED): + return _format_error_response( + "CREATE_FAILED", + "Failed to create task", + {"status_code": create_resp.status_code, "detail": create_resp.text}, + ) - task = create_resp.json() + task = create_resp.json() - # If assigned_to specified, set assignee but keep pending (don't claim) - # Orchestrator will spawn the agent who will then claim it - if input_data.assigned_to: - assigned_task, _ = await _assign_task_to_agent( - client, task["id"], input_data.assigned_to, headers - ) - if assigned_task: - task = assigned_task + # If assigned_to specified, set assignee but keep pending (don't claim) + # Orchestrator will spawn the agent who will then claim it + if input_data.assigned_to: + assigned_task, _ = await _assign_task_to_agent( + client, task["id"], input_data.assigned_to + ) + if assigned_task: + task = assigned_task guidance = f"Task created successfully. ID: {task['id']}. " if input_data.assigned_to: @@ -1362,26 +1184,22 @@ def _validate_cell_pm_assignment( async def _fetch_task_for_assignment( - client: httpx.AsyncClient, + client: ApiClient, task_id: str, - headers: dict[str, str], ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: """Fetch task for assignment. Returns (task, error) tuple.""" try: - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{task_id}", - headers=headers, - ) - except httpx.RequestError as e: + task_resp = await client.get(f"/tasks/{task_id}") + except Exception as e: return None, _format_error_response( "CONNECTION_ERROR", f"Failed to connect to API: {type(e).__name__}", ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: + if task_resp.is_status(status.HTTP_404_NOT_FOUND): return None, _format_error_response("NOT_FOUND", f"Task {task_id} not found") - if task_resp.status_code != status.HTTP_200_OK: + if not task_resp.ok: return None, _format_error_response( "FETCH_FAILED", "Failed to fetch task", @@ -1392,10 +1210,9 @@ async def _fetch_task_for_assignment( async def _assign_task_to_agent( - client: httpx.AsyncClient, + client: ApiClient, task_id: str, assignee: str, - headers: dict[str, str], ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: """ Assign task to agent by setting assigned_to and resetting to pending. @@ -1406,7 +1223,7 @@ async def _assign_task_to_agent( Returns (assigned_task, error) tuple. """ # Resolve assignee slug to UUID - assignee_id = await _resolve_agent_uuid(assignee, headers) + assignee_id = await resolve_agent_uuid(assignee, client._get_headers()) if not assignee_id: return None, _format_error_response( "INVALID_ASSIGNEE", @@ -1417,17 +1234,16 @@ async def _assign_task_to_agent( try: # PATCH to set assigned_to and reset status to pending assign_resp = await client.patch( - f"{settings.internal_api_url}/tasks/{task_id}", + f"/tasks/{task_id}", json={"assigned_to": assignee_id, "status": "pending"}, - headers=headers, ) - except httpx.RequestError as e: + except Exception as e: return None, _format_error_response( "CONNECTION_ERROR", f"Failed to connect to API: {type(e).__name__}", ) - if assign_resp.status_code != status.HTTP_200_OK: + if not assign_resp.ok: return None, _format_error_response( "ASSIGN_FAILED", "Failed to assign task", @@ -1438,11 +1254,11 @@ async def _assign_task_to_agent( async def _handle_task_assign( + client: ApiClient, input_data: TaskAssignInput, agent_id: str, ) -> dict[str, Any]: """Handle task assignment by PM.""" - headers = _get_agent_headers(agent_id) agent_team = get_agent_team(agent_id) role = get_agent_role(agent_id) @@ -1454,28 +1270,25 @@ async def _handle_task_assign( {"role": role}, ) - async with httpx.AsyncClient(timeout=30.0) as client: - # Get task details first - task, error = await _fetch_task_for_assignment( - client, input_data.task_id, headers - ) - if error or task is None: - return error or _format_error_response("FETCH_FAILED", "No task returned") + # Get task details first + task, error = await _fetch_task_for_assignment(client, input_data.task_id) + if error or task is None: + return error or _format_error_response("FETCH_FAILED", "No task returned") - # Validate Cell PM restrictions - validation_error = _validate_cell_pm_assignment( - role, agent_team, task, input_data.assignee - ) - if validation_error: - return validation_error + # Validate Cell PM restrictions + validation_error = _validate_cell_pm_assignment( + role, agent_team, task, input_data.assignee + ) + if validation_error: + return validation_error - # Assign task to agent (sets assigned_to and resets to pending) - # This is NOT claiming - the dev will claim when spawned - assigned_task, assign_error = await _assign_task_to_agent( - client, input_data.task_id, input_data.assignee, headers - ) - if assign_error or assigned_task is None: - return assign_error or _format_error_response("ASSIGN_FAILED", "No task") + # Assign task to agent (sets assigned_to and resets to pending) + # This is NOT claiming - the dev will claim when spawned + assigned_task, assign_error = await _assign_task_to_agent( + client, input_data.task_id, input_data.assignee + ) + if assign_error or assigned_task is None: + return assign_error or _format_error_response("ASSIGN_FAILED", "No task") guidance = ( f"Task assigned to {input_data.assignee} and set to pending. " @@ -1485,12 +1298,11 @@ async def _handle_task_assign( async def _handle_task_escalate( + client: ApiClient, input_data: TaskEscalateInput, agent_id: str, ) -> dict[str, Any]: """Handle task escalation up the hierarchy.""" - headers = _get_agent_headers(agent_id) - # Determine and resolve escalation target upfront target = input_data.escalate_to or get_escalation_target(agent_id) if not target: @@ -1500,7 +1312,7 @@ async def _handle_task_escalate( {"role": get_agent_role(agent_id)}, ) - target_uuid = await _resolve_agent_uuid(target, headers) + target_uuid = await resolve_agent_uuid(target, client._get_headers()) if not target_uuid: return _format_error_response( "INVALID_TARGET", @@ -1508,47 +1320,39 @@ async def _handle_task_escalate( ) try: - async with httpx.AsyncClient(timeout=30.0) as client: - # Get task details - task_resp = await client.get( - f"{settings.internal_api_url}/tasks/{input_data.task_id}", - headers=headers, + # Get task details + task_resp = await client.get(f"/tasks/{input_data.task_id}") + + if task_resp.is_status(status.HTTP_404_NOT_FOUND): + return _format_error_response( + "NOT_FOUND", f"Task {input_data.task_id} not found" ) - if task_resp.status_code == status.HTTP_404_NOT_FOUND: - return _format_error_response( - "NOT_FOUND", f"Task {input_data.task_id} not found" - ) + task = task_resp.json() - task = task_resp.json() + # Create escalation notification + notif_resp = await client.post( + "/notifications", + json={ + "type": "blocker_escalation", + "to_agents": [target_uuid], + "subject": f"Escalation: {task.get('title', 'Unknown task')}", + "body": ( + f"Task {input_data.task_id} escalated by {agent_id}.\n\n" + f"Reason: {input_data.reason}" + ), + "related_task_id": input_data.task_id, + "priority": "high", + }, + ) - # Create escalation notification - notif_resp = await client.post( - f"{settings.internal_api_url}/notifications", - json={ - "type": "blocker_escalation", - "to_agents": [target_uuid], - "subject": f"Escalation: {task.get('title', 'Unknown task')}", - "body": ( - f"Task {input_data.task_id} escalated by {agent_id}.\n\n" - f"Reason: {input_data.reason}" - ), - "related_task_id": input_data.task_id, - "priority": "high", - }, - headers=headers, + if not notif_resp.ok and not notif_resp.is_status(status.HTTP_201_CREATED): + return _format_error_response( + "ESCALATION_FAILED", + "Failed to send escalation notification", + {"status_code": notif_resp.status_code, "detail": notif_resp.text}, ) - - if notif_resp.status_code not in ( - status.HTTP_200_OK, - status.HTTP_201_CREATED, - ): - return _format_error_response( - "ESCALATION_FAILED", - "Failed to send escalation notification", - {"status_code": notif_resp.status_code, "detail": notif_resp.text}, - ) - except httpx.RequestError as e: + except Exception as e: return _format_error_response( "CONNECTION_ERROR", f"Failed to connect to API: {type(e).__name__}", @@ -1580,6 +1384,9 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: """ mcp = FastMCP(f"roboco-task-{agent_id}", json_response=True) + # Create shared API client for this agent + client = ApiClient(agent_id) + @mcp.tool() async def roboco_task_scan( team: str | None = None, @@ -1598,7 +1405,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Dict with paused/assigned/available tasks and guidance """ - return await _handle_task_scan(team, agent_id) + return await _handle_task_scan(client, team, agent_id) @mcp.tool() async def roboco_task_get(task_id: str) -> dict[str, Any]: @@ -1611,7 +1418,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Task details with current status and guidance """ - return await _handle_task_get(task_id, agent_id) + return await _handle_task_get(client, task_id) @mcp.tool() async def roboco_task_claim(task_id: str) -> dict[str, Any]: @@ -1629,7 +1436,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Claimed task with project context and next step guidance """ - return await _handle_task_claim(task_id, agent_id) + return await _handle_task_claim(client, task_id, agent_id) @mcp.tool() async def roboco_task_plan( @@ -1667,7 +1474,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: "risks": risks, "open_questions": open_questions, } - return await _handle_task_plan(task_id, plan_params, agent_id) + return await _handle_task_plan(client, task_id, plan_params, agent_id) @mcp.tool() async def roboco_task_start(task_id: str) -> dict[str, Any]: @@ -1685,7 +1492,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Updated task with execution guidance """ - return await _handle_task_start(task_id, agent_id) + return await _handle_task_start(client, task_id, agent_id) @mcp.tool() async def roboco_task_progress( @@ -1704,7 +1511,9 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Updated task """ - return await _handle_task_progress(task_id, message, percentage, agent_id) + return await _handle_task_progress( + client, task_id, message, percentage, agent_id + ) @mcp.tool() async def roboco_task_block( @@ -1729,9 +1538,13 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Updated task with options """ - return await _handle_task_block( - task_id, reason, blocker_type, what_needed, agent_id + data = TaskBlockInput( + task_id=task_id, + reason=reason, + blocker_type=blocker_type, + what_needed=what_needed, ) + return await _handle_task_block(client, data, agent_id) @mcp.tool() async def roboco_task_unblock(task_id: str) -> dict[str, Any]: @@ -1748,7 +1561,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Updated task ready for work """ - return await _handle_task_unblock(task_id, agent_id) + return await _handle_task_unblock(client, task_id, agent_id) @mcp.tool() async def roboco_task_pause( @@ -1773,9 +1586,13 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Paused task with resume instructions """ - return await _handle_task_pause( - task_id, reason, checkpoint_summary, remaining_work, agent_id + data = TaskPauseInput( + task_id=task_id, + reason=reason, + checkpoint_summary=checkpoint_summary, + remaining_work=remaining_work, ) + return await _handle_task_pause(client, data, agent_id) @mcp.tool() async def roboco_task_submit_verification( @@ -1794,7 +1611,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Task in verifying status with checklist """ - return await _handle_task_submit_verification(task_id, agent_id) + return await _handle_task_submit_verification(client, task_id, agent_id) @mcp.tool() async def roboco_task_submit_qa( @@ -1818,7 +1635,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Task submitted for QA """ return await _handle_task_submit_qa( - task_id, dev_notes, handoff_summary, agent_id + client, task_id, dev_notes, handoff_summary, agent_id ) @mcp.tool() @@ -1841,7 +1658,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Task ready for documentation """ - return await _handle_task_qa_pass(task_id, qa_notes, agent_id) + return await _handle_task_qa_pass(client, task_id, qa_notes, agent_id) @mcp.tool() async def roboco_task_qa_fail( @@ -1865,7 +1682,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Task returned for revision """ - return await _handle_task_qa_fail(task_id, qa_notes, issues, agent_id) + return await _handle_task_qa_fail(client, task_id, qa_notes, issues, agent_id) @mcp.tool() async def roboco_task_complete(task_id: str) -> dict[str, Any]: @@ -1882,7 +1699,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Completed task """ - return await _handle_task_complete(task_id, agent_id) + return await _handle_task_complete(client, task_id) @mcp.tool() async def roboco_agent_idle() -> dict[str, Any]: @@ -1896,7 +1713,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Confirmation of idle state """ - return await _handle_agent_idle(agent_id) + return await _handle_agent_idle(client, agent_id) # ========================================================================= # PM DELEGATION TOOLS @@ -1923,7 +1740,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Returns: Created task with next step guidance """ - return await _handle_task_create(data, agent_id) + return await _handle_task_create(client, data, agent_id) @mcp.tool() async def roboco_task_assign( @@ -1951,7 +1768,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: Updated task with assignment confirmation """ input_data = TaskAssignInput(task_id=task_id, assignee=assignee) - return await _handle_task_assign(input_data, agent_id) + return await _handle_task_assign(client, input_data, agent_id) @mcp.tool() async def roboco_task_escalate( @@ -1986,7 +1803,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP: reason=reason, escalate_to=escalate_to, ) - return await _handle_task_escalate(input_data, agent_id) + return await _handle_task_escalate(client, input_data, agent_id) return mcp diff --git a/roboco/mcp/utils.py b/roboco/mcp/utils.py new file mode 100644 index 00000000..05313ac9 --- /dev/null +++ b/roboco/mcp/utils.py @@ -0,0 +1,422 @@ +""" +MCP Server Utilities + +Shared utilities for all MCP servers to avoid code duplication. +Contains common functions for: +- Agent header generation +- Error response formatting +- Agent UUID resolution +- API client for internal API calls +""" + +from typing import Any + +import httpx + +from roboco.agents_config import get_agent_role, get_agent_team +from roboco.config import settings + +# UUID format constants +_UUID_LENGTH = 36 # Standard UUID string length +_UUID_HYPHEN_COUNT = 4 # Number of hyphens in a UUID + +# HTTP status code ranges +_HTTP_OK = 200 +_HTTP_SUCCESS_MIN = 200 +_HTTP_SUCCESS_MAX = 300 # exclusive + +# Default timeout for API calls (seconds) +DEFAULT_TIMEOUT = 30.0 + + +def get_agent_headers(agent_id: str) -> dict[str, str]: + """ + Get standard headers for API calls from an MCP server. + + Args: + agent_id: The agent's identifier (slug or UUID) + + Returns: + Headers dict with X-Agent-ID, X-Agent-Role, and optionally X-Agent-Team + """ + headers = { + "X-Agent-ID": agent_id, + "X-Agent-Role": get_agent_role(agent_id), + } + team = get_agent_team(agent_id) + if team: + headers["X-Agent-Team"] = team + return headers + + +def format_error_response( + code: str, + message: str, + details: dict[str, Any] | None = None, +) -> dict[str, Any]: + """ + Format a standardized error response for MCP tools. + + Args: + code: Error code (e.g., "NOT_FOUND", "API_ERROR", "PERMISSION_DENIED") + message: Human-readable error message + details: Optional additional error details + + Returns: + Standardized error response dict + """ + response: dict[str, Any] = { + "error": { + "code": code, + "message": message, + } + } + if details: + response["error"]["details"] = details + return response + + +async def resolve_agent_uuid( + agent_id: str, + headers: dict[str, str], +) -> str | None: + """ + Resolve an agent identifier to its UUID. + + If the agent_id is already a UUID, returns it directly. + Otherwise, looks up the agent by slug/name via the API. + + Args: + agent_id: Agent identifier (slug like "be-dev-1" or UUID string) + headers: Request headers for API authentication + + Returns: + UUID string if found, None otherwise + """ + # If it looks like a UUID already, return it + if len(agent_id) == _UUID_LENGTH and agent_id.count("-") == _UUID_HYPHEN_COUNT: + return agent_id + + # Look up by slug + async with httpx.AsyncClient() as client: + try: + resp = await client.get( + f"{settings.internal_api_url}/agents/by-slug/{agent_id}", + headers=headers, + ) + if resp.status_code == _HTTP_OK: + data = resp.json() + agent_id_result: str | None = data.get("id") + return agent_id_result + except Exception: + pass + return None + + +# ============================================================================= +# API CLIENT +# ============================================================================= + + +class ApiResponse: + """Wrapper for API response with convenience methods.""" + + def __init__(self, response: httpx.Response) -> None: + self._response = response + + @property + def ok(self) -> bool: + """True if status code indicates success (2xx).""" + return ( + self._response.status_code >= _HTTP_SUCCESS_MIN + and self._response.status_code < _HTTP_SUCCESS_MAX + ) + + @property + def status_code(self) -> int: + """HTTP status code.""" + return self._response.status_code + + def json(self) -> Any: + """Parse response as JSON.""" + return self._response.json() + + @property + def text(self) -> str: + """Response body as text.""" + return self._response.text + + def is_status(self, *codes: int) -> bool: + """Check if status matches any of the given codes.""" + return self._response.status_code in codes + + +class ApiClient: + """ + HTTP client for MCP servers to call the internal API. + + Provides: + - Connection pooling via shared httpx.AsyncClient + - Automatic URL building from endpoint paths + - Automatic agent header injection + - Standardized error response formatting + - Configurable timeouts + + Usage: + client = ApiClient(agent_id="be-dev-1") + + # GET request + resp = await client.get("/tasks/123") + if resp.ok: + task = resp.json() + + # POST request + resp = await client.post("/tasks", json={"title": "New task"}) + + # With custom timeout + resp = await client.get("/slow-endpoint", timeout=60.0) + """ + + def __init__( + self, + agent_id: str, + timeout: float = DEFAULT_TIMEOUT, + ) -> None: + """ + Initialize API client. + + Args: + agent_id: Agent identifier for header injection + timeout: Default request timeout in seconds + """ + self.agent_id = agent_id + self.timeout = timeout + self.base_url = settings.internal_api_url + self._client: httpx.AsyncClient | None = None + + def _get_headers(self) -> dict[str, str]: + """Get headers with agent context.""" + return get_agent_headers(self.agent_id) + + async def _ensure_client(self) -> httpx.AsyncClient: + """Get or create the httpx client.""" + if self._client is None: + self._client = httpx.AsyncClient(timeout=self.timeout) + return self._client + + async def close(self) -> None: + """Close the HTTP client. Call when done with the client.""" + if self._client: + await self._client.aclose() + self._client = None + + def _build_url(self, endpoint: str) -> str: + """Build full URL from endpoint path.""" + # Remove leading slash if present to avoid double slashes + if endpoint.startswith("/"): + endpoint = endpoint[1:] + return f"{self.base_url}/{endpoint}" + + async def get( + self, + endpoint: str, + params: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> ApiResponse: + """ + Make GET request. + + Args: + endpoint: API endpoint path (e.g., "/tasks/123") + params: Query parameters + timeout: Override default timeout + + Returns: + ApiResponse wrapper + """ + client = await self._ensure_client() + resp = await client.get( + self._build_url(endpoint), + params=params, + headers=self._get_headers(), + timeout=timeout or self.timeout, + ) + return ApiResponse(resp) + + async def post( + self, + endpoint: str, + json: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> ApiResponse: + """ + Make POST request. + + Args: + endpoint: API endpoint path + json: JSON body + params: Query parameters + timeout: Override default timeout + + Returns: + ApiResponse wrapper + """ + client = await self._ensure_client() + resp = await client.post( + self._build_url(endpoint), + json=json, + params=params, + headers=self._get_headers(), + timeout=timeout or self.timeout, + ) + return ApiResponse(resp) + + async def put( + self, + endpoint: str, + json: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> ApiResponse: + """ + Make PUT request. + + Args: + endpoint: API endpoint path + json: JSON body + timeout: Override default timeout + + Returns: + ApiResponse wrapper + """ + client = await self._ensure_client() + resp = await client.put( + self._build_url(endpoint), + json=json, + headers=self._get_headers(), + timeout=timeout or self.timeout, + ) + return ApiResponse(resp) + + async def patch( + self, + endpoint: str, + json: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> ApiResponse: + """ + Make PATCH request (partial update). + + Args: + endpoint: API endpoint path + json: JSON body with fields to update + timeout: Override default timeout + + Returns: + ApiResponse wrapper + """ + client = await self._ensure_client() + resp = await client.patch( + self._build_url(endpoint), + json=json, + headers=self._get_headers(), + timeout=timeout or self.timeout, + ) + return ApiResponse(resp) + + async def delete( + self, + endpoint: str, + timeout: float | None = None, + ) -> ApiResponse: + """ + Make DELETE request. + + Args: + endpoint: API endpoint path + timeout: Override default timeout + + Returns: + ApiResponse wrapper + """ + client = await self._ensure_client() + resp = await client.delete( + self._build_url(endpoint), + headers=self._get_headers(), + timeout=timeout or self.timeout, + ) + return ApiResponse(resp) + + # ========================================================================= + # CONVENIENCE METHODS + # ========================================================================= + + async def get_or_error( + self, + endpoint: str, + error_code: str = "API_ERROR", + error_message: str = "Request failed", + params: dict[str, Any] | None = None, + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """ + GET request returning (data, error) tuple. + + Args: + endpoint: API endpoint + error_code: Error code if request fails + error_message: Error message if request fails + params: Query parameters + + Returns: + (json_data, None) on success, (None, error_response) on failure + """ + try: + resp = await self.get(endpoint, params=params) + if resp.ok: + return resp.json(), None + return None, format_error_response( + error_code, + error_message, + {"status": resp.status_code, "detail": resp.text}, + ) + except Exception as e: + return None, format_error_response( + error_code, + error_message, + {"exception": str(e)}, + ) + + async def post_or_error( + self, + endpoint: str, + json: dict[str, Any] | None = None, + error_code: str = "API_ERROR", + error_message: str = "Request failed", + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """ + POST request returning (data, error) tuple. + + Args: + endpoint: API endpoint + json: JSON body + error_code: Error code if request fails + error_message: Error message if request fails + + Returns: + (json_data, None) on success, (None, error_response) on failure + """ + try: + resp = await self.post(endpoint, json=json) + if resp.ok: + return resp.json(), None + return None, format_error_response( + error_code, + error_message, + {"status": resp.status_code, "detail": resp.text}, + ) + except Exception as e: + return None, format_error_response( + error_code, + error_message, + {"exception": str(e)}, + )