From 206e37e73eacc825c224ab522416ceb86a7c2b17 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 3 Mar 2026 18:28:39 -0800 Subject: [PATCH] =?UTF-8?q?Fix=20circuit=20breaker,=20rate=20limiter,=20an?= =?UTF-8?q?d=20Anthropic=20web=20search=20correctne=E2=80=A6=20(#15)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix circuit breaker, rate limiter, and Anthropic web search correctness (#15) Three tech debt items addressing correctness and security gaps: Circuit breaker HALF_OPEN single-request permit: - Rename should_allow_request property to acquire_request_permit() method to make the side-effecting, non-idempotent nature explicit - Add _half_open_permit flag: exactly one probe request in HALF_OPEN, subsequent callers blocked until probe completes - Explicitly reset permit on all state transitions (record_success, record_failure) for clean state machine invariants - Session uses BaseException catch to ensure record_failure always fires, preventing permanent circuit deadlock on probe crash Rate limiter X-Forwarded-For support: - Add resolve_client_ip() with rightmost-untrusted XFF parsing - Configurable trusted_proxies via --ratelimit-trusted-proxies CLI flag and [ratelimit] trusted_proxies config (comma-separated CIDRs) - IPv4-mapped IPv6 normalization (::ffff:x.x.x.x → IPv4) for dual-stack - Clientless requests (request.client is None) pass through instead of sharing a single "unknown" bucket - Log warning for invalid CIDR entries in trusted_proxies config - Show trusted proxies in startup log when enabled Anthropic web search multi-turn encrypted content: - Capture raw provider content blocks during streaming via _block_to_dict() using model_dump(exclude_none=True) to avoid Anthropic API rejection - Accumulate thinking_delta into raw_blocks (was silently empty on replay) - Store _provider_content on assistant messages, pass through verbatim in _convert_messages() so encrypted_content/encrypted_index survive turns - Persist to SQLite via new provider_data column (auto-migrated) - Add thinking/signature to _block_to_dict fallback attribute list 23 new tests (735 total), ruff + mypy clean. * Fix Copilot PR #15 review issues: provider data, circuit breaker, IP normalization - Persist assistant message when provider_data exists even if text content is empty — prevents losing Anthropic web search encrypted content needed for multi-turn replay (session.py) - Re-raise KeyboardInterrupt/SystemExit immediately after recording failure instead of attempting fallback models (session.py) - Consume HALF_OPEN permit for the transition caller — prevents two concurrent probe requests when only one should be allowed (healthcheck.py) - Normalize IPv4-mapped IPv6 addresses consistently in resolve_client_ip() — prevents duplicate rate-limit buckets for ::ffff:x.x.x.x vs x.x.x.x (ratelimit.py) --- docs/architecture.md | 9 +- tests/test_healthcheck.py | 55 +++++++- tests/test_providers.py | 170 +++++++++++++++++++++++++ tests/test_ratelimit.py | 109 ++++++++++++++++ turnstone/core/config.py | 1 + turnstone/core/healthcheck.py | 28 +++- turnstone/core/memory.py | 29 +++-- turnstone/core/providers/_anthropic.py | 77 ++++++++++- turnstone/core/providers/_protocol.py | 2 + turnstone/core/ratelimit.py | 69 ++++++++++ turnstone/core/session.py | 27 +++- turnstone/server.py | 22 +++- 12 files changed, 568 insertions(+), 30 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 75c495f9..326fc7bc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -819,7 +819,8 @@ HALF_OPEN ──(probe fails)──────────> OPEN - `record_success()` / `record_failure()` update `_consecutive_failures` and transition the `_state` (`CircuitState` enum: `CLOSED`, `OPEN`, `HALF_OPEN`). -- `should_allow_request()` returns `False` when the circuit is `OPEN`, causing +- `acquire_request_permit()` returns `False` when the circuit is `OPEN` or when + in `HALF_OPEN` and the single probe permit has already been consumed. Causes `ChatSession._create_stream_with_retry` to skip the backend and surface an error immediately. - The `/health` endpoint reads the monitor's state: `"status": "ok"` when the @@ -832,8 +833,12 @@ limits using a token-bucket algorithm. Each IP gets a `TokenBucket` with `requests_per_second` (refill rate) and `burst` (bucket capacity) from `[ratelimit]` config. -- Applied in `do_GET` / `do_POST` after authentication but before route dispatch. +- Applied via `RateLimitMiddleware` after authentication but before route dispatch. - `/health` and `/metrics` are exempt (monitoring must always be reachable). +- **X-Forwarded-For support**: when `trusted_proxies` is configured (comma-separated + CIDRs), the middleware parses the `X-Forwarded-For` header using the + rightmost-untrusted approach. IPv4-mapped IPv6 addresses are normalized. + The direct client IP must be in the trusted set before XFF is considered. - On limit exceeded: HTTP 429 with `Retry-After` header and JSON body `{"error": "Rate limit exceeded", "retry_after": N}`. - The `turnstone_ratelimit_rejected_total` counter is incremented on each diff --git a/tests/test_healthcheck.py b/tests/test_healthcheck.py index 842c8362..b20b05ca 100644 --- a/tests/test_healthcheck.py +++ b/tests/test_healthcheck.py @@ -91,7 +91,7 @@ class TestBackendHealthMonitor: mon = _make_monitor(mock_client, failure_threshold=1, cooldown=9999.0) mon.record_failure() assert mon.circuit_state == CircuitState.OPEN - assert mon.should_allow_request is False + assert mon.acquire_request_permit() is False @patch("turnstone.core.healthcheck.time") def test_half_open_after_cooldown( @@ -109,7 +109,7 @@ class TestBackendHealthMonitor: # Advance past cooldown mock_time.monotonic.return_value = t + 61.0 - assert mon.should_allow_request is True + assert mon.acquire_request_permit() is True assert mon.circuit_state == CircuitState.HALF_OPEN # type: ignore[comparison-overlap] def test_success_resets(self, mock_client: MagicMock, mock_metrics: MagicMock) -> None: @@ -126,18 +126,59 @@ class TestBackendHealthMonitor: def test_should_allow_when_closed(self, mock_client: MagicMock) -> None: mon = _make_monitor(mock_client) - assert mon.should_allow_request is True + assert mon.acquire_request_permit() is True - def test_should_allow_half_open(self, mock_client: MagicMock, mock_metrics: MagicMock) -> None: - """HALF_OPEN state allows requests (one probe attempt).""" + def test_half_open_allows_only_one_request( + self, mock_client: MagicMock, mock_metrics: MagicMock + ) -> None: + """HALF_OPEN permits exactly one probe; subsequent callers are blocked.""" mon = _make_monitor(mock_client, failure_threshold=1) mon.record_failure() assert mon.circuit_state == CircuitState.OPEN - # Force into HALF_OPEN + # Force into HALF_OPEN with permit with mon._lock: mon._state = CircuitState.HALF_OPEN - assert mon.should_allow_request is True + mon._half_open_permit = True + + # First caller gets through + assert mon.acquire_request_permit() is True + # Second caller is blocked + assert mon.acquire_request_permit() is False + # Third caller is also blocked + assert mon.acquire_request_permit() is False + + def test_half_open_success_reopens_to_all( + self, mock_client: MagicMock, mock_metrics: MagicMock + ) -> None: + """After probe succeeds in HALF_OPEN, circuit closes and all requests pass.""" + mon = _make_monitor(mock_client, failure_threshold=1) + mon.record_failure() + with mon._lock: + mon._state = CircuitState.HALF_OPEN + mon._half_open_permit = False # permit already consumed + + # Probe succeeds + mon.record_success() + assert mon.circuit_state == CircuitState.CLOSED # type: ignore[comparison-overlap] + # All callers pass now + assert mon.acquire_request_permit() is True + assert mon.acquire_request_permit() is True + + def test_half_open_failure_blocks_all( + self, mock_client: MagicMock, mock_metrics: MagicMock + ) -> None: + """After probe fails in HALF_OPEN, circuit reopens and all requests blocked.""" + mon = _make_monitor(mock_client, failure_threshold=1, cooldown=9999.0) + mon.record_failure() + with mon._lock: + mon._state = CircuitState.HALF_OPEN + mon._half_open_permit = False + + # Probe fails + mon.record_failure() + assert mon.circuit_state == CircuitState.OPEN + assert mon.acquire_request_permit() is False def test_half_open_failure_reopens( self, mock_client: MagicMock, mock_metrics: MagicMock diff --git a/tests/test_providers.py b/tests/test_providers.py index 7debe395..2a651efa 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -1726,3 +1726,173 @@ class TestTavilyFallback: result = provider._apply_web_search(kwargs, caps, tools) assert result is tools assert "web_search_options" not in kwargs + + +# =========================================================================== +# Anthropic provider_blocks / _provider_content round-trip tests +# =========================================================================== + + +class TestAnthropicProviderBlocks: + """Tests for multi-turn web search content preservation.""" + + def setup_method(self) -> None: + from turnstone.core.providers._anthropic import AnthropicProvider + + self.provider = AnthropicProvider() + + def test_convert_messages_uses_provider_content(self) -> None: + """Assistant message with _provider_content passes through verbatim.""" + provider_content = [ + {"type": "text", "text": "Here is what I found."}, + { + "type": "server_tool_use", + "id": "stu_123", + "name": "web_search", + "input": {"query": "turnstone bird"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "stu_123", + "content": [{"type": "web_search_result", "url": "https://example.com"}], + "encrypted_content": "abc123encrypted", + "encrypted_index": "idx456encrypted", + }, + ] + messages = [ + {"role": "user", "content": "Search for turnstone bird"}, + { + "role": "assistant", + "content": "Here is what I found.", + "_provider_content": provider_content, + }, + {"role": "user", "content": "Tell me more"}, + ] + _, converted = self.provider._convert_messages(messages) + # The assistant message should use provider_content verbatim + assistant_msg = converted[1] + assert assistant_msg["role"] == "assistant" + assert assistant_msg["content"] is provider_content + assert assistant_msg["content"][2]["encrypted_content"] == "abc123encrypted" + + def test_convert_messages_without_provider_content_unchanged(self) -> None: + """Assistant message without _provider_content uses normal reconstruction.""" + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + ] + _, converted = self.provider._convert_messages(messages) + assistant_msg = converted[1] + assert assistant_msg["role"] == "assistant" + assert assistant_msg["content"] == [{"type": "text", "text": "Hi there"}] + + def test_block_to_dict_with_model_dump(self) -> None: + """_block_to_dict uses model_dump(exclude_none=True) when available.""" + from turnstone.core.providers._anthropic import _block_to_dict + + class FakeBlock: + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + d = {"type": "text", "text": "hello", "extra": True, "nullable": None} + if kwargs.get("exclude_none"): + return {k: v for k, v in d.items() if v is not None} + return d + + result = _block_to_dict(FakeBlock()) + assert result == {"type": "text", "text": "hello", "extra": True} + assert "nullable" not in result + + def test_block_to_dict_fallback(self) -> None: + """_block_to_dict extracts known attributes as fallback.""" + from turnstone.core.providers._anthropic import _block_to_dict + + class FakeBlock: + type = "web_search_tool_result" + content = [{"type": "web_search_result"}] + encrypted_content = "enc123" + encrypted_index = "idx456" + + result = _block_to_dict(FakeBlock()) + assert result["type"] == "web_search_tool_result" + assert result["encrypted_content"] == "enc123" + assert result["encrypted_index"] == "idx456" + + def test_streaming_captures_provider_blocks(self) -> None: + """Streaming events produce provider_blocks on the final chunk.""" + from turnstone.core.providers._anthropic import AnthropicProvider + + provider = AnthropicProvider() + + # Build mock stream events + events = [] + + # Text block + text_block = MagicMock() + text_block.type = "text" + text_block.text = "" + text_block.model_dump.return_value = {"type": "text", "text": ""} + events.append(MagicMock(type="content_block_start", index=0, content_block=text_block)) + + events.append( + MagicMock( + type="content_block_delta", + index=0, + delta=MagicMock(type="text_delta", text="Hello"), + ) + ) + events.append(MagicMock(type="content_block_stop", index=0)) + + # Server tool use block + stu_block = MagicMock() + stu_block.type = "server_tool_use" + stu_block.name = "web_search" + stu_block.model_dump.return_value = { + "type": "server_tool_use", + "id": "stu_1", + "name": "web_search", + "input": {}, + } + events.append(MagicMock(type="content_block_start", index=1, content_block=stu_block)) + events.append( + MagicMock( + type="content_block_delta", + index=1, + delta=MagicMock(type="input_json_delta", partial_json='{"query":"test"}'), + ) + ) + events.append(MagicMock(type="content_block_stop", index=1)) + + # Web search tool result block + wsr_block = MagicMock() + wsr_block.type = "web_search_tool_result" + wsr_block.model_dump.return_value = { + "type": "web_search_tool_result", + "tool_use_id": "stu_1", + "content": [{"type": "web_search_result", "url": "https://example.com"}], + "encrypted_content": "enc_data", + "encrypted_index": "idx_data", + } + # Make content iterable for count + fake_result = MagicMock() + fake_result.type = "web_search_result" + wsr_block.content = [fake_result] + events.append(MagicMock(type="content_block_start", index=2, content_block=wsr_block)) + events.append(MagicMock(type="content_block_stop", index=2)) + + # Message delta with stop + msg_delta = MagicMock(type="message_delta") + msg_delta.delta = MagicMock(stop_reason="end_turn") + msg_delta.usage = MagicMock(input_tokens=100, output_tokens=50) + events.append(msg_delta) + + chunks = list(provider._iter_anthropic_stream(iter(events))) + + # Find the final chunk with provider_blocks + final_chunks = [c for c in chunks if c.provider_blocks] + assert len(final_chunks) == 1 + blocks = final_chunks[0].provider_blocks + assert len(blocks) == 3 + assert blocks[0]["type"] == "text" + assert blocks[1]["type"] == "server_tool_use" + assert blocks[1]["input"] == {"query": "test"} # parsed from accumulated JSON + assert blocks[2]["type"] == "web_search_tool_result" + assert blocks[2]["encrypted_content"] == "enc_data" diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py index 890c323d..a5849e4b 100644 --- a/tests/test_ratelimit.py +++ b/tests/test_ratelimit.py @@ -132,3 +132,112 @@ class TestRateLimiter: # 10.0.0.2 last_refill=1060, age=0 < 3600 => kept assert removed == 0 assert len(limiter._buckets) == 2 + + +# --------------------------------------------------------------------------- +# resolve_client_ip / parse_trusted_proxies +# --------------------------------------------------------------------------- + + +class TestResolveClientIp: + """X-Forwarded-For parsing with trusted proxy validation.""" + + def test_no_trusted_proxies_returns_direct(self): + from turnstone.core.ratelimit import resolve_client_ip + + result = resolve_client_ip("192.168.1.1", "10.0.0.1", frozenset()) + assert result == "192.168.1.1" + + def test_no_xff_returns_direct(self): + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("127.0.0.0/8") + result = resolve_client_ip("127.0.0.1", "", trusted) + assert result == "127.0.0.1" + + def test_trusted_proxy_extracts_xff(self): + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("127.0.0.1/32") + result = resolve_client_ip("127.0.0.1", "1.2.3.4", trusted) + assert result == "1.2.3.4" + + def test_untrusted_direct_ignores_xff(self): + """If the direct client is not a trusted proxy, XFF is ignored (anti-spoof).""" + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("10.0.0.0/8") + result = resolve_client_ip("203.0.113.5", "1.2.3.4", trusted) + assert result == "203.0.113.5" + + def test_chained_proxies(self): + """XFF: 'client, proxy1, proxy2' with proxy1+proxy2 trusted → returns client.""" + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("10.0.0.0/8") + result = resolve_client_ip("10.0.0.3", "1.2.3.4, 10.0.0.1, 10.0.0.2", trusted) + assert result == "1.2.3.4" + + def test_all_trusted_returns_direct(self): + """If all XFF entries are trusted proxies, fall back to direct IP.""" + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("10.0.0.0/8") + result = resolve_client_ip("10.0.0.3", "10.0.0.1, 10.0.0.2", trusted) + assert result == "10.0.0.3" + + def test_invalid_direct_ip_returns_direct(self): + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("10.0.0.0/8") + result = resolve_client_ip("not-an-ip", "1.2.3.4", trusted) + assert result == "not-an-ip" + + def test_invalid_xff_entry_skipped(self): + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("10.0.0.0/8") + result = resolve_client_ip("10.0.0.1", "garbage, 1.2.3.4", trusted) + assert result == "1.2.3.4" + + def test_ipv6_trusted_proxy(self): + from turnstone.core.ratelimit import parse_trusted_proxies, resolve_client_ip + + trusted = parse_trusted_proxies("::1/128") + result = resolve_client_ip("::1", "2001:db8::1", trusted) + assert result == "2001:db8::1" + + +class TestParseTrustedProxies: + def test_empty_string(self): + from turnstone.core.ratelimit import parse_trusted_proxies + + assert parse_trusted_proxies("") == frozenset() + + def test_single_cidr(self): + from turnstone.core.ratelimit import parse_trusted_proxies + + result = parse_trusted_proxies("10.0.0.0/8") + assert len(result) == 1 + + def test_multiple_cidrs(self): + from turnstone.core.ratelimit import parse_trusted_proxies + + result = parse_trusted_proxies("10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16") + assert len(result) == 3 + + def test_single_ip_becomes_host_network(self): + from turnstone.core.ratelimit import parse_trusted_proxies + + result = parse_trusted_proxies("127.0.0.1") + assert len(result) == 1 + + def test_invalid_entry_skipped(self): + from turnstone.core.ratelimit import parse_trusted_proxies + + result = parse_trusted_proxies("10.0.0.0/8, not-valid, 172.16.0.0/12") + assert len(result) == 2 + + def test_constructor_parses_trusted_proxies(self): + limiter = RateLimiter(enabled=True, rate=10.0, burst=5, trusted_proxies="10.0.0.0/8") + assert len(limiter.trusted_proxies) == 1 diff --git a/turnstone/core/config.py b/turnstone/core/config.py index 972bf596..c73dc760 100644 --- a/turnstone/core/config.py +++ b/turnstone/core/config.py @@ -107,6 +107,7 @@ _CONFIG_MAP: dict[str, dict[str, str]] = { "enabled": "ratelimit_enabled", "requests_per_second": "ratelimit_rps", "burst": "ratelimit_burst", + "trusted_proxies": "ratelimit_trusted_proxies", }, "health": { "backend_probe_interval": "health_probe_interval", diff --git a/turnstone/core/healthcheck.py b/turnstone/core/healthcheck.py index 059aa031..a7d77610 100644 --- a/turnstone/core/healthcheck.py +++ b/turnstone/core/healthcheck.py @@ -47,6 +47,8 @@ class BackendHealthMonitor: self._state = CircuitState.CLOSED self._consecutive_failures = 0 self._last_state_change = time.monotonic() + # Set True on OPEN→HALF_OPEN; consumed by first acquire_request_permit() call + self._half_open_permit = False self._stop_event = threading.Event() self._thread: threading.Thread | None = None @@ -73,8 +75,11 @@ class BackendHealthMonitor: with self._lock: self._consecutive_failures = 0 if self._state != CircuitState.CLOSED: + prev = self._state self._state = CircuitState.CLOSED + self._half_open_permit = False self._last_state_change = time.monotonic() + log.info("Circuit breaker CLOSED (was %s): backend recovered", prev.value) self._update_metrics() def record_failure(self) -> None: @@ -84,6 +89,7 @@ class BackendHealthMonitor: if self._state == CircuitState.HALF_OPEN: # Probe failed in HALF_OPEN — re-open immediately self._state = CircuitState.OPEN + self._half_open_permit = False self._last_state_change = time.monotonic() log.warning("Circuit breaker OPEN: probe failed in HALF_OPEN") self._update_metrics() @@ -113,20 +119,30 @@ class BackendHealthMonitor: with self._lock: return self._state - @property - def should_allow_request(self) -> bool: - """False only if circuit is OPEN (fast-fail). HALF_OPEN allows requests through.""" + def acquire_request_permit(self) -> bool: + """Consume one request permit if available. + + Returns True when the caller may proceed. In HALF_OPEN, only one probe + request is allowed — subsequent callers are blocked until the probe + completes (via ``record_success`` or ``record_failure``). + """ with self._lock: if self._state == CircuitState.OPEN: - # Check if cooldown has elapsed -> transition to HALF_OPEN if (time.monotonic() - self._last_state_change) >= self._cooldown: self._state = CircuitState.HALF_OPEN + self._half_open_permit = False # consumed by this caller self._last_state_change = time.monotonic() - log.info("Circuit breaker HALF_OPEN: cooldown elapsed") + log.info("Circuit breaker HALF_OPEN: cooldown elapsed, one probe permitted") self._update_metrics() + return True # this caller is the probe + return False + if self._state == CircuitState.HALF_OPEN: + # Only one probe request allowed; subsequent callers block + if self._half_open_permit: + self._half_open_permit = False return True return False - return True + return True # CLOSED # ------------------------------------------------------------------ # Background probe diff --git a/turnstone/core/memory.py b/turnstone/core/memory.py index 76019a00..a462274f 100644 --- a/turnstone/core/memory.py +++ b/turnstone/core/memory.py @@ -2,6 +2,8 @@ from __future__ import annotations +import contextlib +import json import os import sqlite3 from datetime import datetime, timedelta @@ -65,6 +67,12 @@ def open_db() -> sqlite3.Connection: except sqlite3.OperationalError: conn.execute("ALTER TABLE conversations ADD COLUMN tool_call_id TEXT") conn.commit() + # Migration: add provider_data column for raw provider content blocks + try: + conn.execute("SELECT provider_data FROM conversations LIMIT 0") + except sqlite3.OperationalError: + conn.execute("ALTER TABLE conversations ADD COLUMN provider_data TEXT") + conn.commit() # Sessions table — maps session_id to human-friendly alias/title conn.execute( "CREATE TABLE IF NOT EXISTS sessions " @@ -122,6 +130,7 @@ def save_message( tool_name: str | None = None, tool_args: str | None = None, tool_call_id: str | None = None, + provider_data: str | None = None, ) -> None: """Log a message to the conversations table.""" global _fts5_available @@ -130,9 +139,9 @@ def save_message( try: conn.execute( "INSERT INTO conversations (session_id, timestamp, role, content, " - "tool_name, tool_args, tool_call_id) " - "VALUES (?, datetime('now'), ?, ?, ?, ?, ?)", - (session_id, role, content, tool_name, tool_args, tool_call_id), + "tool_name, tool_args, tool_call_id, provider_data) " + "VALUES (?, datetime('now'), ?, ?, ?, ?, ?, ?)", + (session_id, role, content, tool_name, tool_args, tool_call_id, provider_data), ) if _fts5_available and content: try: @@ -468,7 +477,7 @@ def load_session_messages(session_id: str) -> list[dict[str, Any]]: conn = open_db() try: rows = conn.execute( - "SELECT role, content, tool_name, tool_args, tool_call_id " + "SELECT role, content, tool_name, tool_args, tool_call_id, provider_data " "FROM conversations WHERE session_id = ? ORDER BY id", (session_id,), ).fetchall() @@ -480,14 +489,18 @@ def load_session_messages(session_id: str) -> list[dict[str, Any]]: messages: list[dict[str, Any]] = [] i = 0 while i < len(rows): - role, content, tool_name, tool_args, tc_id = rows[i] + role, content, tool_name, tool_args, tc_id, provider_data = rows[i] if role == "user": messages.append({"role": "user", "content": content or ""}) i += 1 elif role == "assistant": - messages.append({"role": "assistant", "content": content}) + msg: dict[str, Any] = {"role": "assistant", "content": content} + if provider_data: + with contextlib.suppress(json.JSONDecodeError, TypeError): + msg["_provider_content"] = json.loads(provider_data) + messages.append(msg) i += 1 elif role == "tool_call": @@ -508,7 +521,7 @@ def load_session_messages(session_id: str) -> list[dict[str, Any]]: assistant_msg["tool_calls"] = [] while i < len(rows) and rows[i][0] == "tool_call": - _, _, tn, ta, stored_tc_id = rows[i] + _, _, tn, ta, stored_tc_id, _ = rows[i] call_id = stored_tc_id or f"call_{session_id}_{i}" assistant_msg["tool_calls"].append( { @@ -523,7 +536,7 @@ def load_session_messages(session_id: str) -> list[dict[str, Any]]: # Consume matching tool_result rows result_idx = 0 while i < len(rows) and rows[i][0] == "tool_result": - _, result_content, _, _, result_tc_id = rows[i] + _, result_content, _, _, result_tc_id, _ = rows[i] if result_tc_id: tc_id_to_use = result_tc_id elif result_idx < len(assistant_msg["tool_calls"]): diff --git a/turnstone/core/providers/_anthropic.py b/turnstone/core/providers/_anthropic.py index 270f2c01..4fa5b178 100644 --- a/turnstone/core/providers/_anthropic.py +++ b/turnstone/core/providers/_anthropic.py @@ -253,6 +253,14 @@ class AnthropicProvider: continue if role == "assistant": + # If raw provider content was preserved, pass it through verbatim + # so encrypted_content/encrypted_index from web search are retained + provider_content = msg.get("_provider_content") + if provider_content: + converted.append({"role": "assistant", "content": provider_content}) + i += 1 + continue + content_blocks: list[dict[str, Any]] = [] text = msg.get("content") if text: @@ -390,6 +398,8 @@ class AnthropicProvider: next_tool_index = 0 # Track server-side tool blocks (web search) — accumulate query input server_tool_blocks: dict[int, dict[str, str]] = {} + # Capture raw content blocks for multi-turn preservation + raw_blocks: dict[int, dict[str, Any]] = {} for event in stream: sc = StreamChunk() @@ -397,6 +407,7 @@ class AnthropicProvider: if event_type == "content_block_start": block = event.content_block + raw_blocks[event.index] = _block_to_dict(block) if block.type == "tool_use": idx = next_tool_index tool_block_to_index[event.index] = idx @@ -429,8 +440,18 @@ class AnthropicProvider: delta = event.delta if delta.type == "text_delta": sc.content_delta = delta.text + # Accumulate text into raw block for preservation + if event.index in raw_blocks: + raw_blocks[event.index]["text"] = ( + raw_blocks[event.index].get("text", "") + delta.text + ) elif delta.type == "thinking_delta": sc.reasoning_delta = delta.thinking + # Accumulate thinking text into raw block for round-trip + if event.index in raw_blocks: + raw_blocks[event.index]["thinking"] = ( + raw_blocks[event.index].get("thinking", "") + delta.thinking + ) elif delta.type == "input_json_delta": if event.index in server_tool_blocks: # Accumulate server tool input (search query) @@ -443,8 +464,21 @@ class AnthropicProvider: arguments_delta=delta.partial_json, ) ) + # Accumulate input JSON into raw block for tool_use/server_tool_use + if event.index in raw_blocks: + raw_blocks[event.index]["_input_json"] = ( + raw_blocks[event.index].get("_input_json", "") + delta.partial_json + ) elif event_type == "content_block_stop": + # Finalize accumulated input JSON into parsed input + if event.index in raw_blocks and "_input_json" in raw_blocks[event.index]: + rb = raw_blocks[event.index] + try: + rb["input"] = json.loads(rb.pop("_input_json")) + except (json.JSONDecodeError, TypeError): + rb.pop("_input_json", None) + # When a server tool block completes, emit search query info if event.index in server_tool_blocks: info = server_tool_blocks.pop(event.index) @@ -468,6 +502,9 @@ class AnthropicProvider: ) if hasattr(event.delta, "stop_reason") and event.delta.stop_reason: sc.finish_reason = _normalize_finish_reason(event.delta.stop_reason) + # Emit all raw content blocks for multi-turn preservation + if raw_blocks: + sc.provider_blocks = [raw_blocks[i] for i in sorted(raw_blocks)] elif event_type == "message_start": if hasattr(event.message, "usage") and event.message.usage: @@ -522,7 +559,9 @@ class AnthropicProvider: # which are handled server-side and don't require client execution. content_parts: list[str] = [] tool_calls: list[dict[str, Any]] = [] + provider_blocks: list[dict[str, Any]] = [] for block in response.content: + provider_blocks.append(_block_to_dict(block)) if block.type == "text": content_parts.append(block.text) elif block.type == "tool_use": @@ -536,7 +575,7 @@ class AnthropicProvider: }, } ) - # server_tool_use, web_search_tool_result — skip (server-handled) + # server_tool_use, web_search_tool_result — captured in provider_blocks finish_reason = _normalize_finish_reason(response.stop_reason or "end_turn") @@ -554,6 +593,7 @@ class AnthropicProvider: tool_calls=tool_calls if tool_calls else None, finish_reason=finish_reason, usage=usage, + provider_blocks=provider_blocks, ) # -- retryable errors ---------------------------------------------------- @@ -584,3 +624,38 @@ def _normalize_finish_reason(reason: str) -> str: # Server-side tool (web search) paused a long turn; treat as stop return "stop" return reason + + +def _block_to_dict(block: Any) -> dict[str, Any]: + """Convert an Anthropic SDK content block to a plain dict. + + Preserves all fields including ``encrypted_content`` and ``encrypted_index`` + on ``web_search_tool_result`` blocks, which must be passed back verbatim + on subsequent turns for citation resolution. + """ + if hasattr(block, "model_dump"): + return block.model_dump(exclude_none=True) # type: ignore[no-any-return] + # Fallback: extract known attributes + d: dict[str, Any] = {"type": getattr(block, "type", "")} + for attr in ( + "id", + "name", + "input", + "text", + "thinking", + "signature", + "content", + "encrypted_content", + "encrypted_index", + ): + val = getattr(block, attr, None) + if val is not None: + if hasattr(val, "model_dump"): + d[attr] = val.model_dump() + elif isinstance(val, list): + d[attr] = [ + item.model_dump() if hasattr(item, "model_dump") else item for item in val + ] + else: + d[attr] = val + return d diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index 425e5b86..55be53a0 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -44,6 +44,7 @@ class StreamChunk: finish_reason: str | None = None is_first: bool = False info_delta: str = "" + provider_blocks: list[dict[str, Any]] = field(default_factory=list) @dataclass @@ -54,6 +55,7 @@ class CompletionResult: tool_calls: list[dict[str, Any]] | None = None finish_reason: str = "stop" usage: UsageInfo | None = None + provider_blocks: list[dict[str, Any]] = field(default_factory=list) @dataclass(frozen=True) diff --git a/turnstone/core/ratelimit.py b/turnstone/core/ratelimit.py index 763af154..b4291187 100644 --- a/turnstone/core/ratelimit.py +++ b/turnstone/core/ratelimit.py @@ -6,9 +6,15 @@ Thread-safe. Zero external dependencies. from __future__ import annotations +import ipaddress +import logging import threading import time +log = logging.getLogger(__name__) + +_NetworkType = ipaddress.IPv4Network | ipaddress.IPv6Network + class TokenBucket: """Single token bucket for one client.""" @@ -38,6 +44,67 @@ class TokenBucket: return (1.0 - self.tokens) / self.rate +def parse_trusted_proxies(raw: str) -> frozenset[_NetworkType]: + """Parse a comma-separated list of IPs/CIDRs into a frozen set of networks.""" + if not raw or not raw.strip(): + return frozenset() + nets: list[_NetworkType] = [] + for entry in raw.split(","): + entry = entry.strip() + if not entry: + continue + try: + nets.append(ipaddress.ip_network(entry, strict=False)) + except ValueError: + log.warning("Ignoring invalid trusted_proxies entry: %r", entry) + return frozenset(nets) + + +def _normalize_ip( + addr: ipaddress.IPv4Address | ipaddress.IPv6Address, +) -> ipaddress.IPv4Address | ipaddress.IPv6Address: + """Collapse ``::ffff:x.x.x.x`` to its IPv4 form for dual-stack compatibility.""" + if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: + return addr.ipv4_mapped + return addr + + +def resolve_client_ip( + direct_ip: str, + forwarded_for: str, + trusted_proxies: frozenset[_NetworkType], +) -> str: + """Extract real client IP from X-Forwarded-For, only trusting known proxies. + + Uses the rightmost-untrusted approach: walks the XFF header right-to-left + and returns the first IP not in the trusted set. If the direct client IP + is not a trusted proxy, XFF is ignored entirely (prevents spoofing). + + IPv4-mapped IPv6 addresses (``::ffff:x.x.x.x``) are normalized to IPv4 + before checking against trusted proxies for dual-stack compatibility. + """ + if not trusted_proxies or not forwarded_for: + try: + return str(_normalize_ip(ipaddress.ip_address(direct_ip))) + except ValueError: + return direct_ip + try: + addr = _normalize_ip(ipaddress.ip_address(direct_ip)) + except ValueError: + return direct_ip + if not any(addr in net for net in trusted_proxies): + return str(addr) + parts = [p.strip() for p in forwarded_for.split(",") if p.strip()] + for ip_str in reversed(parts): + try: + ip = _normalize_ip(ipaddress.ip_address(ip_str)) + except ValueError: + continue + if not any(ip in net for net in trusted_proxies): + return str(ip) + return str(addr) + + class RateLimiter: """Per-IP rate limiter using token buckets.""" @@ -49,6 +116,7 @@ class RateLimiter: enabled: bool = False, rate: float = 10.0, burst: int = 20, + trusted_proxies: str = "", ) -> None: if enabled and rate <= 0: raise ValueError(f"rate must be > 0 when enabled, got {rate}") @@ -57,6 +125,7 @@ class RateLimiter: self.enabled = enabled self.rate = rate self.burst = burst + self.trusted_proxies: frozenset[_NetworkType] = parse_trusted_proxies(trusted_proxies) self._buckets: dict[str, TokenBucket] = {} self._lock = threading.Lock() diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 4ee479d3..5926b3d9 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -397,7 +397,7 @@ class ChatSession: before attempting a call — fast-fails when the backend is unreachable. """ # Circuit breaker check — fast-fail if backend is known to be down - if self._health_monitor and not self._health_monitor.should_allow_request: + if self._health_monitor and not self._health_monitor.acquire_request_permit(): raise ConnectionError("Backend unreachable (circuit breaker open)") try: @@ -405,9 +405,11 @@ class ChatSession: if self._health_monitor: self._health_monitor.record_success() return result - except Exception as primary_err: + except BaseException as primary_err: if self._health_monitor: self._health_monitor.record_failure() + if isinstance(primary_err, (KeyboardInterrupt, SystemExit)): + raise if not self._registry or not self._registry.fallback: raise # Try each fallback model. Fallbacks may use different backends; @@ -493,8 +495,15 @@ class ChatSession: # Log assistant message to conversation history content = assistant_msg.get("content", "") tc = assistant_msg.get("tool_calls") - if content: - save_message(self._session_id, "assistant", content) + provider_data = None + if assistant_msg.get("_provider_content"): + import json as _json + + provider_data = _json.dumps(assistant_msg["_provider_content"]) + if content or provider_data is not None: + save_message( + self._session_id, "assistant", content, provider_data=provider_data + ) if tc: for call in tc: fn = call.get("function", {}) @@ -620,6 +629,7 @@ class ChatSession: content_parts: list[str] = [] reasoning_parts: list[str] = [] tool_calls_acc: dict[int, dict[str, Any]] = {} + provider_blocks: list[dict[str, Any]] = [] first_token = True in_think = False # inside a ... block path1_reasoning = False # last reasoning came via reasoning_delta field @@ -775,6 +785,10 @@ class ChatSession: _stop_spinner_once() self.ui.on_info(f"{GRAY}{chunk.info_delta}{RESET}") + # Raw provider content blocks (for multi-turn preservation) + if chunk.provider_blocks: + provider_blocks = chunk.provider_blocks + # Flush any remaining buffered text if pending: _flush_text(pending, in_think) @@ -807,6 +821,11 @@ class ChatSession: if tool_calls_acc: msg["tool_calls"] = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] + # Store raw provider content blocks for multi-turn preservation + # (e.g. Anthropic web_search_tool_result with encrypted_content) + if provider_blocks: + msg["_provider_content"] = provider_blocks + return msg _print_lock = threading.Lock() diff --git a/turnstone/server.py b/turnstone/server.py index 8806d68e..a1824a96 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -37,6 +37,7 @@ from starlette.staticfiles import StaticFiles from turnstone import __version__ from turnstone.core.metrics import metrics as _metrics +from turnstone.core.ratelimit import resolve_client_ip from turnstone.core.session import ChatSession, SessionUI # noqa: F401 from turnstone.core.tools import TOOLS # noqa: F401 — available for introspection from turnstone.core.workstream import Workstream, WorkstreamManager @@ -367,7 +368,13 @@ class RateLimitMiddleware: if limiter is None: await self.app(scope, receive, send) return - client_ip = request.client.host if request.client else "unknown" + if not request.client: + # No peer address — cannot enforce per-IP limit; pass through + await self.app(scope, receive, send) + return + client_ip = request.client.host + xff = request.headers.get("X-Forwarded-For", "") + client_ip = resolve_client_ip(client_ip, xff, limiter.trusted_proxies) path = request.url.path allowed, retry_after = limiter.check(client_ip, path) if not allowed: @@ -1206,6 +1213,11 @@ def main() -> None: default=20, help="Rate limit: burst size (default: 20)", ) + parser.add_argument( + "--ratelimit-trusted-proxies", + default="", + help="Trusted proxy CIDRs for X-Forwarded-For parsing (comma-separated, e.g. '10.0.0.0/8,172.16.0.0/12')", + ) parser.add_argument( "--health-probe-interval", type=float, @@ -1303,6 +1315,7 @@ def main() -> None: enabled=args.ratelimit_enabled, rate=args.ratelimit_rps, burst=args.ratelimit_burst, + trusted_proxies=args.ratelimit_trusted_proxies, ) # Set up global event queue for state-change broadcasts @@ -1398,7 +1411,12 @@ def main() -> None: f"circuit breaker threshold={args.circuit_breaker_threshold}" ) if rate_limiter.enabled: - print(f"Rate limiter: {args.ratelimit_rps} req/s, burst={args.ratelimit_burst}") + proxy_info = ( + f", trusted proxies: {args.ratelimit_trusted_proxies}" + if args.ratelimit_trusted_proxies + else "" + ) + print(f"Rate limiter: {args.ratelimit_rps} req/s, burst={args.ratelimit_burst}{proxy_info}") print(f"Max workstreams: {args.max_workstreams}") print("Press Ctrl+C to stop.")