Fix ACP reconnect broadcaster cleanup#2619
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses audit finding MCP-002 by implementing proper cleanup for temporary reconnect replay broadcasters and connections in the Agent Client Protocol (ACP) stream. It also updates WSBroadcaster to support unique consumer IDs to prevent concurrent reconnects from overwriting each other's subscriptions. Feedback on these changes identifies a critical scoping bug where the lack of a nonlocal declaration in the nested callback prevents the cleanup logic from executing. Additionally, the new regression test needs to be updated to actually invoke the callback and trigger the broadcaster initialization, as it currently passes trivially without exercising the cleanup path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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}") |
There was a problem hiding this comment.
In Python, assigning to variables inside a nested function (like _send_callback) creates local variables within that function's scope unless they are explicitly declared as nonlocal. Because reconnect_broadcaster and reconnect_conn_id are assigned here without a nonlocal declaration, the outer variables in acp_session_stream remain None. As a result, the cleanup logic in the finally block is completely bypassed when a connection disconnects, leaving the broadcaster and subscriptions active.
Add a nonlocal declaration for both variables at the start of the block to ensure they bind to the outer scope.
| 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}") | |
| if bus is not None and last_sequence > 0: | |
| nonlocal reconnect_broadcaster, reconnect_conn_id | |
| from tldw_Server_API.app.core.Agent_Client_Protocol.consumers.ws_broadcaster import WSBroadcaster | |
| reconnect_conn_id = f"ws-{session_id}-{id(stream)}" | |
| reconnect_broadcaster = WSBroadcaster(consumer_id=f"ws_broadcaster:{reconnect_conn_id}") |
There was a problem hiding this comment.
Checked against the rebased branch. I did not apply this suggestion because the assignments to reconnect_broadcaster and reconnect_conn_id are in the acp_session_stream function scope, not inside _send_callback; adding nonlocal at that location would be invalid Python. The branch now verifies cleanup from the outer-scope variables with the strengthened reconnect replay regression in 7d82013.
| async def test_acp_session_stream_reconnect_cleans_up_replay_broadcaster(monkeypatch): | ||
| """Reconnect replay broadcasters must not leave event-bus subscribers behind.""" |
There was a problem hiding this comment.
The test currently passes trivially because _send_callback is never invoked during the test execution. Since the callback is never called, the temporary reconnect broadcaster is never initialized, meaning bus._subscribers is empty from the start and the cleanup logic is not actually exercised.
Consider publishing an event to the bus or directly invoking the registered callback during the test to ensure the broadcaster is created and active, thereby verifying that the cleanup logic successfully stops and removes it.
There was a problem hiding this comment.
Addressed in 7d82013. The reconnect cleanup test now preloads the session event bus with a buffered completion event, reconnects with last_sequence=1, asserts the stream receives the replayed payload, and then verifies the temporary broadcaster subscription and runner websocket callback are cleaned up.
Address MCP-002 by giving reconnect replay broadcasters unique consumer ids and stopping them during WebSocket cleanup. Add regressions for endpoint subscriber cleanup and broadcaster subscription isolation.
Update TASK-12142 with the draft PR reference so the audit remediation record is traceable.
rmusser01
force-pushed
the
codex/audit-acp-reconnect-cleanup-2026-07-04
branch
from
5cbd9d0 to
7d82013
Compare
PR Summary by QodoFix ACP reconnect replay broadcaster lifecycle cleanup
AI Description
Diagram
High-Level Assessment
Files changed (5)
|
Code Review by Qodo
Context used✅ Compliance rules (platform):
74 rules 1. WSBroadcaster.__init__ missing docstring
|
![qodo-code-review[bot]](https://avatars.githubusercontent.com/in/484649?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
1. wsbroadcaster.init missing docstring 📘 Rule violation ✧ Quality
The modified WSBroadcaster.__init__ method has no docstring, violating the requirement that all functions/methods include docstrings. Missing docstrings reduce maintainability and make the new consumer_id behavior harder to understand and audit.
Agent Prompt
## Issue description
`WSBroadcaster.__init__` was modified but still lacks a docstring. The compliance checklist requires docstrings for all functions/methods.
## Issue Context
A new optional `consumer_id` parameter was added and should be documented.
## Fix Focus Areas
- tldw_Server_API/app/core/Agent_Client_Protocol/consumers/ws_broadcaster.py[59-66]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| 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() |
There was a problem hiding this comment.
2. Cleanup suppresses endpoint exceptions 📘 Rule violation ☼ Reliability
The new cleanup logic uses contextlib.suppress(...) around broadcaster teardown, which can silently swallow exceptions without logging or justification. This can hide real cleanup failures and complicate operational debugging.
Agent Prompt
## Issue description
Cleanup logic in `acp_session_stream` suppresses exceptions via `contextlib.suppress(...)`, which violates the rule against silently swallowing exceptions.
## Issue Context
If teardown fails (e.g., unsubscribe/stop issues), the exception is currently dropped without any log line or explicit justification.
## Fix Focus Areas
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[1338-1343]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| 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, | ||
| ) |
There was a problem hiding this comment.
3. Reconnect workflow added in endpoint 📘 Rule violation ⌂ Architecture
Reconnect replay/broadcaster lifecycle logic is newly implemented directly inside the v1 endpoint module, rather than being encapsulated in a core feature module. This increases coupling and makes the endpoint responsible for business/workflow logic beyond request/response orchestration.
Agent Prompt
## Issue description
New reconnect replay workflow logic (broadcaster creation, unique consumer id selection, replay connection management) was added inside the endpoint module.
## Issue Context
Endpoints should primarily validate/route and delegate workflow/business logic to `/app/core/...` modules so the logic is reusable and testable without HTTP/WebSocket endpoint coupling.
## Fix Focus Areas
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[1277-1292]
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[1338-1343]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| @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.""" |
There was a problem hiding this comment.
4. Async test missing required marker 📘 Rule violation ▣ Testability
The newly added async tests use @pytest.mark.asyncio, but this marker is not in the policy’s approved set. This violates the rule that every test must have exactly one accepted marker.
Agent Prompt
## Issue description
New async tests are marked with `@pytest.mark.asyncio`, but the compliance rule requires exactly one approved marker (`unit`, `integration`, `external_api`, or `local_llm_service`) and `asyncio` is not allowed.
## Issue Context
The repository already configures `asyncio_mode = "auto"`, so async tests do not require the `asyncio` marker to run.
## Fix Focus Areas
- tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py[294-296]
- tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py[59-61]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| 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) |
There was a problem hiding this comment.
5. New test functions lack hints 📘 Rule violation ✧ Quality
The newly added async tests and related helpers introduce multiple function and method definitions without complete parameter and/or return type annotations (including the test function itself and nested helpers). This violates the requirement that all function parameters and return values be explicitly type hinted.
Agent Prompt
## Issue description
Newly added test code (including async tests and helper functions/methods) defines functions without complete type annotations for all parameters and explicit return types.
## Issue Context
PR Compliance ID 224215 requires explicit type hints for all function parameters and return values, and this applies to tests and any helper functions defined within them as well.
## Fix Focus Areas
- tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py[295-347]
- tldw_Server_API/tests/Agent_Client_Protocol/test_ws_broadcaster.py[60-64]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| 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() |
There was a problem hiding this comment.
6. Cleanup breaks on cancellation 🐞 Bug ☼ Reliability
In acp_session_stream(), cleanup awaits reconnect_broadcaster.stop() without shielding against task cancellation, and asyncio.CancelledError is not part of _ACP_ENDPOINT_NONCRITICAL_EXCEPTIONS. If the handler is cancelled mid-finalizer, the broadcaster may remain subscribed and the code after that await (unregister_websocket, stream.stop, quota release) may not run, undermining the reconnect cleanup goal.
Agent Prompt
### Issue description
`acp_session_stream()` now performs reconnect replay broadcaster cleanup in `finally`, but the awaited cleanup call (`await reconnect_broadcaster.stop()`) is not protected from cancellation. Since `_ACP_ENDPOINT_NONCRITICAL_EXCEPTIONS` does not include `asyncio.CancelledError`, a cancellation raised during cleanup can abort the finalizer early and skip:
- stopping/unsubscribing the replay broadcaster (subscriber leak)
- `unregister_websocket`
- `stream.stop()`
- quota release
This can reintroduce the lifecycle leak under cancellation scenarios (e.g., app shutdown / task cancellation).
### Issue Context
The SSE endpoint explicitly catches `asyncio.CancelledError` and still stops its consumer in `finally`, which is the more robust lifecycle pattern.
### Fix Focus Areas
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[1338-1352]
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[125-145]
- tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py[3363-3372]
### Suggested fix approach
1. Ensure finalizer cleanup cannot be interrupted mid-way by task cancellation:
- Wrap awaited cleanup calls in `asyncio.shield(...)`, and/or
- Catch `asyncio.CancelledError` around cleanup awaits, continue the remaining cleanup steps, then re-raise cancellation after critical cleanup completes.
2. Apply this to at least:
- `reconnect_broadcaster.stop()`
- `client.unregister_websocket(...)`
- `stream.stop()`
3. Keep cancellation propagation semantics correct (cleanup first, then re-raise) similar to how the SSE endpoint ensures `consumer.stop()` runs even when cancelled.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…s; sync docs - Add tests/Chat_Macros/property/test_macro_parser_properties.py: 10 hypothesis properties for parse_macro_args/_coerce_bool/ load_macro_definition, importorskip-guarded so they skip on dev today and activate when codex/chat-macros-v1 (PR #2618) merges. Validated against that branch: 10/10 pass. - ci.yml: add tests/RPG (11 files, silently unsharded on dev — shard guard was red repo-wide; 177/177 pass locally) and tests/Chat_Macros (covers the prep file now and the macro branch's tests on merge) to the gap-verified-7 shard. - Audit/plan: egress (#2604) and ACP (#2619) fixes merged to dev today — update live-unfixed wording; chat-macros re-added as prep-done pending PR #2618. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Follow-up PR #2634 addresses the remaining review comments from this PR: reconnect-replay lifecycle extracted to core (ws_broadcaster.start/stop_reconnect_replay) with logged cleanup instead of contextlib.suppress, WSBroadcaster.init docstring, marker/type-hint fixes on the new tests — plus a fix for the stale test_determine_permission_tier_batch that was red on dev. |
…s; sync docs - Add tests/Chat_Macros/property/test_macro_parser_properties.py: 10 hypothesis properties for parse_macro_args/_coerce_bool/ load_macro_definition, importorskip-guarded so they skip on dev today and activate when codex/chat-macros-v1 (PR #2618) merges. Validated against that branch: 10/10 pass. - ci.yml: add tests/RPG (11 files, silently unsharded on dev — shard guard was red repo-wide; 177/177 pass locally) and tests/Chat_Macros (covers the prep file now and the macro branch's tests on merge) to the gap-verified-7 shard. - Audit/plan: egress (#2604) and ACP (#2619) fixes merged to dev today — update live-unfixed wording; chat-macros re-added as prep-done pending PR #2618. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 participant
Addresses AUDIT-2026-06-27-MCP-002 / TASK-12142.
Change summary:
Latest-dev validation:
.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)..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.jsonwith 0 findings.git diff --check.TestACPRunnerClientPermissions::test_determine_permission_tier_batchexpectsfs.writeto bebatch, while current dev returnsindividual.Review notes:
nonlocalsuggestion was checked and not applied because the assignments are in theacp_session_streamfunction scope, not inside_send_callback; addingnonlocalthere would be invalid Python.Draft until the required human-authored Change summary is added.