Fix ACP reconnect broadcaster cleanup

backlog/tasks/task-12142 - Fix-audit-ACP-reconnect-broadcaster-cleanup.md

@@ -0,0 +1,88 @@
+---
+id: TASK-12142
+title: Fix audit ACP reconnect broadcaster cleanup
+status: Done
+assignee: []
+created_date: 2026-07-04 06:55
+updated_date: 2026-07-04 18:59
+labels:
+- audit
+- remediation
+- mcp
+- acp
+- websocket
+- reliability
+dependencies: []
+references:
+- AUDIT-2026-06-27-MCP-002
+- https://github.com/rmusser01/tldw_server/pull/2619
+documentation:
+- Docs/superpowers/reviews/2026-06-27-repo-audit/domains/mcp-sandbox-agent-protocol.md
+- Docs/superpowers/reviews/2026-06-27-repo-audit/specialists/reliability-lifecycle.md
+priority: low
+modified_files:
+- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py
+- tldw_Server_API/app/core/Agent_Client_Protocol/consumers/ws_broadcaster.py
+- tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py
+- tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py
+---
+
+## Description
+
+<!-- SECTION:DESCRIPTION:BEGIN -->
+Address audit finding MCP-002: ACP reconnect WebSocket replay creates a temporary WSBroadcaster but the endpoint finalizer does not remove the replay connection or stop the broadcaster, leaving event-bus subscribers/tasks behind after disconnect.
+<!-- SECTION:DESCRIPTION:END -->
+
+## Acceptance Criteria
+<!-- AC:BEGIN -->
+- [x] #1 Reconnect stream retains enough broadcaster lifecycle state for finalizer cleanup.
+- [x] #2 Endpoint finalizer removes the replay connection and stops the temporary broadcaster after disconnect/error.
+- [x] #3 Session event-bus subscriber count returns to baseline after last_sequence reconnect disconnect.
+- [x] #4 WSBroadcaster supports unique consumer ids so concurrent reconnect replay broadcasters do not overwrite each other.
+- [x] #5 Focused endpoint-level regression covers last_sequence reconnect cleanup.
+<!-- AC:END -->
+
+## Implementation Plan
+
+<!-- SECTION:PLAN:BEGIN -->
+1. Add focused regression coverage for endpoint reconnect-disconnect cleanup and broadcaster consumer-id isolation.
+2. Make the reconnect broadcaster lifecycle explicit in acp_session_stream, including connection removal and stop in the finalizer.
+3. Allow WSBroadcaster instances to use unique consumer ids while preserving the default consumer id for existing callers.
+4. Run focused ACP websocket/broadcaster tests, Bandit on touched production scope, and diff whitespace validation.
+<!-- SECTION:PLAN:END -->
+
+## Implementation Notes
+
+<!-- SECTION:NOTES:BEGIN -->
+Validation notes:
+- Original PR branch was stale after latest dev advanced; fetched origin/dev and rebased cleanly onto fd5c152b065c408e4e8ee5f08da41589f21cb7f5. Post-rebase merge-base matched origin/dev before validation.
+- Red check failed as expected before the original implementation: reconnect endpoint left `ws_broadcaster` in the session event bus and WSBroadcaster did not accept a custom `consumer_id`.
+- Reviewed Gemini feedback on PR #2619. The `nonlocal` suggestion is technically invalid for the current code because `reconnect_broadcaster` and `reconnect_conn_id` are assigned in the `acp_session_stream` function scope, not inside `_send_callback`. No endpoint code change was made for that suggestion.
+- Strengthened the endpoint reconnect regression so it preloads the session event bus with a buffered completion event and asserts the reconnect stream receives the replayed payload before verifying cleanup. This prevents the test from passing trivially when the reconnect broadcaster path is not exercised.
+- Targeted regression command passed after rebase and review follow-up: `.venv/bin/python -m pytest tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py::test_acp_session_stream_reconnect_cleans_up_replay_broadcaster tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py::test_ws_broadcaster_allows_unique_consumer_ids -q` (2 passed).
+- Bandit passed with 0 findings: `.venv/bin/python -m bandit -r tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py tldw_Server_API/app/core/Agent_Client_Protocol/consumers/ws_broadcaster.py -f json -o /tmp/bandit_acp_reconnect_cleanup_latest_dev.json`.
+- `git diff --check` passed.
+- Broader file command on latest dev still has one unrelated baseline failure: `test_acp_websocket.py::TestACPRunnerClientPermissions::test_determine_permission_tier_batch` expects `fs.write` to be `batch`, but current dev returns `individual`; the focused ACP reconnect and broadcaster regressions pass.
+PR opened: https://github.com/rmusser01/tldw_server/pull/2619 (draft against dev). Draft status is intentional until the human-authored Change summary required by project policy is added.
+<!-- SECTION:NOTES:END -->
+
+## Final Summary
+
+<!-- SECTION:FINAL_SUMMARY:BEGIN -->
+Implemented MCP-002 remediation for ACP reconnect replay cleanup and refreshed it against latest dev. `acp_session_stream` retains the temporary replay broadcaster and connection id, uses a per-connection WSBroadcaster consumer id, removes the replay connection, and stops the broadcaster in the endpoint finalizer. `WSBroadcaster` accepts an optional consumer id while preserving the existing default. The endpoint regression now proves the reconnect broadcaster path by replaying a buffered completion event and asserting the stream receives it before disconnect cleanup.
+
+Verification on latest dev fd5c152b065c408e4e8ee5f08da41589f21cb7f5:
+- Merge-base matched origin/dev after rebase.
+- Passed: `.venv/bin/python -m pytest tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py::test_acp_session_stream_reconnect_cleans_up_replay_broadcaster tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py::test_ws_broadcaster_allows_unique_consumer_ids -q` (2 passed).
+- Passed: `.venv/bin/python -m bandit -r tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py tldw_Server_API/app/core/Agent_Client_Protocol/consumers/ws_broadcaster.py -f json -o /tmp/bandit_acp_reconnect_cleanup_latest_dev.json` with 0 findings.
+- Passed: `git diff --check`.
+- Broader run note: `test_acp_websocket.py test_ws_broadcaster.py -q` has one unrelated current-dev failure in `TestACPRunnerClientPermissions::test_determine_permission_tier_batch`.
+<!-- SECTION:FINAL_SUMMARY:END -->
+
+## Definition of Done
+<!-- DOD:BEGIN -->
+- [x] #1 Focused ACP websocket and broadcaster tests pass.
+- [x] #2 Bandit over touched production files reports no new issues.
+- [x] #3 git diff --check passes.
+- [x] #4 Backlog task contains verification evidence and final summary.
+<!-- DOD:END -->

tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py

@@ -1255,6 +1255,8 @@ async def acp_session_stream(
         labels={"component": "acp", "endpoint": "acp_session_stream"},
     )
     send_callback: Any | None = None
+    reconnect_broadcaster: Any | None = None
+    reconnect_conn_id: str | None = None
 
     try:
         await stream.start()
@@ -1276,14 +1278,15 @@ async def _send_callback(message: dict[str, Any]) -> None:
         if bus is not None and last_sequence > 0:
             from tldw_Server_API.app.core.Agent_Client_Protocol.consumers.ws_broadcaster import WSBroadcaster
 
-            broadcaster = WSBroadcaster()
-            await broadcaster.start(bus)
+            reconnect_conn_id = f"ws-{session_id}-{id(stream)}"
+            reconnect_broadcaster = WSBroadcaster(consumer_id=f"ws_broadcaster:{reconnect_conn_id}")
+            await reconnect_broadcaster.start(bus)
 
             async def _ws_send_str(msg: str) -> None:
                 await stream.send_json(json.loads(msg))
 
-            await broadcaster.add_connection(
-                conn_id=f"ws-{session_id}-{id(stream)}",
+            await reconnect_broadcaster.add_connection(
+                conn_id=reconnect_conn_id,
                 send_callback=_ws_send_str,
                 from_sequence=last_sequence,
             )
@@ -1332,6 +1335,12 @@ async def _ws_send_str(msg: str) -> None:
     except _ACP_ENDPOINT_NONCRITICAL_EXCEPTIONS:
         logger.exception("WebSocket error for ACP session {}", session_id)
     finally:
+        if reconnect_broadcaster is not None:
+            if reconnect_conn_id is not None:
+                with contextlib.suppress(_ACP_ENDPOINT_NONCRITICAL_EXCEPTIONS):
+                    reconnect_broadcaster.remove_connection(reconnect_conn_id)
+            with contextlib.suppress(_ACP_ENDPOINT_NONCRITICAL_EXCEPTIONS):
+                await reconnect_broadcaster.stop()
         # Unregister WebSocket
         if send_callback is not None:
             try:

tldw_Server_API/app/core/Agent_Client_Protocol/consumers/ws_broadcaster.py

@@ -56,7 +56,8 @@ class WSBroadcaster(EventConsumer):
 
     consumer_id: str = "ws_broadcaster"
 
-    def __init__(self) -> None:
+    def __init__(self, consumer_id: str | None = None) -> None:
+        self.consumer_id = consumer_id or self.consumer_id
         self._connections: dict[str, _ConnectionInfo] = {}
         self._bus: SessionEventBus | None = None
         self._queue: asyncio.Queue[AgentEvent] | None = None

tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py

@@ -25,6 +25,7 @@
     SessionWebSocketRegistry,
 )
 from tldw_Server_API.app.core.Agent_Client_Protocol.event_bus import SessionEventBus
+from tldw_Server_API.app.core.Agent_Client_Protocol.events import AgentEvent, AgentEventKind
 from tldw_Server_API.app.core.Agent_Client_Protocol.stdio_client import ACPMessage
 from tldw_Server_API.app.services.acp_runtime_policy_service import (
     ACPRuntimePolicySnapshot,
@@ -290,6 +291,91 @@ async def _fake_get_runner_client():
         pytest.fail(f"Expected exactly one stream.stop() call, got {lifecycle['stop_calls']}")
 
 
+@pytest.mark.asyncio
+async def test_acp_session_stream_reconnect_cleans_up_replay_broadcaster(monkeypatch):
+    """Reconnect replay broadcasters must not leave event-bus subscribers behind."""
+    import tldw_Server_API.app.api.v1.endpoints.agent_client_protocol as acp_endpoints
+
+    session_id = "session-reconnect-cleanup"
+    bus = SessionEventBus(session_id=session_id)
+    await bus.publish(
+        AgentEvent(
+            session_id=session_id,
+            kind=AgentEventKind.COMPLETION,
+            payload={"detail": "replay-me"},
+        )
+    )
+    stub_client = MockRunnerClient()
+    release_calls: list[str] = []
+    streams: list[Any] = []
+
+    class _DisconnectingStream:
+        def __init__(self, *args, **kwargs) -> None:
+            self.sent_payloads: list[dict[str, Any]] = []
+            self.stop_calls = 0
+            streams.append(self)
+
+        async def start(self) -> None:
+            return None
+
+        async def stop(self) -> None:
+            self.stop_calls += 1
+
+        async def send_json(self, payload: dict[str, Any]) -> None:
+            self.sent_payloads.append(payload)
+
+        async def receive_json(self) -> dict[str, Any]:
+            raise WebSocketDisconnect(code=1000)
+
+    class _FakeWebSocket:
+        async def close(self, code: int = 1000) -> None:
+            return None
+
+    async def _fake_authenticate_ws(websocket, token=None, api_key=None, required_scope="read"):
+        return 1
+
+    async def _fake_get_runner_client():
+        return stub_client
+
+    async def _fake_resolve_persona_id(client, *, session_id: str, user_id: int):
+        return None
+
+    def _fake_acquire_quota(*, user_id: int, session_id: str, persona_id: str | None):
+        return "quota-token", None
+
+    def _fake_release_quota(token: str) -> None:
+        release_calls.append(token)
+
+    monkeypatch.setattr(acp_endpoints, "_authenticate_ws", _fake_authenticate_ws)
+    monkeypatch.setattr(acp_endpoints, "get_runner_client", _fake_get_runner_client)
+    monkeypatch.setattr(acp_endpoints, "get_session_event_bus", lambda _: bus)
+    monkeypatch.setattr(acp_endpoints, "_resolve_acp_session_persona_id", _fake_resolve_persona_id)
+    monkeypatch.setattr(acp_endpoints, "_acp_ws_try_acquire_quota", _fake_acquire_quota)
+    monkeypatch.setattr(acp_endpoints, "_acp_ws_release_quota", _fake_release_quota)
+    monkeypatch.setattr(acp_endpoints, "WebSocketStream", _DisconnectingStream)
+
+    await acp_endpoints.acp_session_stream(
+        _FakeWebSocket(),
+        session_id,
+        token="valid-token",
+        last_sequence=1,
+    )
+
+    assert bus._subscribers == {}
+    assert not stub_client.has_websocket_connections(session_id)
+    assert release_calls == ["quota-token"]
+    assert streams and streams[0].stop_calls == 1
+    replayed_payloads = [
+        payload
+        for payload in streams[0].sent_payloads
+        if payload.get("kind") == AgentEventKind.COMPLETION.value
+    ]
+    assert len(replayed_payloads) == 1
+    assert replayed_payloads[0]["session_id"] == session_id
+    assert replayed_payloads[0]["sequence"] == 1
+    assert replayed_payloads[0]["payload"] == {"detail": "replay-me"}
+
+
 class TestACPWebSocketConnection:
     """Tests for WebSocket connection handling."""
 

tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py

@@ -56,6 +56,27 @@ async def fake_send(msg: str) -> None:
     assert received[2]["kind"] == "completion"
 
 
+@pytest.mark.asyncio
+async def test_ws_broadcaster_allows_unique_consumer_ids():
+    """Distinct broadcasters should not overwrite each other's bus subscription."""
+    bus = SessionEventBus(session_id="sess-unique-consumers")
+    first = WSBroadcaster(consumer_id="ws_broadcaster:conn-1")
+    second = WSBroadcaster(consumer_id="ws_broadcaster:conn-2")
+
+    try:
+        await first.start(bus)
+        await second.start(bus)
+
+        assert set(bus._subscribers) == {
+            "ws_broadcaster:conn-1",
+            "ws_broadcaster:conn-2",
+        }
+    finally:
+        await first.stop()
+        await second.stop()
+    assert bus._subscribers == {}
+
+
 @pytest.mark.asyncio
 async def test_ws_broadcaster_summary_filters_thinking():
     """Summary verbosity should drop thinking/tool_call/etc., keep completion."""