Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,6 @@

## Stage 4: Verification And Task Record
**Goal**: Prove the remediation and record it for repeatable audit follow-up.
**Success Criteria**: Focused pytest passes, `git diff --check` passes, Bandit runs on touched production files, and `TASK-12138` records final verification.
**Success Criteria**: Focused pytest passes, `git diff --check` passes, Bandit runs on touched production files, and `TASK-12146` records final verification.
**Tests**: `python -m pytest` focused files, `git diff --check`, and Bandit on touched production files.
**Status**: Complete
23 changes: 12 additions & 11 deletions tldw_Server_API/app/api/v1/endpoints/agent_client_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -1276,16 +1276,16 @@ async def _send_callback(message: dict[str, Any]) -> None:
# available; otherwise fall back to the existing path.
bus = get_session_event_bus(session_id)
if bus is not None and last_sequence > 0:
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}")
await reconnect_broadcaster.start(bus)
from tldw_Server_API.app.core.Agent_Client_Protocol.consumers.ws_broadcaster import (
start_reconnect_replay,
)

async def _ws_send_str(msg: str) -> None:
await stream.send_json(json.loads(msg))

await reconnect_broadcaster.add_connection(
reconnect_conn_id = f"ws-{session_id}-{id(stream)}"
reconnect_broadcaster = await start_reconnect_replay(
bus,
conn_id=reconnect_conn_id,
send_callback=_ws_send_str,
from_sequence=last_sequence,
Expand Down Expand Up @@ -1336,11 +1336,12 @@ async def _ws_send_str(msg: str) -> None:
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()
from tldw_Server_API.app.core.Agent_Client_Protocol.consumers.ws_broadcaster import (
stop_reconnect_replay,
)

# logs (never silently suppresses) cleanup failures
await stop_reconnect_replay(reconnect_broadcaster, reconnect_conn_id)
# Unregister WebSocket
if send_callback is not None:
try:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,15 @@ class WSBroadcaster(EventConsumer):
consumer_id: str = "ws_broadcaster"

def __init__(self, consumer_id: str | None = None) -> None:
"""Initialize the broadcaster.

Args:
consumer_id: Optional unique identifier used when subscribing to
the session event bus. Defaults to the class-level
``"ws_broadcaster"``. Pass a unique value (e.g. per
reconnect-replay connection) so concurrent broadcasters on the
same bus do not overwrite each other's subscriptions.
"""
self.consumer_id = consumer_id or self.consumer_id
self._connections: dict[str, _ConnectionInfo] = {}
self._bus: SessionEventBus | None = None
Expand Down Expand Up @@ -195,3 +204,52 @@ def _passes_filter(kind: AgentEventKind, verbosity: str) -> bool:
return kind in _SUMMARY_KINDS
# Unknown verbosity -- default to full
return True


async def start_reconnect_replay(
bus: SessionEventBus,
*,
conn_id: str,
send_callback: SendCallback,
from_sequence: int,
) -> WSBroadcaster:
"""Create and start a dedicated broadcaster that replays missed events.

Used on WebSocket reconnect: a per-connection broadcaster (unique
``consumer_id`` so concurrent reconnects don't clobber each other's bus
subscriptions) is subscribed to *bus* and replays buffered events from
*from_sequence* to *send_callback*. Callers must pair this with
:func:`stop_reconnect_replay` on disconnect.
"""
Comment thread
qodo-code-review[bot] marked this conversation as resolved.
broadcaster = WSBroadcaster(consumer_id=f"ws_broadcaster:{conn_id}")
await broadcaster.start(bus)
await broadcaster.add_connection(
conn_id=conn_id,
send_callback=send_callback,
from_sequence=from_sequence,
)
return broadcaster
Comment on lines +209 to +264

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If add_connection raises an exception after broadcaster.start(bus) has succeeded, the broadcaster will remain active and subscribed to the event bus, leading to a resource/memory leak. Since the caller (agent_client_protocol.py) only receives the broadcaster reference upon successful return, it won't be able to clean it up in its finally block. Wrapping add_connection in a try...except block to stop the broadcaster on failure ensures proper cleanup.

async def start_reconnect_replay(
    bus: SessionEventBus,
    *,
    conn_id: str,
    send_callback: SendCallback,
    from_sequence: int,
) -> WSBroadcaster:
    """Create and start a dedicated broadcaster that replays missed events.

    Used on WebSocket reconnect: a per-connection broadcaster (unique
    ``consumer_id`` so concurrent reconnects don't clobber each other's bus
    subscriptions) is subscribed to *bus* and replays buffered events from
    *from_sequence* to *send_callback*. Callers must pair this with
    :func:`stop_reconnect_replay` on disconnect.
    """
    broadcaster = WSBroadcaster(consumer_id=f"ws_broadcaster:{conn_id}")
    await broadcaster.start(bus)
    try:
        await broadcaster.add_connection(
            conn_id=conn_id,
            send_callback=send_callback,
            from_sequence=from_sequence,
        )
    except Exception:
        await broadcaster.stop()
        raise
    return broadcaster

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in dd8660b: add_connection is now wrapped so a mid-replay failure stops the broadcaster (best-effort, logged) before re-raising; regression test test_start_reconnect_replay_cleans_up_when_registration_fails asserts no leaked bus subscriber.

Comment thread
qodo-code-review[bot] marked this conversation as resolved.


async def stop_reconnect_replay(broadcaster: WSBroadcaster, conn_id: str | None) -> None:
"""Tear down a reconnect-replay broadcaster created by :func:`start_reconnect_replay`.

Cleanup failures are logged (never silently suppressed) but do not
propagate: this runs in disconnect paths where raising would mask the
original connection error.
"""
if conn_id is not None:
try:
broadcaster.remove_connection(conn_id)
except Exception as exc:
logger.warning(
"Failed to remove reconnect replay connection {}: {}", conn_id, exc
)
try:
await broadcaster.stop()
except Exception as exc:
logger.warning(
"Failed to stop reconnect replay broadcaster {}: {}",
broadcaster.consumer_id,
exc,
)
Comment thread
qodo-code-review[bot] marked this conversation as resolved.
Outdated
14 changes: 12 additions & 2 deletions tldw_Server_API/app/core/LLM_Calls/tokenizer_resolver.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
from typing import Any, Callable
from urllib.parse import parse_qsl, quote, quote_plus, urlparse

from loguru import logger

from tldw_Server_API.app.core.config import load_comprehensive_config
from tldw_Server_API.app.core.custom_openai_providers import (
custom_openai_config_option_names,
Expand Down Expand Up @@ -654,13 +656,17 @@ def _http_post(*, url: str, payload: dict[str, Any], headers: dict[str, str], ti
from tldw_Server_API.app.core.http_client import fetch
except Exception as exc:
raise TokenizerUnavailable("Provider tokenizer HTTP client unavailable") from exc
# allow_redirects=False: these POSTs carry provider credentials
# (Authorization / x-api-key); following a cross-origin redirect would
# resend the secrets to the redirect target. Provider tokenizer APIs
# never legitimately redirect.
return fetch(
method="POST",
url=url,
json=payload,
headers=headers,
timeout=timeout,
allow_redirects=True,
allow_redirects=False,
)


Expand Down Expand Up @@ -891,8 +897,12 @@ def _runtime_probe_exact_resolution(
)
finally:
if original_timeout is not None:
with suppress(Exception):
try:
setattr(encoding, "timeout_seconds", original_timeout)
except Exception as exc:
logger.warning(
f"Failed to restore tokenizer encoding timeout_seconds={original_timeout!r}: {exc}"
)

return resolution

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

from __future__ import annotations

import asyncio
import json
import tempfile
import time
Expand Down Expand Up @@ -204,7 +205,8 @@ async def run_arxiv_download_adapter(config: dict[str, Any], context: dict[str,
response.raise_for_status()
filename = pdf_url.split("/")[-1] or "paper.pdf"
output_path = str(art_dir / filename)
Path(output_path).write_bytes(response.content)
# blocking file write must not stall the event loop
await asyncio.to_thread(Path(output_path).write_bytes, response.content)

if callable(context.get("add_artifact")):
context["add_artifact"](
Expand Down
38 changes: 27 additions & 11 deletions tldw_Server_API/tests/Agent_Client_Protocol/test_acp_websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,8 +291,9 @@ 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):
async def test_acp_session_stream_reconnect_cleans_up_replay_broadcaster(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Reconnect replay broadcasters must not leave event-bus subscribers behind."""
import tldw_Server_API.app.api.v1.endpoints.agent_client_protocol as acp_endpoints

Expand All @@ -310,7 +311,7 @@ async def test_acp_session_stream_reconnect_cleans_up_replay_broadcaster(monkeyp
streams: list[Any] = []

class _DisconnectingStream:
def __init__(self, *args, **kwargs) -> None:
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.sent_payloads: list[dict[str, Any]] = []
self.stop_calls = 0
streams.append(self)
Expand All @@ -331,16 +332,23 @@ 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"):
async def _fake_authenticate_ws(
websocket: Any,
token: str | None = None,
api_key: str | None = None,
required_scope: str = "read",
) -> int:
return 1

async def _fake_get_runner_client():
async def _fake_get_runner_client() -> MockRunnerClient:
return stub_client

async def _fake_resolve_persona_id(client, *, session_id: str, user_id: int):
async def _fake_resolve_persona_id(client: Any, *, session_id: str, user_id: int) -> None:
return None

def _fake_acquire_quota(*, user_id: int, session_id: str, persona_id: str | None):
def _fake_acquire_quota(
*, user_id: int, session_id: str, persona_id: str | None
) -> tuple[str, None]:
return "quota-token", None

def _fake_release_quota(token: str) -> None:
Expand Down Expand Up @@ -981,7 +989,12 @@ async def test_determine_permission_tier_individual(self):

@pytest.mark.asyncio
async def test_determine_permission_tier_batch(self):
"""Test that write operations get batch tier."""
"""Batch tier is the fallback for tools that are neither destructive nor read-only.

Write/modify operations were promoted to the stricter "individual"
tier by the ACP runner security hardening (see permission_tiers.py);
this test previously pinned the pre-hardening behavior.
"""
from tldw_Server_API.app.core.Agent_Client_Protocol.runner_client import ACPRunnerClient

mock_config = MagicMock()
Expand All @@ -993,10 +1006,13 @@ async def test_determine_permission_tier_batch(self):

client = ACPRunnerClient(mock_config)

# Write operations should be batch
assert client._determine_permission_tier("fs.write") == "batch"
# Write/modify operations are individual tier post-hardening
assert client._determine_permission_tier("fs.write") == "individual"
assert client._determine_permission_tier("modify_file") == "individual"

# Tools matching neither the individual nor auto token sets fall back to batch
assert client._determine_permission_tier("git.commit") == "batch"
assert client._determine_permission_tier("modify_file") == "batch"
assert client._determine_permission_tier("fs.chmod") == "batch"

@pytest.mark.asyncio
async def test_permission_request_message_includes_runtime_policy_metadata(self):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
WSBroadcaster,
)

pytestmark = pytest.mark.unit


def _make_event(kind: AgentEventKind, session_id: str = "sess-1") -> AgentEvent:
return AgentEvent(session_id=session_id, kind=kind, payload={"data": kind.value})
Expand Down Expand Up @@ -56,8 +58,7 @@ async def fake_send(msg: str) -> None:
assert received[2]["kind"] == "completion"


@pytest.mark.asyncio
async def test_ws_broadcaster_allows_unique_consumer_ids():
async def test_ws_broadcaster_allows_unique_consumer_ids() -> None:
"""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")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,9 @@ def _fake_fetch(**kwargs):
"json": {"content": "hello"},
"headers": {"X-Test": "1"},
"timeout": 7.0,
"allow_redirects": True,
# redirects disabled: these POSTs carry provider credentials and
# must not follow cross-origin redirects (PR #2604 review)
"allow_redirects": False,
}
]

Expand Down
Loading