diff --git a/src/api/chat_routes.py b/src/api/chat_routes.py index 8047cb2..a0f1a22 100644 --- a/src/api/chat_routes.py +++ b/src/api/chat_routes.py @@ -58,6 +58,12 @@ from src.services.adk.memory import MemoryLimitExceeded from src.services.session_service import SessionLimitExceeded from src.middleware.permissions import RequirePermission +from src.services.permission_service import PermissionService + +# Runtime permission an agent-bot key must hold to invoke an agent over a chat +# WebSocket. WebSocket routes bypass the HTTP RequirePermission gate, so this +# confinement is enforced inline against the same allowlist. +WS_RUNTIME_PERMISSION = "ai_agent_processor.execute" import logging import json @@ -93,11 +99,18 @@ async def get_jwt_token_ws(token: str, skip_validation: bool = False) -> Optiona if auth_response and auth_response.user: # Return user context similar to what middleware does user = auth_response.user.dict() if hasattr(auth_response.user, 'dict') else auth_response.user + # Surface agent-bot identity when the token is bound to an agent + # (metadata.agent_id), mirroring EvoAuthMiddleware, so WebSocket + # handlers can enforce the same per-agent confinement. + metadata = getattr(auth_response, "metadata", None) or {} + bot_agent_id = metadata.get("agent_id") return { "sub": user.get("email") or user.get("id"), "email": user.get("email"), "user_id": user.get("id"), "user": user, + "is_agent_bot": bool(bot_agent_id), + "agent_id": bot_agent_id, } except Exception as e: logger.warning(f"EvoAuth token validation failed: {str(e)}") @@ -144,12 +157,14 @@ async def websocket_chat( # Verify authentication is_authenticated = False + auth_payload = None # Try with token (Bearer token from EvoAuth) if auth_data.get("token"): try: payload = await get_jwt_token_ws(auth_data["token"], True) if payload: + auth_payload = payload user_id = payload.get("user_id") or payload.get("sub") is_authenticated = True logger.info(f"WebSocket: User {user_id} authenticated for agent {agent_id}") @@ -165,6 +180,18 @@ async def websocket_chat( await websocket.close(code=status.WS_1008_POLICY_VIOLATION) return + # Confine agent-bot keys: a bot key may only invoke its OWN agent, and + # only for the runtime action. Reject at the handshake before any work. + if not PermissionService.is_agent_bot_permission_allowed( + auth_payload, websocket.url.path, WS_RUNTIME_PERMISSION + ): + logger.warning( + f"WebSocket: Agent Bot for agent {(auth_payload or {}).get('agent_id')} " + f"denied access to agent {agent_id}" + ) + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + return + logger.info(f"WebSocket authenticated for agent {agent_id}") # Main message loop @@ -402,6 +429,7 @@ async def websocket_live_chat( # Start authentication process logger.info("=== STARTING AUTHENTICATION PROCESS ===") is_authenticated = False + auth_payload = None # Try with JWT token first if auth_data.get("token"): @@ -414,6 +442,7 @@ async def websocket_live_chat( logger.info(f"JWT payload received: {bool(payload)}") if payload: + auth_payload = payload logger.info(f"JWT payload keys: {list(payload.keys())}") logger.info(f"JWT payload user: {payload.get('sub', 'N/A')}") logger.info(f"JWT payload email: {payload.get('email', 'N/A')}") @@ -514,6 +543,21 @@ async def websocket_live_chat( pass # Already disconnected return + # Confine agent-bot keys: a bot key may only invoke its OWN agent, and + # only for the runtime action. Reject at the handshake before any work. + if not PermissionService.is_agent_bot_permission_allowed( + auth_payload, websocket.url.path, WS_RUNTIME_PERMISSION + ): + logger.warning( + f"❌ Live WebSocket: Agent Bot for agent " + f"{(auth_payload or {}).get('agent_id')} denied access to agent {agent_id}" + ) + try: + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + except: + pass # Already disconnected + return + logger.info( f"✅ Live WebSocket authenticated successfully for agent {agent_id}" ) diff --git a/src/api/dependencies.py b/src/api/dependencies.py index 0c3c01e..6d27649 100644 --- a/src/api/dependencies.py +++ b/src/api/dependencies.py @@ -35,17 +35,27 @@ async def verify_agent_access( db: Session, agent: Any, # Agent object required_permission: str = "read", + user_context: Optional[Dict[str, Any]] = None, ) -> Tuple[bool, bool]: """ - Checks if the user has access to an agent, either by: - 1. Direct client ownership (admin or client user) - 2. Folder sharing permissions (for agents in shared folders) + Object-level access check for an agent. + + The Community box is single-tenant: agents carry no owner/account column, + so every authenticated user shares one agent pool and the coarse RBAC + lives in the route-level RequirePermission gates. What this enforces: + + - Agent-bot credentials (is_agent_bot) act only on their OWN agent: a + bot key for agent A is denied on agent B. + - A missing user_context fails closed, so call sites cannot skip the + check accidentally. + - Regular users keep pool access; the return flags it as shared access + when it flows through an active folder share for the user's email. Args: - #Removed for further handling - payload: JWT payload with user information db: Database session agent: Agent object to be checked required_permission: Required permission ("read" or "write") + user_context: Authenticated context set by EvoAuthMiddleware Returns: tuple: (has_access: bool, is_shared_access: bool) @@ -55,33 +65,46 @@ async def verify_agent_access( Raises: HTTPException: If access is denied """ - try: - return True, False # Access granted by direct ownership - except HTTPException as client_error: - # If direct access fails, check folder sharing - if agent.folder_id: - # Waiting for token implementation to get the user's email - user_email = None - if user_email: - has_folder_access = folder_share_service.check_folder_access( - db, agent.folder_id, user_email, required_permission + if not user_context: + logger.error("verify_agent_access called without user_context - denying access") + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Access denied", + ) + + if user_context.get("is_agent_bot"): + bot_agent_id = str(user_context.get("agent_id") or "") + if bot_agent_id and bot_agent_id == str(agent.id): + return True, False + + logger.warning( + f"Agent Bot for agent {bot_agent_id or ''} denied access to agent {agent.id}" + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Agent API Key can only access its own agent resources", + ) + + user_email = user_context.get("email") + if agent.folder_id and user_email: + try: + if folder_share_service.check_folder_access(db, agent.folder_id, user_email, required_permission): + logger.info( + f"User {user_email} granted {required_permission} access to agent {agent.id} " + f"via shared folder {agent.folder_id}" ) - if has_folder_access: - logger.info( - f"Usuário {user_email} recebeu acesso {required_permission} ao agente {agent.id} via pasta compartilhada {agent.folder_id}" - ) - return True, True - else: - logger.warning( - f"Usuário {user_email} negado ao agente {agent.id} - sem permissão de pasta compartilhada" - ) - else: - logger.warning("Nenhum e-mail de usuário encontrado no token para verificação de pasta compartilhada") - else: - logger.info( - f"Agente {agent.id} não está em uma pasta, não é possível verificar compartilhamento de pasta" - ) - raise client_error + return True, True + except HTTPException: + # An explicit denial from the share service must fail closed: it + # must not be swallowed and downgraded into a silent grant. + raise + except Exception as e: + # Transient infra errors do not revoke the single-tenant pool + # access every authenticated user already has; only the shared + # upgrade is withheld. + logger.warning(f"Folder share lookup failed for agent {agent.id}: {e}") + + return True, False def get_request_optional(request: Request) -> Request: """Dependency to provide the Request object, making it optional in endpoint signatures.""" diff --git a/src/api/session_routes.py b/src/api/session_routes.py index f710f63..dc3949c 100644 --- a/src/api/session_routes.py +++ b/src/api/session_routes.py @@ -185,9 +185,8 @@ async def create_new_session( status_code=status.HTTP_404_NOT_FOUND ) - # Verify access (skip for agent bots as they're already validated) - if not is_agent_bot: - has_access, is_shared_access = await verify_agent_access(db, agent, "read") + # Verify access (agent bots are constrained to their own agent inside) + has_access, is_shared_access = await verify_agent_access(db, agent, "read", current_user) # Get user identifier default_user_id = str(current_user.get("user_id") or current_user.get("email") or "") if current_user else "" @@ -373,7 +372,7 @@ async def get_agent_sessions( ) # Verify if the user has access to the agent (including shared folder access) - has_access, is_shared_access = await verify_agent_access(db, agent, "read") + has_access, is_shared_access = await verify_agent_access(db, agent, "read", current_user) # List ALL sessions for the agent (both test sessions and real sessions) @@ -448,7 +447,7 @@ async def bulk_delete_sessions( if agent: try: has_access, is_shared_access = await verify_agent_access( - db, agent, "read" + db, agent, "read", current_user ) validated_session_ids.append(session_id) except HTTPException as e: @@ -542,7 +541,7 @@ async def get_session( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "read" + db, agent, "read", current_user ) return success_response( @@ -595,7 +594,7 @@ async def get_agent_messages( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "read" + db, agent, "read", current_user ) # Get app_name and user_id from the session object instead of parsing session_id @@ -787,7 +786,7 @@ async def remove_session( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "read" + db, agent, "read", current_user ) # Delete the session (from both database and ADK) @@ -850,7 +849,7 @@ async def get_session_metadata_endpoint( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "read" + db, agent, "read", current_user ) # Get metadata @@ -907,7 +906,7 @@ async def update_session_metadata_endpoint( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "write" + db, agent, "write", current_user ) # Use user_id from the authenticated user @@ -978,7 +977,7 @@ async def delete_session_metadata_endpoint( agent = await agent_service.get_agent(db, agent_id) if agent: has_access, is_shared_access = await verify_agent_access( - db, agent, "write" + db, agent, "write", current_user ) # Use user_id from the authenticated user diff --git a/src/services/permission_service.py b/src/services/permission_service.py index 51d4a98..2d760e4 100644 --- a/src/services/permission_service.py +++ b/src/services/permission_service.py @@ -41,6 +41,21 @@ class PermissionService: across all authentication and authorization operations. """ + # Permissions an agent-bot key may exercise on its OWN agent. A bot key is + # a runtime credential: it exists to let the agent be invoked and to read + # its own conversational state. Management/credential actions (integrations + # connect/read/update/disconnect/create, tool config, session lifecycle + # mutation, agent deletion) are intentionally absent so a leaked bot key + # cannot delete the agent or read/write OAuth credentials. Fail closed: + # anything not listed here is denied even on the bot's own agent path. + _BOT_RUNTIME_PERMISSIONS = frozenset({ + "ai_agent_processor.execute", + "ai_a2a_protocol.execute", + "ai_a2a_protocol.read", + "ai_a2a_protocol.task_management", + "ai_chat_sessions.read", + }) + def __init__(self, evo_auth_base_url: str): self.evo_auth_service = EvoAuthService(evo_auth_base_url) logger.info(f"Permission service initialized with EvoAuthService") @@ -60,7 +75,46 @@ async def check_permission(self, auth_token: str, permission_key: str, token_typ """ return await self.evo_auth_service.check_permission(auth_token, permission_key, token_type) - async def validate_permission(self, request: Request, resource: str, action: str) -> None: + @staticmethod + def _path_scoped_to_agent(path: str, agent_id: str) -> bool: + """True when the request path targets the given agent: the agent_id is + a full path segment (/agents/{id}/..., /chat/{id}, /a2a/{id}/...) or the + trailing token of a session-id segment ({display_id}_{agent_id}). + + Canonical segment extraction (no substring matching), so an agent id + embedded mid-token in an unrelated value cannot spoof scope. + """ + if not agent_id: + return False + for segment in path.split("/"): + if not segment: + continue + if segment == agent_id: + return True + # Session ids embed the agent id as the suffix after the last "_". + if "_" in segment and segment.rsplit("_", 1)[-1] == agent_id: + return True + return False + + @classmethod + def is_agent_bot_permission_allowed( + cls, user_context: Optional[dict], path: str, permission_key: str + ) -> bool: + """Shared agent-bot confinement decision, usable outside the HTTP RBAC + gate (e.g. WebSocket handshakes that never reach validate_permission). + + Regular (non-bot) contexts are unaffected and always return True. An + agent-bot key is allowed only when the path targets its OWN agent and + the permission is a runtime action. Fail closed on everything else. + """ + if not user_context or not user_context.get("is_agent_bot"): + return True + bot_agent_id = str(user_context.get("agent_id") or "") + if not bot_agent_id or not cls._path_scoped_to_agent(path, bot_agent_id): + return False + return permission_key in cls._BOT_RUNTIME_PERMISSIONS + + async def validate_permission(self, request: Request, resource: str, action: str) -> None: # Build permission key permission_key = f"{resource}.{action}" @@ -82,11 +136,50 @@ async def validate_permission(self, request: Request, resource: str, action: str token_info = user_context.get("token_info", {}) token_type = token_info.get("type", "bearer") - # Agent Bots have full access (validated by middleware) + # Agent Bots bypass RBAC only on routes scoped to their OWN agent + # (the middleware validated the key against that agent) AND only for + # runtime permissions. A bot key must not act as a wildcard credential + # on other agents, nor perform management/credential actions (delete + # the agent, read/write OAuth credentials) even on its own agent. if user_context.get("is_agent_bot"): + bot_agent_id = str(user_context.get("agent_id") or "") + path = str(request.url.path) + scoped_to_own_agent = bool(bot_agent_id) and self._path_scoped_to_agent(path, bot_agent_id) + + if not scoped_to_own_agent: + logger.warning( + f"Permission: Agent Bot for agent {bot_agent_id or ''} denied " + f"{permission_key} on unscoped path {path}" + ) + raise HTTPException( + status_code=403, + detail={ + "error": "Insufficient permissions", + "code": "ERR_FORBIDDEN", + "message": "Agent API Key can only access its own agent resources", + "permission": permission_key + } + ) + + if permission_key not in self._BOT_RUNTIME_PERMISSIONS: + logger.warning( + f"Permission: Agent Bot for agent {bot_agent_id} denied non-runtime " + f"permission {permission_key}" + ) + raise HTTPException( + status_code=403, + detail={ + "error": "Insufficient permissions", + "code": "ERR_FORBIDDEN", + "message": "Agent API Key is limited to runtime actions on its own agent", + "permission": permission_key + } + ) + logger.info(f"Permission: Agent Bot access granted for permission {permission_key}") return - + + auth_token = token_info.get("access_token") if not auth_token: logger.error("Token not found in user context") diff --git a/tests/unit/test_agent_object_authz.py b/tests/unit/test_agent_object_authz.py new file mode 100644 index 0000000..7ce01bc --- /dev/null +++ b/tests/unit/test_agent_object_authz.py @@ -0,0 +1,331 @@ +"""Object-level authorization for agents. + +The Community box is single-tenant (agents have no owner column), so the +object boundary enforced here is: agent-bot credentials act only on their own +agent, a missing user context fails closed, regular users keep pool access +(flagged as shared when granted via folder share), and a bot token only +bypasses route RBAC on paths scoped to its own agent. +""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from src.api.dependencies import verify_agent_access +from src.services.permission_service import PermissionService + +AGENT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_AGENT_ID = "22222222-2222-2222-2222-222222222222" + + +def make_agent(agent_id=AGENT_ID, folder_id=None): + return SimpleNamespace(id=agent_id, folder_id=folder_id) + + +def run(coro): + return asyncio.run(coro) + + +class TestVerifyAgentAccess: + def test_missing_user_context_fails_closed(self): + with pytest.raises(HTTPException) as exc: + run(verify_agent_access(MagicMock(), make_agent(), "read", None)) + assert exc.value.status_code == 403 + + def test_agent_bot_can_access_its_own_agent(self): + context = {"is_agent_bot": True, "agent_id": AGENT_ID} + has_access, is_shared = run( + verify_agent_access(MagicMock(), make_agent(AGENT_ID), "read", context) + ) + assert has_access is True + assert is_shared is False + + def test_agent_bot_is_denied_on_another_agent(self): + context = {"is_agent_bot": True, "agent_id": AGENT_ID} + with pytest.raises(HTTPException) as exc: + run(verify_agent_access(MagicMock(), make_agent(OTHER_AGENT_ID), "write", context)) + assert exc.value.status_code == 403 + + def test_agent_bot_without_agent_id_is_denied(self): + context = {"is_agent_bot": True} + with pytest.raises(HTTPException) as exc: + run(verify_agent_access(MagicMock(), make_agent(), "read", context)) + assert exc.value.status_code == 403 + + def test_user_has_pool_access_without_folder(self): + context = {"is_agent_bot": False, "email": "user@example.com"} + has_access, is_shared = run( + verify_agent_access(MagicMock(), make_agent(), "read", context) + ) + assert has_access is True + assert is_shared is False + + def test_user_access_is_flagged_shared_via_folder_share(self): + context = {"is_agent_bot": False, "email": "user@example.com"} + agent = make_agent(folder_id="33333333-3333-3333-3333-333333333333") + with patch( + "src.api.dependencies.folder_share_service.check_folder_access", return_value=True + ): + has_access, is_shared = run(verify_agent_access(MagicMock(), agent, "read", context)) + assert has_access is True + assert is_shared is True + + def test_folder_share_lookup_failure_keeps_pool_access(self): + context = {"is_agent_bot": False, "email": "user@example.com"} + agent = make_agent(folder_id="33333333-3333-3333-3333-333333333333") + with patch( + "src.api.dependencies.folder_share_service.check_folder_access", + side_effect=RuntimeError("db down"), + ): + has_access, is_shared = run(verify_agent_access(MagicMock(), agent, "read", context)) + assert has_access is True + assert is_shared is False + + def test_folder_share_explicit_denial_fails_closed(self): + # An explicit HTTPException from the share service must propagate as a + # denial, not be swallowed into a silent grant. + context = {"is_agent_bot": False, "email": "user@example.com"} + agent = make_agent(folder_id="33333333-3333-3333-3333-333333333333") + with patch( + "src.api.dependencies.folder_share_service.check_folder_access", + side_effect=HTTPException(status_code=403, detail="denied"), + ): + with pytest.raises(HTTPException) as exc: + run(verify_agent_access(MagicMock(), agent, "read", context)) + assert exc.value.status_code == 403 + + +class TestAgentBotPermissionScope: + def make_service(self): + with patch("src.services.permission_service.EvoAuthService"): + return PermissionService("http://auth.test") + + def make_request(self, path, context): + request = MagicMock() + request.state.user_context = context + request.url.path = path + return request + + def test_path_scoped_to_agent_matches_segment_and_session_suffix(self): + assert PermissionService._path_scoped_to_agent( + f"/api/v1/agents/{AGENT_ID}/integrations/github/status", AGENT_ID + ) + assert PermissionService._path_scoped_to_agent(f"/api/v1/chat/{AGENT_ID}", AGENT_ID) + assert PermissionService._path_scoped_to_agent( + f"/api/v1/sessions/sync/display_{AGENT_ID}", AGENT_ID + ) + assert not PermissionService._path_scoped_to_agent("/api/v1/clients/usage", AGENT_ID) + assert not PermissionService._path_scoped_to_agent( + f"/api/v1/agents/{OTHER_AGENT_ID}/integrations/github/status", AGENT_ID + ) + + def test_path_scoped_rejects_mid_token_substring(self): + # The agent id embedded inside a larger token must not be treated as + # scope (canonical segment extraction, not substring matching). + assert not PermissionService._path_scoped_to_agent( + f"/api/v1/sessions/sync/prefix_{AGENT_ID}_suffix", AGENT_ID + ) + assert not PermissionService._path_scoped_to_agent( + f"/api/v1/agents/x{AGENT_ID}x/status", AGENT_ID + ) + assert not PermissionService._path_scoped_to_agent("/api/v1/agents//status", "") + + def test_bot_bypasses_rbac_for_runtime_action_on_its_own_agent(self): + service = self.make_service() + context = {"is_agent_bot": True, "agent_id": AGENT_ID, "token_info": {}} + + request = self.make_request(f"/api/v1/a2a/{AGENT_ID}", context) + assert run(service.validate_permission(request, "ai_a2a_protocol", "execute")) is None + + def test_bot_denied_credential_read_on_its_own_agent(self): + # Reading integration config exposes OAuth credentials; a runtime bot + # key must not reach it even on its own agent. + service = self.make_service() + context = {"is_agent_bot": True, "agent_id": AGENT_ID, "token_info": {}} + + request = self.make_request(f"/api/v1/agents/{AGENT_ID}/integrations/github/status", context) + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(request, "integrations", "read")) + assert exc.value.status_code == 403 + + def test_bot_denied_management_action_on_its_own_agent(self): + # Disconnecting an integration writes/destroys credentials — a + # management action outside the runtime allowlist. + service = self.make_service() + context = {"is_agent_bot": True, "agent_id": AGENT_ID, "token_info": {}} + + request = self.make_request(f"/api/v1/agents/{AGENT_ID}/integrations/github", context) + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(request, "integrations", "disconnect")) + assert exc.value.status_code == 403 + + def test_bot_is_denied_rbac_bypass_on_unscoped_path(self): + service = self.make_service() + context = {"is_agent_bot": True, "agent_id": AGENT_ID, "token_info": {}} + + request = self.make_request("/api/v1/clients/usage", context) + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(request, "ai_clients", "usage")) + assert exc.value.status_code == 403 + + def test_bot_is_denied_any_action_on_another_agents_integration_route(self): + # P1 confinement on an integration router: a bot key for agent A cannot + # act on agent B, for a runtime action or a management action. + service = self.make_service() + context = {"is_agent_bot": True, "agent_id": AGENT_ID, "token_info": {}} + + read_request = self.make_request( + f"/api/v1/agents/{OTHER_AGENT_ID}/integrations/github/status", context + ) + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(read_request, "integrations", "read")) + assert exc.value.status_code == 403 + + disconnect_request = self.make_request( + f"/api/v1/agents/{OTHER_AGENT_ID}/integrations/github", context + ) + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(disconnect_request, "integrations", "disconnect")) + assert exc.value.status_code == 403 + + +class TestUserPermissionOnIntegrationWrites: + """A regular user without the granted key is denied on gated routes — + the 403 path RequirePermission enforces on the integration credential + writes (save credentials / connect / disconnect).""" + + def make_service(self): + with patch("src.services.permission_service.EvoAuthService"): + service = PermissionService("http://auth.test") + return service + + def make_request(self, context): + request = MagicMock() + request.state.user_context = context + request.url.path = f"/api/v1/agents/{AGENT_ID}/integrations/github/credentials" + return request + + def user_context(self): + return { + "is_agent_bot": False, + "user_id": "u-1", + "token_info": {"access_token": "tok", "type": "bearer"}, + } + + def test_user_without_permission_gets_403(self): + from unittest.mock import AsyncMock + + service = self.make_service() + service.evo_auth_service.check_permission = AsyncMock(return_value=False) + + with pytest.raises(HTTPException) as exc: + run(service.validate_permission(self.make_request(self.user_context()), "integrations", "update")) + assert exc.value.status_code == 403 + assert exc.value.detail["permission"] == "integrations.update" + + def test_user_with_permission_passes(self): + from unittest.mock import AsyncMock + + service = self.make_service() + service.evo_auth_service.check_permission = AsyncMock(return_value=True) + + assert run(service.validate_permission(self.make_request(self.user_context()), "integrations", "update")) is None + + +RUNTIME_PERMISSION = "ai_agent_processor.execute" + + +class TestIsAgentBotPermissionAllowed: + """Shared confinement decision reused by the WebSocket handshake.""" + + def test_bot_allowed_runtime_on_own_agent(self): + context = {"is_agent_bot": True, "agent_id": AGENT_ID} + assert PermissionService.is_agent_bot_permission_allowed( + context, f"/chat/ws/{AGENT_ID}/user/sess", RUNTIME_PERMISSION + ) + + def test_bot_denied_on_another_agent(self): + context = {"is_agent_bot": True, "agent_id": OTHER_AGENT_ID} + assert not PermissionService.is_agent_bot_permission_allowed( + context, f"/chat/ws/{AGENT_ID}/user/sess", RUNTIME_PERMISSION + ) + + def test_bot_denied_non_runtime_on_own_agent(self): + context = {"is_agent_bot": True, "agent_id": AGENT_ID} + assert not PermissionService.is_agent_bot_permission_allowed( + context, f"/agents/{AGENT_ID}/integrations/github", "integrations.disconnect" + ) + + def test_bot_without_agent_id_denied(self): + context = {"is_agent_bot": True} + assert not PermissionService.is_agent_bot_permission_allowed( + context, f"/chat/ws/{AGENT_ID}/user/sess", RUNTIME_PERMISSION + ) + + def test_regular_user_unaffected(self): + assert PermissionService.is_agent_bot_permission_allowed( + {"is_agent_bot": False}, f"/chat/ws/{OTHER_AGENT_ID}/user/sess", RUNTIME_PERMISSION + ) + assert PermissionService.is_agent_bot_permission_allowed( + None, f"/chat/ws/{OTHER_AGENT_ID}/user/sess", RUNTIME_PERMISSION + ) + + +class TestChatWebSocketAgentBotConfinement: + """The chat WebSocket routes bypass the HTTP RequirePermission gate, so the + agent-bot confinement is enforced inline at the handshake.""" + + def build_client(self, monkeypatch, bot_agent_id): + import src.api.chat_routes as chat_routes + from fastapi import FastAPI + from fastapi.testclient import TestClient + from src.config.database import get_db + + fake_agent = SimpleNamespace(id=AGENT_ID, name="agent-a", folder_id=None, config=None) + monkeypatch.setattr( + chat_routes.agent_service, "get_agent", AsyncMock(return_value=fake_agent) + ) + monkeypatch.setattr( + chat_routes, + "get_jwt_token_ws", + AsyncMock(return_value={ + "is_agent_bot": True, + "agent_id": bot_agent_id, + "user_id": "bot", + }), + ) + + async def fake_stream(**kwargs): + yield json.dumps({"reply": "ok"}) + + monkeypatch.setattr(chat_routes, "run_agent_stream", fake_stream) + + app = FastAPI() + app.include_router(chat_routes.router) + app.dependency_overrides[get_db] = lambda: None + return TestClient(app) + + def test_bot_for_other_agent_rejected_at_handshake(self, monkeypatch): + from starlette.websockets import WebSocketDisconnect + + client = self.build_client(monkeypatch, bot_agent_id=OTHER_AGENT_ID) + with client.websocket_connect(f"/chat/ws/{AGENT_ID}/user/sess") as ws: + ws.send_json({"type": "authorization", "token": "t"}) + with pytest.raises(WebSocketDisconnect) as exc: + ws.receive_json() + assert exc.value.code == 1008 + + def test_bot_for_own_agent_accepted_and_runs(self, monkeypatch): + client = self.build_client(monkeypatch, bot_agent_id=AGENT_ID) + with client.websocket_connect(f"/chat/ws/{AGENT_ID}/user/sess") as ws: + ws.send_json({"type": "authorization", "token": "t"}) + ws.send_json({"message": "hi"}) + first = ws.receive_json() + assert first["turn_complete"] is False + assert first["message"] == {"reply": "ok"}