Fixed MCP connection and other MCP issues + API fixes

roboco/mcp/journal_server.py

@@ -20,6 +20,7 @@ import httpx
 from fastapi import status
 from mcp.server.fastmcp import FastMCP
 
+from roboco.agents_config import get_agent_role
 from roboco.config import settings
 from roboco.llm import ToonAdapter
 from roboco.mcp.schemas import (
@@ -54,6 +55,14 @@ def _format_error_response(
     }
 
 
+def _get_agent_headers(agent_id: str) -> dict[str, str]:
+    """Get standard headers for API calls."""
+    return {
+        "X-Agent-Id": agent_id,
+        "X-Agent-Role": get_agent_role(agent_id),
+    }
+
+
 async def _post_journal_entry(
     endpoint: str,
     payload: dict[str, Any],
@@ -64,7 +73,7 @@ async def _post_journal_entry(
         resp = await client.post(
             f"{settings.internal_api_url}/journals/me/{endpoint}",
             json=payload,
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
         if resp.status_code not in [200, 201]:
             return None, _format_error_response(
@@ -221,7 +230,7 @@ async def _handle_search(query: str, top_k: int, agent_id: str) -> dict[str, Any
         resp = await client.post(
             f"{settings.internal_api_url}/journals/me/search",
             json=payload,
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code != status.HTTP_200_OK:
@@ -249,11 +258,11 @@ async def _handle_stats(agent_id: str) -> dict[str, Any]:
     async with httpx.AsyncClient() as client:
         stats_resp = await client.get(
             f"{settings.internal_api_url}/journals/me/stats",
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
         growth_resp = await client.get(
             f"{settings.internal_api_url}/journals/me/growth",
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         stats = (
@@ -299,7 +308,7 @@ async def _handle_recent(
         resp = await client.get(
             f"{settings.internal_api_url}/journals/me/entries",
             params=params,
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code != status.HTTP_200_OK:

roboco/mcp/message_server.py

@@ -20,7 +20,7 @@ import httpx
 from fastapi import status
 from mcp.server.fastmcp import FastMCP
 
-from roboco.agents_config import CHANNEL_ACCESS
+from roboco.agents_config import CHANNEL_ACCESS, get_agent_role
 from roboco.config import settings
 from roboco.llm import ToonAdapter
 from roboco.mcp.schemas import (
@@ -33,6 +33,14 @@ from roboco.mcp.schemas import (
 _toon = ToonAdapter()
 
 
+def _get_agent_headers(agent_id: str) -> dict[str, str]:
+    """Get standard headers for API calls."""
+    return {
+        "X-Agent-Id": agent_id,
+        "X-Agent-Role": get_agent_role(agent_id),
+    }
+
+
 # =============================================================================
 # HELPER FUNCTIONS
 # =============================================================================
@@ -112,26 +120,76 @@ def _validate_message_send(
     return None
 
 
+async def _get_default_group(
+    client: httpx.AsyncClient,
+    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,
+    )
+
+    if groups_resp.status_code != status.HTTP_200_OK:
+        return _format_error_response(
+            "GROUPS_ERROR",
+            "Failed to get channel groups",
+            {"status": groups_resp.status_code},
+        )
+
+    groups = groups_resp.json()
+    if not 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:
+        if group.get("is_active", True):
+            return str(group["id"])
+    return str(groups[0]["id"])
+
+
 async def _get_or_create_session(
     client: httpx.AsyncClient,
     channel_id: str,
+    headers: dict[str, str],
 ) -> str | dict[str, Any]:
     """Get or create session for channel. Returns session_id or error dict."""
-    session_resp = await client.get(
-        f"{settings.internal_api_url}/channels/{channel_id}/session"
+    # First get the default group for this channel
+    group_result = await _get_default_group(client, channel_id, headers)
+    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,
     )
 
-    if session_resp.status_code == status.HTTP_200_OK:
-        return str(session_resp.json()["id"])
+    if sessions_resp.status_code == status.HTTP_200_OK:
+        data = sessions_resp.json()
+        items = data.get("items", [])
+        # Find an active session
+        for session in items:
+            if session.get("status") == "active":
+                return str(session["id"])
 
+    # Create new session
     create_resp = await client.post(
         f"{settings.internal_api_url}/sessions",
-        json={"channel_id": channel_id},
+        json={"group_id": group_id},
+        headers=headers,
     )
     if create_resp.status_code in [status.HTTP_200_OK, status.HTTP_201_CREATED]:
         return str(create_resp.json()["id"])
 
-    return _format_error_response("SESSION_ERROR", "Failed to get or create session")
+    return _format_error_response(
+        "SESSION_ERROR",
+        "Failed to create session",
+        {"api_error": create_resp.text},
+    )
 
 
 # =============================================================================
@@ -164,7 +222,7 @@ async def _handle_channel_list(agent_id: str) -> dict[str, Any]:
     }
 
 
-async def _handle_channel_history(
+async def _handle_channel_history(  # noqa: PLR0911
     agent_id: str,
     channel_slug: str,
     limit: int,
@@ -179,38 +237,86 @@ async def _handle_channel_history(
     limit = min(limit, 100)
     since = datetime.now(UTC) - timedelta(hours=hours_back)
 
+    headers = _get_agent_headers(agent_id)
     async with httpx.AsyncClient() as client:
+        # Get channel by slug
         channels_resp = await client.get(
             f"{settings.internal_api_url}/channels",
             params={"slug": channel_slug},
+            headers=headers,
         )
 
         if channels_resp.status_code != status.HTTP_200_OK:
             return _format_error_response("API_ERROR", "Failed to fetch channels")
 
-        channels = channels_resp.json()
-        if not channels:
+        channels_data = channels_resp.json()
+        items = channels_data.get("items", channels_data)
+        if not items:
             return _format_error_response(
                 "NOT_FOUND", f"Channel #{channel_slug} not found"
             )
 
-        channel_id = channels[0]["id"]
+        channel = items[0] if isinstance(items, list) else items
+        channel_id = channel["id"]
 
-        messages_resp = await client.get(
-            f"{settings.internal_api_url}/channels/{channel_id}/messages",
-            params={"after": since.isoformat(), "limit": limit},
+        # Get groups for this channel
+        group_result = await _get_default_group(client, channel_id, headers)
+        if isinstance(group_result, dict):
+            return group_result  # Error response
+        group_id = group_result
+
+        # Get sessions for this group
+        sessions_resp = await client.get(
+            f"{settings.internal_api_url}/sessions",
+            params={"group_id": group_id, "limit": 5},
+            headers=headers,
         )
 
-        if messages_resp.status_code != status.HTTP_200_OK:
-            return _format_error_response("API_ERROR", "Failed to fetch messages")
+        if sessions_resp.status_code != status.HTTP_200_OK:
+            return _format_error_response("API_ERROR", "Failed to fetch sessions")
 
-        messages = messages_resp.json()
+        sessions_data = sessions_resp.json()
+        sessions = sessions_data.get("items", [])
+
+        if not sessions:
+            return {
+                "channel": channel_slug,
+                "messages": [],
+                "total": 0,
+                "has_more": False,
+                "since": since.isoformat(),
+            }
+
+        # Get messages from all recent sessions
+        all_messages = []
+        for session in sessions:
+            session_id = session["id"]
+            messages_resp = await client.get(
+                f"{settings.internal_api_url}/messages",
+                params={
+                    "session_id": session_id,
+                    "after": since.isoformat(),
+                    "limit": limit,
+                },
+                headers=headers,
+            )
+
+            if messages_resp.status_code == status.HTTP_200_OK:
+                msg_data = messages_resp.json()
+                all_messages.extend(msg_data.get("items", []))
+
+            if len(all_messages) >= limit:
+                break
+
+        # Sort by timestamp descending and limit
+        all_messages.sort(key=lambda m: m.get("timestamp", ""), reverse=True)
+        all_messages = all_messages[:limit]
 
     return {
         "channel": channel_slug,
-        "messages": messages.get("items", []),
-        "total": messages.get("total", 0),
-        "has_more": messages.get("has_more", False),
+        "messages": all_messages,
+        "total": len(all_messages),
+        "has_more": len(all_messages) >= limit,
         "since": since.isoformat(),
     }
 
@@ -225,10 +331,12 @@ async def _handle_message_send(
     ):
         return validation_error
 
+    headers = _get_agent_headers(agent_id)
     async with httpx.AsyncClient() as client:
         channels_resp = await client.get(
             f"{settings.internal_api_url}/channels",
             params={"slug": data.channel_slug},
+            headers=headers,
         )
 
         if channels_resp.status_code != status.HTTP_200_OK or not channels_resp.json():
@@ -239,7 +347,7 @@ async def _handle_message_send(
         channel = channels_resp.json()[0]
         channel_id = channel["id"]
 
-        session_result = await _get_or_create_session(client, channel_id)
+        session_result = await _get_or_create_session(client, channel_id, headers)
         if isinstance(session_result, dict):
             return session_result
         session_id = session_result
@@ -257,7 +365,7 @@ async def _handle_message_send(
         send_resp = await client.post(
             f"{settings.internal_api_url}/messages",
             json=message_data,
-            headers={"X-Agent-Id": agent_id},
+            headers=headers,
         )
 
         if send_resp.status_code not in [status.HTTP_200_OK, status.HTTP_201_CREATED]:
@@ -273,10 +381,14 @@ async def _handle_message_send(
         }
 
 
-async def _handle_message_get(message_id: str) -> dict[str, Any]:
+async def _handle_message_get(message_id: str, agent_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}")
+        resp = await client.get(
+            f"{settings.internal_api_url}/messages/{message_id}",
+            headers=headers,
+        )
 
         if resp.status_code == status.HTTP_404_NOT_FOUND:
             return _format_error_response(
@@ -396,7 +508,7 @@ def create_message_mcp_server(agent_id: str) -> FastMCP:
     @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)
+        return await _handle_message_get(message_id, agent_id)
 
     @mcp.tool()
     async def roboco_ask_question(

roboco/mcp/notify_server.py

@@ -30,6 +30,14 @@ from roboco.mcp.schemas import SendNotificationInput
 # =============================================================================
 
 
+def _get_agent_headers(agent_id: str) -> dict[str, str]:
+    """Get standard headers for API calls."""
+    return {
+        "X-Agent-Id": agent_id,
+        "X-Agent-Role": get_agent_role(agent_id),
+    }
+
+
 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)
@@ -126,7 +134,7 @@ async def _handle_list(
         resp = await client.get(
             f"{settings.internal_api_url}/notifications",
             params=params,
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code != status.HTTP_200_OK:
@@ -162,7 +170,7 @@ async def _handle_get(agent_id: str, notification_id: str) -> dict[str, Any]:
     async with httpx.AsyncClient() as client:
         resp = await client.get(
             f"{settings.internal_api_url}/notifications/{notification_id}",
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code == status.HTTP_404_NOT_FOUND:
@@ -193,7 +201,7 @@ async def _handle_ack(agent_id: str, notification_id: str) -> dict[str, Any]:
     async with httpx.AsyncClient() as client:
         resp = await client.post(
             f"{settings.internal_api_url}/notifications/{notification_id}/ack",
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code == status.HTTP_404_NOT_FOUND:
@@ -280,7 +288,7 @@ async def _handle_send(agent_id: str, data: SendNotificationInput) -> dict[str,
         resp = await client.post(
             f"{settings.internal_api_url}/notifications",
             json=payload,
-            headers={"X-Agent-Id": agent_id},
+            headers=_get_agent_headers(agent_id),
         )
 
         if resp.status_code not in [status.HTTP_200_OK, status.HTTP_201_CREATED]:

roboco/mcp/task_server.py

@@ -26,9 +26,55 @@ import httpx
 from fastapi import status
 from mcp.server.fastmcp import FastMCP
 
+from roboco.agents_config import get_agent_role
 from roboco.config import settings
 from roboco.llm import ToonAdapter
 
+
+def _get_agent_headers(agent_id: str) -> dict[str, str]:
+    """Get standard headers for API calls."""
+    return {
+        "X-Agent-Id": agent_id,
+        "X-Agent-Role": get_agent_role(agent_id),
+    }
+
+
+# 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
+    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
+
+
 # Global TOON adapter for encoding task data
 _toon = ToonAdapter()
 
@@ -145,11 +191,13 @@ def _get_next_step_guidance(status: str) -> tuple[str, str]:
 
 async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]:
     """Handle task scanning."""
+    headers = _get_agent_headers(agent_id)
     async with httpx.AsyncClient() as client:
         # Get paused tasks for this agent
         paused_resp = await client.get(
             f"{settings.internal_api_url}/tasks",
             params={"assigned_to": agent_id, "status": "paused"},
+            headers=headers,
         )
         paused_tasks = (
             paused_resp.json() if paused_resp.status_code == status.HTTP_200_OK else []
@@ -159,6 +207,7 @@ async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]:
         assigned_resp = await client.get(
             f"{settings.internal_api_url}/tasks",
             params={"assigned_to": agent_id},
+            headers=headers,
         )
         assigned_data = (
             assigned_resp.json()
@@ -179,6 +228,7 @@ async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]:
         available_resp = await client.get(
             f"{settings.internal_api_url}/tasks",
             params=params,
+            headers=headers,
         )
         available_tasks = (
             available_resp.json()
@@ -216,10 +266,14 @@ async def _handle_task_scan(team: str | None, agent_id: str) -> dict[str, Any]:
     }
 
 
-async def _handle_task_get(task_id: str) -> dict[str, Any]:
+async def _handle_task_get(task_id: str, agent_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}")
+        resp = await client.get(
+            f"{settings.internal_api_url}/tasks/{task_id}",
+            headers=headers,
+        )
 
         if resp.status_code == status.HTTP_404_NOT_FOUND:
             return _format_error_response(
@@ -272,10 +326,14 @@ def _validate_task_claimable(task: dict) -> dict[str, Any] | None:
     return None
 
 
-async def _get_project_context(project_id: str) -> dict[str, Any] | None:
+async def _get_project_context(project_id: str, agent_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}")
+        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
@@ -284,10 +342,12 @@ async def _get_project_context(project_id: str) -> dict[str, Any] | None:
 
 async def _handle_task_claim(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",
             params={"assigned_to": agent_id},
+            headers=headers,
         )
         if active_resp.status_code == status.HTTP_200_OK:
             active_tasks = active_resp.json()
@@ -296,7 +356,10 @@ async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]:
             if error := _check_paused_tasks(active_tasks):
                 return error
 
-        task_resp = await client.get(f"{settings.internal_api_url}/tasks/{task_id}")
+        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")
 
@@ -307,6 +370,7 @@ async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]:
         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(
@@ -319,7 +383,7 @@ async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]:
 
     project = None
     if claimed_task.get("project_id"):
-        project = await _get_project_context(claimed_task["project_id"])
+        project = await _get_project_context(claimed_task["project_id"], agent_id)
 
     return _format_task_response(
         claimed_task,
@@ -332,13 +396,30 @@ async def _handle_task_claim(task_id: str, agent_id: str) -> dict[str, Any]:
     )
 
 
-def _validate_task_ownership(task: dict, agent_id: str) -> dict[str, Any] | None:
+async def _validate_task_ownership(
+    task: dict, agent_id: str, headers: dict[str, str]
+) -> dict[str, Any] | None:
     """Validate agent owns the task. Returns error or None."""
-    if task.get("assigned_to") != agent_id:
+    assigned_to = task.get("assigned_to")
+    if not assigned_to:
+        return _format_error_response(
+            "NOT_ASSIGNED",
+            "This task is not assigned to anyone",
+        )
+
+    # Resolve agent_id (which may be a slug) to UUID for comparison
+    agent_uuid = await _resolve_agent_uuid(agent_id, headers)
+    if not agent_uuid:
+        return _format_error_response(
+            "AGENT_NOT_FOUND",
+            f"Could not resolve agent: {agent_id}",
+        )
+
+    if str(assigned_to) != agent_uuid:
         return _format_error_response(
             "NOT_OWNER",
             "You are not assigned to this task",
-            {"assigned_to": task.get("assigned_to")},
+            {"assigned_to": assigned_to},
         )
     return None
 
@@ -381,13 +462,17 @@ async def _handle_task_plan(
     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}")
+        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()
-        if error := _validate_task_ownership(task, agent_id):
+        if error := await _validate_task_ownership(task, agent_id, headers):
             return error
         if error := _validate_task_status_claimed(task):
             return error
@@ -396,6 +481,7 @@ async def _handle_task_plan(
         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(
@@ -423,10 +509,12 @@ async def _handle_task_plan(
     )
 
 
-def _validate_task_start(task: dict[str, Any], agent_id: str) -> dict[str, Any] | None:
+async def _validate_task_start(
+    task: dict[str, Any], agent_id: str, headers: dict[str, str]
+) -> dict[str, Any] | None:
     """Validate task can be started. Returns error dict or None."""
-    if task.get("assigned_to") != agent_id:
-        return _format_error_response("NOT_OWNER", "You are not assigned to this task")
+    if error := await _validate_task_ownership(task, agent_id, headers):
+        return error
 
     task_status = task.get("status")
     if task_status not in ["claimed", "paused"]:
@@ -459,19 +547,24 @@ def _validate_task_start(task: dict[str, Any], agent_id: str) -> dict[str, Any]
 
 async def _handle_task_start(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}")
+        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()
 
-        if validation_error := _validate_task_start(task, agent_id):
+        if validation_error := await _validate_task_start(task, agent_id, headers):
             return validation_error
 
         # Start the task
         start_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/start"
+            f"{settings.internal_api_url}/tasks/{task_id}/start",
+            headers=headers,
         )
 
         if start_resp.status_code != status.HTTP_200_OK:
@@ -500,17 +593,19 @@ async def _handle_task_progress(
     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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "in_progress":
             return _format_error_response(
@@ -526,6 +621,7 @@ async def _handle_task_progress(
                 "message": message,
                 "percentage": percentage,
             },
+            headers=headers,
         )
 
         if progress_resp.status_code != status.HTTP_200_OK:
@@ -557,17 +653,19 @@ async def _handle_task_block(
             "Both 'reason' and 'what_needed' are required to block a task.",
         )
 
+    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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "in_progress":
             return _format_error_response(
@@ -583,6 +681,7 @@ async def _handle_task_block(
                 "blocker_type": blocker_type,
                 "what_needed": what_needed,
             },
+            headers=headers,
         )
 
         if block_resp.status_code != status.HTTP_200_OK:
@@ -605,17 +704,19 @@ async def _handle_task_block(
 
 async def _handle_task_unblock(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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "blocked":
             return _format_error_response(
@@ -624,7 +725,8 @@ async def _handle_task_unblock(task_id: str, agent_id: str) -> dict[str, Any]:
             )
 
         unblock_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/unblock"
+            f"{settings.internal_api_url}/tasks/{task_id}/unblock",
+            headers=headers,
         )
 
         if unblock_resp.status_code != status.HTTP_200_OK:
@@ -647,17 +749,19 @@ async def _handle_task_pause(
     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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "in_progress":
             return _format_error_response(
@@ -674,11 +778,13 @@ async def _handle_task_pause(
                 "remaining_work": remaining_work,
                 "notes": reason,
             },
+            headers=headers,
         )
 
         # Pause the task
         pause_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/pause"
+            f"{settings.internal_api_url}/tasks/{task_id}/pause",
+            headers=headers,
         )
 
         if pause_resp.status_code != status.HTTP_200_OK:
@@ -700,17 +806,19 @@ async def _handle_task_submit_verification(
     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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "in_progress":
             return _format_error_response(
@@ -728,7 +836,8 @@ async def _handle_task_submit_verification(
             )
 
         verify_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/verify"
+            f"{settings.internal_api_url}/tasks/{task_id}/verify",
+            headers=headers,
         )
 
         if verify_resp.status_code != status.HTTP_200_OK:
@@ -765,17 +874,19 @@ 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}")
+        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()
 
-        if task.get("assigned_to") != agent_id:
-            return _format_error_response(
-                "NOT_OWNER", "You are not assigned to this task"
-            )
+        if error := await _validate_task_ownership(task, agent_id, headers):
+            return error
 
         if task.get("status") != "verifying":
             return _format_error_response(
@@ -790,11 +901,13 @@ async def _handle_task_submit_qa(
                 "dev_notes": dev_notes,
                 "documenter_handoff": handoff_summary,
             },
+            headers=headers,
         )
 
         # Submit for QA
         qa_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/submit-qa"
+            f"{settings.internal_api_url}/tasks/{task_id}/submit-qa",
+            headers=headers,
         )
 
         if qa_resp.status_code != status.HTTP_200_OK:
@@ -824,8 +937,12 @@ 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}")
+        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")
 
@@ -847,6 +964,7 @@ async def _handle_task_qa_pass(
         pass_resp = await client.post(
             f"{settings.internal_api_url}/tasks/{task_id}/pass-qa",
             json={"notes": qa_notes},
+            headers=headers,
         )
 
         if pass_resp.status_code != status.HTTP_200_OK:
@@ -881,8 +999,12 @@ 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}")
+        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")
 
@@ -899,6 +1021,7 @@ async def _handle_task_qa_fail(
         fail_resp = await client.post(
             f"{settings.internal_api_url}/tasks/{task_id}/fail-qa",
             json={"notes": full_notes},
+            headers=headers,
         )
 
         if fail_resp.status_code != status.HTTP_200_OK:
@@ -915,10 +1038,14 @@ async def _handle_task_qa_fail(
     )
 
 
-async def _handle_task_complete(task_id: str) -> dict[str, Any]:
+async def _handle_task_complete(task_id: str, agent_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}")
+        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")
 
@@ -931,7 +1058,8 @@ async def _handle_task_complete(task_id: str) -> dict[str, Any]:
             )
 
         complete_resp = await client.post(
-            f"{settings.internal_api_url}/tasks/{task_id}/complete"
+            f"{settings.internal_api_url}/tasks/{task_id}/complete",
+            headers=headers,
         )
 
         if complete_resp.status_code != status.HTTP_200_OK:
@@ -948,11 +1076,13 @@ async def _handle_task_complete(task_id: str) -> dict[str, Any]:
 
 async def _handle_agent_idle(agent_id: str) -> dict[str, Any]:
     """Handle agent going idle (no work available)."""
+    headers = _get_agent_headers(agent_id)
     async with httpx.AsyncClient() as client:
         # 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,
         )
 
         if resp.status_code == status.HTTP_204_NO_CONTENT:
@@ -1022,7 +1152,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)
+        return await _handle_task_get(task_id, agent_id)
 
     @mcp.tool()
     async def roboco_task_claim(task_id: str) -> dict[str, Any]:
@@ -1288,7 +1418,7 @@ def create_task_mcp_server(agent_id: str) -> FastMCP:
         Returns:
             Completed task
         """
-        return await _handle_task_complete(task_id)
+        return await _handle_task_complete(task_id, agent_id)
 
     @mcp.tool()
     async def roboco_agent_idle() -> dict[str, Any]: