pass roles + debugging

web/agent/Dockerfile

@@ -25,4 +25,4 @@ EXPOSE 8000
 ENV PORT=8000
 ENV HOST="0.0.0.0"
 
-CMD ["uv", "run", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]
+CMD ["uv", "run", "--frozen", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]

web/agent/pyproject.toml

@@ -5,9 +5,10 @@ description = "Cavepedia AI Agent with MCP tools"
 requires-python = ">=3.13"
 dependencies = [
     "uvicorn",
+    "starlette",
     "pydantic-ai",
-    "google-genai",
     "mcp",
+    "google-genai",
     "ag-ui-protocol",
     "python-dotenv",
     "httpx",

web/agent/src/agent.py

@@ -35,32 +35,52 @@ def check_mcp_available(url: str, timeout: float = 5.0) -> bool:
         logger.warning(f"MCP server not reachable: {e}")
         return False
 
-# Try to configure MCP if server is available
-toolsets = []
-if check_mcp_available(CAVE_MCP_URL):
-    try:
-        from pydantic_ai.mcp import MCPServerStreamableHTTP
-        mcp_server = MCPServerStreamableHTTP(
-            url=CAVE_MCP_URL,
-            timeout=30.0,
-        )
-        toolsets.append(mcp_server)
-        logger.info(f"MCP server configured: {CAVE_MCP_URL}")
-    except Exception as e:
-        logger.warning(f"Could not configure MCP server: {e}")
-else:
-    logger.info("MCP server unavailable - running without MCP tools")
+# Check if MCP is available at startup
+MCP_AVAILABLE = check_mcp_available(CAVE_MCP_URL)
+logger.info(f"MCP server available: {MCP_AVAILABLE}")
 
-# Create the agent with Google Gemini model
-agent = Agent(
-    model=GoogleModel("gemini-3-pro-preview"),
-    toolsets=toolsets if toolsets else None,
-    instructions="""You are a helpful caving assistant. Help users with all aspects of caving including cave exploration, safety, surveying techniques, cave locations, geology, equipment, history, conservation, and any other caving-related topics.
+AGENT_INSTRUCTIONS = """You are a helpful caving assistant. Help users with all aspects of caving including cave exploration, safety, surveying techniques, cave locations, geology, equipment, history, conservation, and any other caving-related topics.
 
 IMPORTANT RULES:
 1. Always cite your sources at the end of each response when possible.
 2. If you're not certain about information, say so clearly. Do NOT make up information or hallucinate facts.
-3. Provide accurate, helpful, and safety-conscious information.""",
-)
+3. Provide accurate, helpful, and safety-conscious information."""
 
-logger.info(f"Agent initialized successfully (MCP: {'enabled' if toolsets else 'disabled'})")
+
+def create_agent(user_roles: list[str] | None = None):
+    """Create an agent with MCP tools configured for the given user roles."""
+    toolsets = []
+
+    if MCP_AVAILABLE and user_roles:
+        try:
+            import json
+            from pydantic_ai.mcp import MCPServerStreamableHTTP
+
+            roles_header = json.dumps(user_roles)
+            logger.info(f"Creating MCP server with roles: {roles_header}")
+
+            mcp_server = MCPServerStreamableHTTP(
+                url=CAVE_MCP_URL,
+                headers={"x-user-roles": roles_header},
+                timeout=30.0,
+            )
+            toolsets.append(mcp_server)
+            logger.info(f"MCP server configured with roles: {user_roles}")
+        except Exception as e:
+            logger.warning(f"Could not configure MCP server: {e}")
+    elif not user_roles:
+        logger.info("No user roles provided - MCP tools disabled")
+    else:
+        logger.info("MCP server unavailable - running without MCP tools")
+
+    return Agent(
+        model=GoogleModel("gemini-2.5-flash"),
+        toolsets=toolsets if toolsets else None,
+        instructions=AGENT_INSTRUCTIONS,
+    )
+
+
+# Create a default agent for health checks etc
+agent = create_agent()
+
+logger.info("Agent module initialized successfully")

web/agent/src/main.py

@@ -4,6 +4,7 @@ Self-hosted PydanticAI agent server using AG-UI protocol.
 
 import os
 import sys
+import json
 import logging
 from dotenv import load_dotenv
 
@@ -24,13 +25,53 @@ if not os.getenv("GOOGLE_API_KEY"):
     sys.exit(1)
 
 import uvicorn
-from src.agent import agent
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import Response, JSONResponse
+from starlette.routing import Route
+
+from pydantic_ai.ui.ag_ui import AGUIAdapter
+
+from src.agent import create_agent
 
 logger.info("Creating AG-UI app...")
 
-# Convert PydanticAI agent to ASGI app with AG-UI protocol
-debug_mode = os.getenv("DEBUG", "").lower() in ("true", "1", "yes")
-app = agent.to_ag_ui(debug=debug_mode)
+
+async def handle_agent_request(request: Request) -> Response:
+    """Handle incoming AG-UI requests with dynamic role-based MCP configuration."""
+    # Debug: log all incoming headers
+    logger.info(f"DEBUG: All request headers: {dict(request.headers)}")
+
+    # Extract user roles from request headers
+    roles_header = request.headers.get("x-user-roles", "")
+    logger.info(f"DEBUG: x-user-roles header value: '{roles_header}'")
+    user_roles = []
+
+    if roles_header:
+        try:
+            user_roles = json.loads(roles_header)
+            logger.info(f"Request received with roles: {user_roles}")
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse x-user-roles header: {e}")
+
+    # Create agent with the user's roles
+    agent = create_agent(user_roles)
+
+    # Dispatch the request using AGUIAdapter
+    return await AGUIAdapter.dispatch_request(request, agent=agent)
+
+
+async def health(request: Request) -> Response:
+    """Health check endpoint."""
+    return JSONResponse({"status": "ok"})
+
+
+app = Starlette(
+    routes=[
+        Route("/", handle_agent_request, methods=["POST"]),
+        Route("/health", health, methods=["GET"]),
+    ],
+)
 
 logger.info("AG-UI app created successfully")
 

web/agent/uv.lock

@@ -231,6 +231,7 @@ dependencies = [
     { name = "mcp" },
     { name = "pydantic-ai" },
     { name = "python-dotenv" },
+    { name = "starlette" },
     { name = "uvicorn" },
 ]
 
@@ -242,6 +243,7 @@ requires-dist = [
     { name = "mcp" },
     { name = "pydantic-ai" },
     { name = "python-dotenv" },
+    { name = "starlette" },
     { name = "uvicorn" },
 ]