From e4265ce38a7910b4819942657a64d7d337d0c7b4 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 20 Jul 2026 20:35:18 +0100 Subject: [PATCH] fix: mini issues --- application/agents/tools/read_webpage.py | 11 ++++ application/llm/openai.py | 7 +++ application/security/safe_url.py | 70 +++++++++++++----------- tests/agents/test_read_webpage_tool.py | 27 ++++++++- tests/llm/test_openai_responses.py | 34 ++++++++++++ 5 files changed, 115 insertions(+), 34 deletions(-) diff --git a/application/agents/tools/read_webpage.py b/application/agents/tools/read_webpage.py index c8b25227..4c0203df 100644 --- a/application/agents/tools/read_webpage.py +++ b/application/agents/tools/read_webpage.py @@ -138,6 +138,17 @@ class ReadWebpageTool(Tool): headers={'User-Agent': 'DocsGPT-Agent/1.0'}, timeout=10, ) + # Redirects are not followed (each hop would need its own SSRF + # validation); without this the redirect body comes back as + # near-empty markdown with no hint of where the page went. + if 300 <= response.status_code < 400: + location = response.headers.get("Location", "") + if location: + return ( + f"Error: URL redirects to '{location}'. " + "Fetch that URL directly instead." + ) + return "Error: URL responded with a redirect but no target location." response.raise_for_status() content_type = response.headers.get("Content-Type", "") diff --git a/application/llm/openai.py b/application/llm/openai.py index 48b90726..e078a839 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -205,6 +205,10 @@ class OpenAILLM(BaseLLM): "response_id": ( self._last_response_id if settings.OPENAI_RESPONSES_STORE else None ), + # Without these the coverage guard in ``_build_responses_input`` + # is inert after a cross-process resume, and the trim's carry + # loop re-sends every earlier round's outputs unfiltered. + "call_ids": sorted(self._last_response_call_ids or ()), "reasoning_items": self._last_reasoning_items, "reasoning_for_calls": self._reasoning_for_calls, } @@ -216,6 +220,9 @@ class OpenAILLM(BaseLLM): if state.get("chain_key") != self.responses_chain_key(): return False self._imported_response_id = state.get("response_id") + # Tolerant of rows persisted before ``call_ids`` existed: an empty + # set keeps the guard disabled rather than rejecting the state. + self._last_response_call_ids = set(state.get("call_ids") or ()) self._last_reasoning_items = list(state.get("reasoning_items") or []) self._reasoning_for_calls = dict(state.get("reasoning_for_calls") or {}) return True diff --git a/application/security/safe_url.py b/application/security/safe_url.py index 974f928f..73af4c97 100644 --- a/application/security/safe_url.py +++ b/application/security/safe_url.py @@ -305,6 +305,42 @@ def _url_userinfo_prefix(netloc: str) -> str: return f"{netloc.rsplit('@', 1)[0]}@" +def _pinned_session( + url: str, headers: dict[str, str] | None +) -> tuple[requests.Session, str, dict[str, str]]: + """Run the SSRF guard once and build a session pinned to the chosen IP. + + Single home for the pinning preamble (validate/pick IP, netloc rewrite, + ``Host`` header, HTTPS adapter mount) shared by :func:`pinned_request` + and :func:`pinned_fetch_bytes`, so the guard cannot drift between them. + + Returns: + Tuple of the session (caller must close it), the IP-literal URL, + and the request headers carrying the original ``Host``. + + Raises: + UnsafeUserUrlError: If the URL fails the SSRF guard. + """ + + host, ip, parts = _validate_and_pick_ip(url) + + netloc = f"{_url_userinfo_prefix(parts.netloc)}{_ip_to_url_host(ip)}" + if parts.port is not None: + netloc = f"{netloc}:{parts.port}" + pinned_url = urlunsplit( + (parts.scheme, netloc, parts.path, parts.query, parts.fragment) + ) + + request_headers = dict(headers or {}) + host_header = host if parts.port is None else f"{host}:{parts.port}" + request_headers["Host"] = host_header + + session = requests.Session() + if parts.scheme == "https": + session.mount("https://", _PinnedHostAdapter(host)) + return session, pinned_url, request_headers + + def pinned_request( method: str, url: str, @@ -324,22 +360,7 @@ def pinned_request( requests.RequestException: For network-level failures. """ - host, ip, parts = _validate_and_pick_ip(url) - - netloc = f"{_url_userinfo_prefix(parts.netloc)}{_ip_to_url_host(ip)}" - if parts.port is not None: - netloc = f"{netloc}:{parts.port}" - pinned_url = urlunsplit( - (parts.scheme, netloc, parts.path, parts.query, parts.fragment) - ) - - request_headers = dict(headers or {}) - host_header = host if parts.port is None else f"{host}:{parts.port}" - request_headers["Host"] = host_header - - session = requests.Session() - if parts.scheme == "https": - session.mount("https://", _PinnedHostAdapter(host)) + session, pinned_url, request_headers = _pinned_session(url, headers) try: return session.request( method=method.upper(), @@ -379,22 +400,7 @@ def pinned_fetch_bytes( requests.RequestException: For network-level failures. """ - host, ip, parts = _validate_and_pick_ip(url) - - netloc = f"{_url_userinfo_prefix(parts.netloc)}{_ip_to_url_host(ip)}" - if parts.port is not None: - netloc = f"{netloc}:{parts.port}" - pinned_url = urlunsplit( - (parts.scheme, netloc, parts.path, parts.query, parts.fragment) - ) - - request_headers = dict(headers or {}) - host_header = host if parts.port is None else f"{host}:{parts.port}" - request_headers["Host"] = host_header - - session = requests.Session() - if parts.scheme == "https": - session.mount("https://", _PinnedHostAdapter(host)) + session, pinned_url, request_headers = _pinned_session(url, headers) try: response = session.request( method="GET", diff --git a/tests/agents/test_read_webpage_tool.py b/tests/agents/test_read_webpage_tool.py index f0560f00..ed98673e 100644 --- a/tests/agents/test_read_webpage_tool.py +++ b/tests/agents/test_read_webpage_tool.py @@ -19,8 +19,8 @@ class _FakeResponse: raise requests.exceptions.HTTPError(str(self.status_code)) -def _fetch_result(body: bytes, content_type=None, status_code: int = 200): - headers = {} +def _fetch_result(body: bytes, content_type=None, status_code: int = 200, headers=None): + headers = dict(headers or {}) if content_type is not None: headers["Content-Type"] = content_type return body, _FakeResponse(status_code=status_code, headers=headers) @@ -236,6 +236,29 @@ class TestReadWebpageContentGuards: assert "widget" in result + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_redirect_reported_with_target(self, mock_fetch, tool): + # allow_redirects=False (SSRF) means a 3xx would otherwise return + # the redirect body as near-empty markdown with no hint. + mock_fetch.return_value = _fetch_result( + b"", status_code=301, + headers={"Location": "https://example.com/moved-here"}, + ) + + result = tool.execute_action("read_webpage", url="https://example.com/old") + + assert result.startswith("Error") + assert "https://example.com/moved-here" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_redirect_without_location_reported(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result(b"", status_code=302) + + result = tool.execute_action("read_webpage", url="https://example.com/old") + + assert result.startswith("Error") + assert "redirect" in result.lower() + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") def test_rss_feed_allowed(self, mock_fetch, tool): mock_fetch.return_value = _fetch_result( diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 7373a4b7..1d2bb5d2 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -286,6 +286,40 @@ def test_import_rejects_state_for_different_target(monkeypatch): assert second.import_responses_state(state) is False +@pytest.mark.unit +def test_export_import_round_trips_call_ids_so_filter_survives_resume(monkeypatch): + """Cross-process resume: without the call-id set the coverage guard is + inert, and the trim's carry loop re-sends every earlier round's outputs + from a grouped multi-round history.""" + first = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + first._last_response_id = "resp_prev" + first._last_response_call_ids = {"r2c0", "r2c1"} + + second = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + assert second.import_responses_state(first.export_responses_state()) is True + assert second._last_response_call_ids == {"r2c0", "r2c1"} + + items, prev = second._build_responses_input( + _multi_round_messages(3), "resp_prev" + ) + assert prev == "resp_prev" + outputs = [ + i["call_id"] for i in items if i.get("type") == "function_call_output" + ] + assert sorted(outputs) == ["r2c0", "r2c1"] + + +@pytest.mark.unit +def test_import_tolerates_state_without_call_ids(monkeypatch): + """Rows persisted before the call-id key existed must still import.""" + first = _make_llm(monkeypatch, _responses_caps()) + state = first.export_responses_state() + state.pop("call_ids", None) + second = _make_llm(monkeypatch, _responses_caps()) + assert second.import_responses_state(state) is True + assert second._last_response_call_ids == set() + + @pytest.mark.unit def test_start_responses_turn_drops_prior_turn_call_map(monkeypatch): llm = _make_llm(monkeypatch, _responses_caps())