1. roboco/mcp/utils.py - Created shared utilities:

- ApiClient class with connection pooling and automatic header injection
    - ApiResponse wrapper with .ok, .is_status(), .json() methods
    - get_agent_headers(), format_error_response(), resolve_agent_uuid()
2. roboco/mcp/schemas/__init__.py - Added new input schemas:
    - TaskBlockInput (groups blocker parameters)
    - TaskPauseInput (groups pause parameters)
3. Refactored all 4 MCP servers:
    - journal_server.py - Uses shared ApiClient
    - notify_server.py - Uses shared ApiClient
    - message_server.py - Uses shared ApiClient
    - task_server.py - Uses shared ApiClient, reduced argument count with dataclasses
This commit is contained in:
Renn F
2025-12-21 03:20:13 +01:00
parent b04993d8b8
commit 5d71f2fe9d
6 changed files with 1325 additions and 1289 deletions
+90 -174
View File
@@ -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
+130 -222
View File
@@ -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
+72 -121
View File
@@ -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
+22
View File
@@ -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"
)
+589 -772
View File
File diff suppressed because it is too large Load Diff
+422
View File
@@ -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)},
)