NetDevAutomate
diff --git a/‎packages/studyloop/src/studyloop/web/routes/session.py‎
Lines changed: 154 additions & 63 deletions b/‎packages/studyloop/src/studyloop/web/routes/session.py‎
Lines changed: 154 additions & 63 deletions
diff --git a/‎packages/studyloop/tests/test_web_live_session.py‎
Lines changed: 8 additions & 73 deletions b/‎packages/studyloop/tests/test_web_live_session.py‎
Lines changed: 8 additions & 73 deletions
@@ -576,79 +576,170 @@ def get_session_options() -> dict[str, list[dict]]:
     }
 
 
+_WS_CLOSE_POLICY = 1008  # RFC 6455 Policy Violation
+
+
+def _allowed_ws_origins() -> set[str]:
+    """Return the allow-list for Origin on the live session WS.
+
+    Localhost only by default — this is a single-user tool. Expand via
+    ``STUDYLOOP_ALLOWED_ORIGINS`` (comma-separated) if the operator fronts
+    the app behind a reverse proxy. See plan Blocker B1 (CSRF via cross-
+    origin WS upgrade; RFC 6455 has no same-origin enforcement and CORS
+    does not apply to WS).
+    """
+    import os
+
+    defaults = {
+        "http://127.0.0.1:8788",
+        "http://localhost:8788",
+        "http://127.0.0.1",
+        "http://localhost",
+    }
+    extra = os.environ.get("STUDYLOOP_ALLOWED_ORIGINS", "").strip()
+    if extra:
+        return defaults | {o.strip() for o in extra.split(",") if o.strip()}
+    return defaults
+
+
+def _origin_allowed(origin: str) -> bool:
+    if not origin:
+        return False
+    allowed = _allowed_ws_origins()
+    if origin in allowed:
+        return True
+    # Allow any port on localhost / 127.0.0.1 (dev server port may vary).
+    for prefix in ("http://127.0.0.1", "http://localhost"):
+        if origin.startswith(prefix + ":") or origin == prefix:
+            return True
+    return False
+
+
 @router.websocket("/session/ws")
 async def live_session_socket(websocket: WebSocket) -> None:
-    """Bidirectional live agent session socket.
-
-    Client messages:
-    - {"type": "start", "topic": "...", "energy": 5, "agent": "codex", "transport": "pty"}
-    - {"type": "input", "text": "..."}
-    - {"type": "stop"}
+    """Bidirectional live agent session socket (plan §1.5).
+
+    The route binds an already-acquired ``active.ActiveSession`` to the
+    WS client. Sessions are acquired via ``POST /api/session/start``
+    (which calls ``active.acquire(config, PTYTransport)``); the WS then
+    streams transport events out and pumps control frames in.
+
+    Inbound JSON control frames:
+    - ``{"type": "input", "data": "..."}``   → ``transport.send_input``
+    - ``{"type": "resize", "cols": N, "rows": N}`` → ``transport.resize``
+    - ``{"type": "stop"}``                    → ``transport.cancel``
+
+    Outbound framing:
+    - ``OutputBytes.data`` → **binary** frame (verbatim PTY bytes)
+    - ``Started``/``Stopped``/``TransportError``/``AgentMessage`` → text
+      JSON frames (``{"type": ..., ...}``)
+
+    Close codes:
+    - 1008 (Policy Violation) if Origin is disallowed, if no session is
+      active, or if ``?study_session_id`` does not match the active one.
+    - 1000 normal close on ``stop`` frame or transport-emitted ``Stopped``.
     """
-    await websocket.accept()
-    manager = websocket.app.state.agent_session_manager
-    session_id: str | None = None
-    event_task: asyncio.Task[None] | None = None
+    from studyloop.session import active as session_active
+    from studyloop.session.transport import (
+        AgentMessage,
+        OutputBytes,
+        Started,
+        Stopped,
+        TransportError,
+    )
 
-    async def forward_events(events) -> None:  # type: ignore[no-untyped-def]
-        async for event in events:
-            await websocket.send_json(event.as_dict())
+    # --- Pre-accept guards -----------------------------------------------
 
-    try:
-        while True:
-            message = await websocket.receive_json()
-            message_type = message.get("type")
-            if message_type == "start":
-                requested_transport = str(message.get("transport") or "pty")
-                if requested_transport == "acp":
-                    await websocket.send_json(
-                        {
-                            "type": "error",
-                            "data": {
-                                "message": (
-                                    "ACP transport is scaffolded but not enabled yet. "
-                                    "Use Browser terminal for Gemini or Kiro."
-                                )
-                            },
-                        }
-                    )
-                    continue
-                if session_id is not None:
-                    await websocket.send_json(
-                        {"type": "error", "data": {"message": "Session already started"}}
-                    )
-                    continue
-                session_id, events = await manager.start_session(
-                    topic=str(message.get("topic") or "Study Session"),
-                    energy=int(message.get("energy") or 5),
-                    agent=message.get("agent"),
-                    transport=requested_transport,
+    # Origin check (plan Blocker B1). Must happen before accept() — the
+    # handshake has not yet completed, so close-without-accept sends an
+    # HTTP 403 response per Starlette semantics, which the client sees as
+    # a failed upgrade. After accept() we can only send a WS close frame.
+    origin = websocket.headers.get("origin", "")
+    if not _origin_allowed(origin):
+        logger.warning("WS /session/ws rejected: disallowed origin=%r", origin)
+        await websocket.close(code=_WS_CLOSE_POLICY)
+        return
+
+    requested = websocket.query_params.get("study_session_id")
+    current = await session_active.current()
+    if current is None:
+        await websocket.close(code=_WS_CLOSE_POLICY)
+        return
+    if requested and requested != current.study_session_id:
+        logger.warning(
+            "WS /session/ws rejected: requested=%r active=%r",
+            requested,
+            current.study_session_id,
+        )
+        await websocket.close(code=_WS_CLOSE_POLICY)
+        return
+
+    await websocket.accept()
+    transport = current.transport
+
+    async def pty_to_ws() -> None:
+        """Pump transport events → WS frames until the session ends."""
+        async for event in transport.events():
+            if isinstance(event, OutputBytes):
+                await websocket.send_bytes(event.data)
+            elif isinstance(event, Started):
+                await websocket.send_json({"type": "started", "agent": event.agent})
+            elif isinstance(event, Stopped):
+                await websocket.send_json(
+                    {
+                        "type": "stopped",
+                        "returncode": event.returncode,
+                        "reason": event.reason,
+                    }
                 )
-                event_task = asyncio.create_task(forward_events(events))
-            elif message_type == "input":
-                if session_id is None:
-                    await websocket.send_json(
-                        {"type": "error", "data": {"message": "No active session"}}
-                    )
-                    continue
-                await manager.send(session_id, str(message.get("text") or ""))
-            elif message_type == "stop":
-                if session_id:
-                    await manager.stop(session_id)
-                    session_id = None
-                await websocket.close()
-                return
-            else:
+                return  # Stopped is terminal — stop pumping.
+            elif isinstance(event, TransportError):
+                await websocket.send_json({"type": "transport_error", "message": event.message})
+            elif isinstance(event, AgentMessage):
                 await websocket.send_json(
-                    {"type": "error", "data": {"message": "Unknown message type"}}
+                    {"type": "agent_message", "kind": event.kind, "payload": event.payload}
                 )
-    except WebSocketDisconnect:
+
+    async def ws_to_pty() -> None:
+        """Read WS control frames and forward to transport."""
+        while True:
+            frame = await websocket.receive_json()
+            ftype = frame.get("type")
+            if ftype == "input":
+                data = frame.get("data", "")
+                if isinstance(data, str):
+                    await transport.send_input(data.encode("utf-8"))
+            elif ftype == "resize":
+                try:
+                    cols = int(frame.get("cols", 80))
+                    rows = int(frame.get("rows", 24))
+                except (TypeError, ValueError):
+                    continue
+                await transport.resize(cols, rows)
+            elif ftype == "stop":
+                await transport.cancel()
+                return
+            # Silently drop unknown frame types — no error channel needed.
+
+    # --- Pump with TaskGroup (plan Blocker B5) ---------------------------
+    #
+    # TaskGroup raises ExceptionGroup on any child exception; the WS
+    # disconnect paths come up as ExceptionGroup[WebSocketDisconnect,
+    # ConnectionClosedOK, ...]. ``except*`` unpacks cleanly without
+    # reaching for .exceptions.
+
+    try:
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(pty_to_ws(), name="ws-pty-to-ws")
+            tg.create_task(ws_to_pty(), name="ws-ws-to-pty")
+    except* WebSocketDisconnect:
         pass
+    except* OSError as eg:
+        logger.error("PTY I/O error on /session/ws: %s", eg.exceptions)
+    except* Exception as eg:  # pragma: no cover — defensive
+        logger.exception("unexpected error on /session/ws: %s", eg.exceptions)
     finally:
-        if event_task:
-            event_task.cancel()
-        if session_id:
-            await manager.stop(session_id)
+        await session_active.release()
 
 
 @router.post("/session/end")
 
@@ -1,58 +1,26 @@
-"""Tests for web live-session WebSocket and picker endpoints."""
+"""Tests for web session-picker endpoints.
+
+Prior to §1.5 wiring this file also held WS route tests exercising the
+legacy ``{type: "start", ...}`` protocol and ``agent_session_manager``.
+Those were removed when the WS route migrated to the
+``active.acquire`` + ``PTYTransport`` flow; current WS coverage lives in
+``test_web_session_ws.py``.
+"""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
 from fastapi.testclient import TestClient
 
-from studyloop.session_runtime import SessionEvent
 from studyloop.web.app import create_app
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator
     from pathlib import Path
 
     import pytest
 
 
-class FakeManager:
-    """Minimal manager used by the WebSocket route tests."""
-
-    def __init__(self) -> None:
-        self.sent: list[str] = []
-        self.stopped: list[str] = []
-
-    async def start_session(
-        self,
-        *,
-        topic: str,
-        energy: int,
-        agent: str | None = None,
-        transport: str = "pty",
-    ) -> tuple[str, AsyncIterator[SessionEvent]]:
-        async def events() -> AsyncIterator[SessionEvent]:
-            yield SessionEvent(
-                "started",
-                "fake-session",
-                {
-                    "topic": topic,
-                    "energy": energy,
-                    "agent": agent or "shell",
-                    "transport": transport,
-                },
-            )
-            yield SessionEvent("output", "fake-session", {"text": "ready"})
-
-        return "fake-session", events()
-
-    async def send(self, session_id: str, text: str) -> None:
-        self.sent.append(f"{session_id}:{text}")
-
-    async def stop(self, session_id: str) -> None:
-        self.stopped.append(session_id)
-
-
 def test_session_options_returns_course_hierarchy(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -83,36 +51,3 @@ def __init__(self) -> None:
     assert "agents" in body
     assert all(agent["recommended_transport"] == "ttyd" for agent in body["agents"])
     assert all(agent["acp_ready"] is False for agent in body["agents"])
-
-
-def test_live_session_websocket_streams_events_and_accepts_input() -> None:
-    app = create_app()
-    manager = FakeManager()
-    app.state.agent_session_manager = manager
-    client = TestClient(app)
-
-    with client.websocket_connect("/api/session/ws") as websocket:
-        websocket.send_json({"type": "start", "topic": "Python", "energy": 6, "transport": "pty"})
-        assert websocket.receive_json()["type"] == "started"
-        assert websocket.receive_json()["type"] == "output"
-        websocket.send_json({"type": "input", "text": "what is a decorator?"})
-        websocket.send_json({"type": "stop"})
-
-    assert manager.sent == ["fake-session:what is a decorator?"]
-    assert manager.stopped == ["fake-session"]
-
-
-def test_live_session_websocket_rejects_acp_until_handshake_is_implemented() -> None:
-    app = create_app()
-    manager = FakeManager()
-    app.state.agent_session_manager = manager
-    client = TestClient(app)
-
-    with client.websocket_connect("/api/session/ws") as websocket:
-        websocket.send_json({"type": "start", "topic": "Python", "energy": 6, "transport": "acp"})
-        payload = websocket.receive_json()
-        websocket.close()
-
-    assert payload["type"] == "error"
-    assert "ACP transport is scaffolded" in payload["data"]["message"]
-    assert manager.sent == []