diff --git a/api/routes/bot.py b/api/routes/bot.py index 7d6bfa0..552a97e 100644 --- a/api/routes/bot.py +++ b/api/routes/bot.py @@ -8,11 +8,46 @@ from api.auth import auth_dependency, ws_auth_dependency from bot.agent_dispatch import dispatch_message from bot.message_handler import handle_message +from config.loader import config +from interactive_agent_layer.client import LayerClient router = APIRouter() logger = logging.getLogger(__name__) +async def _respond_to_permission(request_id, decision, user_id: str) -> None: + """POST the user's permission decision to the interactive-agent-layer. + + decision is "approve" or "deny" (chatStore.respondToPermission in the + frontend); mapped here to the boolean PermissionRespondRequest.approve + the layer expects (interactive_agent_layer/server.py: PermissionRespondRequest). + + Failures (layer unreachable, request_id already resolved/timed out, bad + request_id, etc.) are logged and swallowed — never raised to the caller. + The gate's own permission timeout (interactive_agent_layer/permissions.py) + is the backstop: a lost POST here just means the user waits for that + timeout to fire (denial) instead of getting an instant response. + """ + if not request_id: + logger.warning( + "bot_chat: permission_response missing request_id, dropping. user_id=%s", + user_id, + ) + return + + approve = decision == "approve" + base_url: str = config.get("agent_layer.base_url", "http://127.0.0.1:5003") + try: + layer = LayerClient(base_url) + await layer.respond_to_permission(request_id, approve) + except Exception as exc: + logger.warning( + "bot_chat: permission_response POST failed request_id=%s decision=%r" + " user_id=%s: %s", + request_id, decision, user_id, exc, + ) + + async def _cancel_and_wait(task: asyncio.Task, timeout: float = 5.0) -> None: """Cancel a task and wait for it to finish (best-effort, no raise). @@ -316,15 +351,43 @@ def capture_send_fn(msg: str) -> None: ) ) - # Race session completion against the next incoming client message. - recv_task = asyncio.create_task(websocket.receive_json()) - done, _pending = await asyncio.wait( - {session_task, recv_task}, - return_when=asyncio.FIRST_COMPLETED, - ) + # Race session completion against incoming client frames — re-arming + # recv_task after each one so MULTIPLE mid-session frames (e.g. a + # permission_response, possibly followed by another) are all + # consumed within this turn. The old single-shot fall-through only + # ever looked at the FIRST non-stop frame and then just awaited the + # session to completion — any second frame (like a second + # permission_response, or the response to the layer's permission + # request at all) was silently dropped. That drop is exactly why + # permission_response never reached the layer before this fix. + stopped = False + while not session_task.done(): + recv_task = asyncio.create_task(websocket.receive_json()) + done, _pending = await asyncio.wait( + {session_task, recv_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + + if recv_task not in done: + # session_task finished first — recv_task is still pending. + # Cancel and clean it up before leaving the loop. + recv_task.cancel() + try: + await recv_task + except asyncio.CancelledError: + pass + except WebSocketDisconnect: + # Client disconnected while awaiting: propagate so the + # outer handler can cancel session_task and return cleanly. + raise + except Exception as e: + logger.warning( + "bot_chat: unexpected error awaiting cancelled recv_task," + " user_id=%s: %s", + user_id, e, exc_info=True, + ) + break - # --- recv_task won (client sent something mid-session) --- - if recv_task in done and session_task not in done: try: incoming = recv_task.result() except Exception as exc: @@ -333,50 +396,38 @@ def capture_send_fn(msg: str) -> None: await _cancel_and_wait(session_task) raise exc - if incoming.get("type") == "stop": + itype = incoming.get("type") + + if itype == "stop": # Await cancellation before sending "Stopped." to guarantee no # further status frames from the session arrive after the ack. await _cancel_and_wait(session_task) await websocket.send_json({"type": "status", "content": "Stopped."}) await websocket.send_json({"type": "turn_complete", "final_text": "", "session_id": ""}) - session_task = None - continue # Back to outer loop — ready for next message + stopped = True + break + + if itype == "permission_response": + # Forward to the layer and keep looping — the session is + # still running, awaiting this (or another) decision. + await _respond_to_permission( + incoming.get("request_id"), incoming.get("decision"), user_id, + ) + continue - # Non-stop message mid-session (UI race / protocol error): - # drop it, let the session finish. recv_task is done (no orphan). + # Any other non-stop, non-permission_response frame (UI race or + # protocol error): drop it and keep looping so the session can + # still finish, or the client can still stop it / respond to a + # permission request afterwards. logger.warning( "bot_chat: mid-session message type=%r dropped (session continues)," " user_id=%s", - incoming.get("type"), user_id, + itype, user_id, ) - # Fall through to await session_task, with recv_task cleanup in finally. - # --- session_task won (or fell through from non-stop mid-session case) --- - # Await session if not yet done; always clean up recv_task. - try: - if not session_task.done(): - await session_task - finally: - # Cancel and clean up recv_task in all cases (normal, exception, cancel). - if not recv_task.done(): - recv_task.cancel() - try: - await recv_task - except asyncio.CancelledError: - pass - except WebSocketDisconnect: - # Client disconnected while awaiting session: propagate so - # the outer handler can cancel session_task and return cleanly. - raise - except Exception as e: - logger.warning( - "bot_chat: unexpected error awaiting cancelled recv_task," - " user_id=%s: %s", - user_id, e, exc_info=True, - ) - - if session_task is None or session_task.cancelled(): - continue + if stopped: + session_task = None + continue # Back to outer loop — ready for next message # Re-raise any exception from handle_message so error handlers below # can send a user-facing error frame and keep the connection alive. diff --git a/interactive_agent_layer/client.py b/interactive_agent_layer/client.py index 537aba5..e533d91 100644 --- a/interactive_agent_layer/client.py +++ b/interactive_agent_layer/client.py @@ -58,6 +58,21 @@ async def interrupt(self, session_id: str) -> None: resp = await self.client.post(f"{self.base_url}/session/{session_id}/interrupt") resp.raise_for_status() + async def respond_to_permission(self, request_id: str, approve: bool) -> None: + """POST the user's decision for a pending permission_request. + + Matches PermissionRespondRequest (interactive_agent_layer/server.py) — + body is exactly {"approve": bool}, nothing else. Raises on 404 (the + request already resolved, timed out, or never existed) or any other + HTTP error; callers should treat that as best-effort and swallow it — + the gate's own permission timeout is the backstop. + """ + resp = await self.client.post( + f"{self.base_url}/permission/{request_id}/respond", + json={"approve": approve}, + ) + resp.raise_for_status() + async def get_status(self, session_id: str) -> dict: resp = await self.client.get(f"{self.base_url}/session/{session_id}/status") resp.raise_for_status() diff --git a/tests/api/test_bot_ws_permission_response.py b/tests/api/test_bot_ws_permission_response.py new file mode 100644 index 0000000..c71ccc9 --- /dev/null +++ b/tests/api/test_bot_ws_permission_response.py @@ -0,0 +1,388 @@ +"""Tests for the bot_chat WebSocket permission_response link (brief 1d). + +Covers the keystone fix: {"type": "permission_response", request_id, decision} +frames sent by the client mid-session must reach the interactive-agent-layer's +/permission/{request_id}/respond endpoint (via LayerClient.respond_to_permission), +and the mid-session recv loop must be able to consume MULTIPLE such frames in +a single turn instead of dropping everything after the first non-stop frame. + +LayerClient is faked at the api.routes.bot import site (not real HTTP) — the +real POST body / endpoint contract is covered by +tests/interactive_agent_layer/test_client.py::test_respond_to_permission_*. +This file verifies bot.py maps decision -> approve correctly and drives the +loop correctly. +""" +from __future__ import annotations + +import asyncio +import os +from contextlib import asynccontextmanager +from unittest.mock import patch + +import pytest +from starlette.testclient import TestClient + +os.environ.setdefault("TETHER_DISABLE_RATE_LIMITS", "1") +os.environ.setdefault("TETHER_COOKIE_SECURE", "false") +os.environ.setdefault("TETHER_JWT_SECRET", "test-secret-for-ws-tests") + +from api.auth import create_jwt + +TEST_USER_ID = "00000000-0000-0000-0000-000000000099" +TEST_USERNAME = "ws_test_user" + + +class _FakePool: + pass + + +@asynccontextmanager +async def _noop_lifespan(app): + app.state.pool = _FakePool() + app.state.vault = None + yield + + +def _make_app(): + from api.main import create_app + return create_app(lifespan_override=_noop_lifespan) + + +def _valid_token(): + return create_jwt(TEST_USER_ID, TEST_USERNAME, is_admin=False) + + +class _FakeLayerClient: + """Records (base_url) at construction and (request_id, approve) per call. + + Tests register an asyncio.Event per request_id in `events` before the + permission_response frame is sent; respond_to_permission sets it so the + fake dispatch_message coroutine (which awaits that event) can resume — + this proves the call genuinely came from the server-side WS handler + processing the incoming frame, not test-side coordination. + """ + + instances: list = [] + events: dict = {} + + def __init__(self, base_url): + self.base_url = base_url + self.calls: list = [] + _FakeLayerClient.instances.append(self) + + async def respond_to_permission(self, request_id, approve): + self.calls.append((request_id, approve)) + ev = _FakeLayerClient.events.get(request_id) + if ev is not None: + ev.set() + + +class _RaisingLayerClient: + """Simulates the layer being unreachable — respond_to_permission always raises.""" + + def __init__(self, base_url): + pass + + async def respond_to_permission(self, request_id, approve): + raise RuntimeError("layer unreachable") + + +def setup_function(_fn): + _FakeLayerClient.instances = [] + _FakeLayerClient.events = {} + + +def test_permission_response_forwards_to_layer_and_resumes_session(): + """A single permission_request/permission_response round trip: + + dispatch_message emits permission_request via event_fn, blocks until the + approval lands, then completes the turn. bot_chat must forward the + permission_response frame to the layer (via LayerClient) instead of + dropping it, letting the session proceed. + """ + app = _make_app() + token = _valid_token() + approved = asyncio.Event() + _FakeLayerClient.events["req-1"] = approved + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + event_fn = kwargs["event_fn"] + await event_fn({ + "type": "permission_request", + "request_id": "req-1", + "kind": "user_action", + "target": "delete_tasks", + "session_id": "s1", + "reason_from_bot": None, + }) + await asyncio.wait_for(approved.wait(), timeout=5) + send_fn("proceeded after approval") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch), \ + patch("api.routes.bot.LayerClient", new=_FakeLayerClient): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "do something", + "agent_version": "tether-agent-2.0", + }) + + perm = ws.receive_json() + assert perm["type"] == "permission_request" + assert perm["request_id"] == "req-1" + + ws.send_json({ + "type": "permission_response", + "request_id": "req-1", + "decision": "approve", + }) + # Give the server-side handler a beat to invoke the fake layer + # client and set the event before we assert on it below. + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + assert tc["final_text"] == "proceeded after approval" + + assert len(_FakeLayerClient.instances) >= 1 + all_calls = [c for inst in _FakeLayerClient.instances for c in inst.calls] + assert ("req-1", True) in all_calls + + +def test_permission_response_deny_maps_to_approve_false(): + app = _make_app() + token = _valid_token() + approved = asyncio.Event() + _FakeLayerClient.events["req-deny"] = approved + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + event_fn = kwargs["event_fn"] + await event_fn({ + "type": "permission_request", + "request_id": "req-deny", + "kind": "destructive", + "target": "delete_context", + "session_id": "s1", + "reason_from_bot": None, + }) + await asyncio.wait_for(approved.wait(), timeout=5) + send_fn("handled denial") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch), \ + patch("api.routes.bot.LayerClient", new=_FakeLayerClient): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "do something risky", + "agent_version": "tether-agent-2.0", + }) + perm = ws.receive_json() + assert perm["request_id"] == "req-deny" + + ws.send_json({ + "type": "permission_response", + "request_id": "req-deny", + "decision": "deny", + }) + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + + all_calls = [c for inst in _FakeLayerClient.instances for c in inst.calls] + assert ("req-deny", False) in all_calls + + +def test_two_permission_responses_in_one_session(): + """Multiple mid-session frames must ALL be consumed in one turn — not just the first. + + This is the acceptance bar for the recv-loop restructure: the old + single-shot fall-through only ever looked at one mid-session frame per + turn, so a second permission_response would be silently dropped (or + misread by the next outer-loop iteration as a new user message). + """ + app = _make_app() + token = _valid_token() + approved_1 = asyncio.Event() + approved_2 = asyncio.Event() + _FakeLayerClient.events["req-1"] = approved_1 + _FakeLayerClient.events["req-2"] = approved_2 + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + event_fn = kwargs["event_fn"] + await event_fn({ + "type": "permission_request", "request_id": "req-1", + "kind": "user_action", "target": "a", "session_id": "s", + "reason_from_bot": None, + }) + await asyncio.wait_for(approved_1.wait(), timeout=5) + + await event_fn({ + "type": "permission_request", "request_id": "req-2", + "kind": "user_action", "target": "b", "session_id": "s", + "reason_from_bot": None, + }) + await asyncio.wait_for(approved_2.wait(), timeout=5) + + send_fn("both handled") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch), \ + patch("api.routes.bot.LayerClient", new=_FakeLayerClient): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "go", + "agent_version": "tether-agent-2.0", + }) + + p1 = ws.receive_json() + assert p1["request_id"] == "req-1" + ws.send_json({ + "type": "permission_response", + "request_id": "req-1", "decision": "approve", + }) + + p2 = ws.receive_json() + assert p2["request_id"] == "req-2" + ws.send_json({ + "type": "permission_response", + "request_id": "req-2", "decision": "deny", + }) + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + assert tc["final_text"] == "both handled" + + all_calls = [c for inst in _FakeLayerClient.instances for c in inst.calls] + assert ("req-1", True) in all_calls + assert ("req-2", False) in all_calls + + +def test_layer_error_logged_and_session_continues(): + """respond_to_permission raising must be swallowed — logged, not crashed. + + The gate's own permission timeout is the backstop; a failed POST here + should not tear down the WebSocket or the session. + """ + app = _make_app() + token = _valid_token() + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + event_fn = kwargs["event_fn"] + await event_fn({ + "type": "permission_request", "request_id": "req-err", + "kind": "user_action", "target": "x", "session_id": "s", + "reason_from_bot": None, + }) + # Session does not wait on approval here — simulates the gate's own + # timeout eventually firing and letting the turn finish regardless. + await asyncio.sleep(0.05) + send_fn("finished despite layer error") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch), \ + patch("api.routes.bot.LayerClient", new=_RaisingLayerClient): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "go", + "agent_version": "tether-agent-2.0", + }) + perm = ws.receive_json() + assert perm["request_id"] == "req-err" + + ws.send_json({ + "type": "permission_response", + "request_id": "req-err", "decision": "approve", + }) + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + assert tc["final_text"] == "finished despite layer error" + + # Connection must stay alive for another message. + ws.send_json({ + "type": "user", "content": "still alive?", + "agent_version": "tether-agent-2.0", + }) + # fake_dispatch runs again — will emit another permission_request + # first; drain it before the final turn_complete. + perm2 = ws.receive_json() + assert perm2["request_id"] == "req-err" + tc2 = ws.receive_json() + assert tc2["type"] == "turn_complete" + + +def test_permission_response_missing_request_id_dropped_not_crashed(): + """A malformed permission_response (no request_id) must not crash the handler.""" + app = _make_app() + token = _valid_token() + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + event_fn = kwargs["event_fn"] + await event_fn({ + "type": "permission_request", "request_id": "req-1", + "kind": "user_action", "target": "a", "session_id": "s", + "reason_from_bot": None, + }) + await asyncio.sleep(0.05) + send_fn("done") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch), \ + patch("api.routes.bot.LayerClient", new=_FakeLayerClient): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "go", + "agent_version": "tether-agent-2.0", + }) + perm = ws.receive_json() + assert perm["request_id"] == "req-1" + + # Malformed: no request_id at all. + ws.send_json({"type": "permission_response", "decision": "approve"}) + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + assert tc["final_text"] == "done" + + all_calls = [c for inst in _FakeLayerClient.instances for c in inst.calls] + assert all_calls == [] # never called — dropped before reaching the layer + + +def test_unknown_mid_session_frame_still_dropped_with_warning(): + """A frame type that isn't stop or permission_response keeps the old drop+warn behaviour.""" + app = _make_app() + token = _valid_token() + + async def fake_dispatch(agent_version, text, send_fn, pool, user_id, **kwargs): + await asyncio.sleep(0.05) + send_fn("done") + + with patch("api.routes.bot.dispatch_message", new=fake_dispatch): + with TestClient(app, raise_server_exceptions=False) as client: + with client.websocket_connect( + "/api/bot/chat", + cookies={"tether_token": token}, + ) as ws: + ws.send_json({ + "type": "user", "content": "go", + "agent_version": "tether-agent-2.0", + }) + ws.send_json({"type": "some_unknown_frame", "foo": "bar"}) + + tc = ws.receive_json() + assert tc["type"] == "turn_complete" + assert tc["final_text"] == "done" diff --git a/tests/interactive_agent_layer/test_client.py b/tests/interactive_agent_layer/test_client.py index 5058d6b..7778431 100644 --- a/tests/interactive_agent_layer/test_client.py +++ b/tests/interactive_agent_layer/test_client.py @@ -1,6 +1,8 @@ """Tests for LayerClient (contracts 7-8).""" from __future__ import annotations +import asyncio + import httpx import pytest from httpx import ASGITransport @@ -106,6 +108,47 @@ async def test_get_status(layer_client_with_app): assert status["user_id"] == "user4" +async def test_respond_to_permission_resolves_future_approve(layer_client_with_app, layer): + """LayerClient.respond_to_permission POSTs {"approve": True} and resolves the future.""" + lc = layer_client_with_app + sid = await lc.start_session( + user_id="user6", user_ws_id="wsid6", agent_version="v1", + options={}, user_message="approve me", + ) + session = layer.sessions[sid] + fut = asyncio.get_event_loop().create_future() + session.permission_pending["req-42"] = fut + + await lc.respond_to_permission("req-42", True) + + assert fut.done() + assert fut.result() is True + + +async def test_respond_to_permission_resolves_future_deny(layer_client_with_app, layer): + """LayerClient.respond_to_permission POSTs {"approve": False} and resolves the future.""" + lc = layer_client_with_app + sid = await lc.start_session( + user_id="user7", user_ws_id="wsid7", agent_version="v1", + options={}, user_message="deny me", + ) + session = layer.sessions[sid] + fut = asyncio.get_event_loop().create_future() + session.permission_pending["req-43"] = fut + + await lc.respond_to_permission("req-43", False) + + assert fut.done() + assert fut.result() is False + + +async def test_respond_to_permission_not_found_raises(layer_client_with_app): + """A request_id with no pending future (already resolved / never existed) raises.""" + lc = layer_client_with_app + with pytest.raises(httpx.HTTPStatusError): + await lc.respond_to_permission("no-such-request", True) + + async def test_turn_yields_events(layer_client_with_app): lc = layer_client_with_app sid = await lc.start_session(