From 0a9f18c58f4f8026db1083bf7c4b8945b52d09c8 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 20 Jul 2026 12:18:36 +0100 Subject: [PATCH 1/4] fix: add benchmark to ignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 948884e6..b6a4a4e1 100644 --- a/.gitignore +++ b/.gitignore @@ -203,3 +203,4 @@ tests/e2e/playwright-report/ tests/e2e/test-results/ tests/e2e/.e2e-last-run.json application/core/models/foundry_test_internal.yaml +bench/ From 3cbd284948682115a9c991c73fbe34d5c79f3ad7 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 20 Jul 2026 18:29:24 +0100 Subject: [PATCH 2/4] fix: multitool call openai issue --- application/llm/handlers/base.py | 91 +++---- application/llm/openai.py | 133 +++++++++- tests/llm/handlers/test_llm_handlers.py | 189 +++++++++++++- tests/llm/test_openai_responses.py | 315 ++++++++++++++++++++++++ 4 files changed, 664 insertions(+), 64 deletions(-) diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 44cb558d..22a9fef5 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -836,6 +836,35 @@ class LLMHandler(ABC): updated_messages = messages.copy() pending_actions: List[Dict] = [] + # One assistant message carries the WHOLE parallel batch, followed by + # one tool message per call — the layout every provider expects, and + # the one the resume (``BaseAgent._resume_*``) and history-replay + # (``BaseAgent._build_messages``) paths already emit. + # + # Emitting a separate assistant message per call breaks chained + # Responses requests: ``_trim_for_previous_response`` re-sends only + # what follows the LAST assistant message, so every output except the + # final one was dropped and the provider rejected the request with + # "No tool output found for function call ". + batch_assistant: Optional[Dict[str, Any]] = None + + def _declare_call(tool_call_obj: Dict[str, Any]) -> None: + """Add ``tool_call_obj`` to the batch's assistant message, + creating (and appending) that message on first use.""" + nonlocal batch_assistant + if batch_assistant is None: + batch_assistant = { + "role": "assistant", + "content": None, + "tool_calls": [], + } + # Reasoning rides once on the batch message; DeepSeek thinking + # mode rejects an active-turn assistant message without it. + if reasoning_content: + batch_assistant["reasoning_content"] = reasoning_content + updated_messages.append(batch_assistant) + batch_assistant["tool_calls"].append(tool_call_obj) + for i, call in enumerate(tool_calls): # Check context limit before executing tool call if hasattr(agent, '_check_context_limit') and agent._check_context_limit(updated_messages): @@ -865,6 +894,12 @@ class LLMHandler(ABC): if compression_successful and rebuilt_messages is not None: # Update the messages list with rebuilt compressed version updated_messages = rebuilt_messages + # The rebuilt list no longer contains the batch + # message we were appending to, so mutating it + # would be a silent no-op. Start a fresh one for + # the remaining calls; each assistant message + # still owns exactly the results that follow it. + batch_assistant = None # Yield compression success message yield { @@ -969,14 +1004,7 @@ class LLMHandler(ABC): } if getattr(call, "thought_signature", None): tool_call_obj["thought_signature"] = call.thought_signature - assistant_msg: Dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [tool_call_obj], - } - if reasoning_content: - assistant_msg["reasoning_content"] = reasoning_content - updated_messages.append(assistant_msg) + _declare_call(tool_call_obj) denial_call = ToolCall( id=pause_info["call_id"], name=call.name, @@ -1081,19 +1109,7 @@ class LLMHandler(ABC): if call.thought_signature: tool_call_obj["thought_signature"] = call.thought_signature - assistant_msg: Dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [tool_call_obj], - } - # Each call in a parallel batch becomes its own - # assistant message here, so the same per-round - # reasoning has to ride on every one — DeepSeek - # thinking mode rejects any assistant message in the - # active turn that's missing reasoning_content. - if reasoning_content: - assistant_msg["reasoning_content"] = reasoning_content - updated_messages.append(assistant_msg) + _declare_call(tool_call_obj) assistant_appended = True # The tool result's tool_call_id must match the id put on the @@ -1110,24 +1126,22 @@ class LLMHandler(ABC): ) except Exception as e: logger.error(f"Error executing tool: {str(e)}", exc_info=True) - # The error tool message's tool_call_id must match the assistant - # tool_call id that precedes it. When the success path already - # appended one (create_tool_message failed afterwards) that id is - # the executor-returned ``call_id``; otherwise the except builds - # its own assistant message below from ``call.id``. + # The error tool message's tool_call_id must match the + # tool_call id declared on the batch assistant message. When + # the success path already declared this call that id is the + # executor-returned ``call_id``; otherwise the except declares + # it below from ``call.id``. error_id = call_id if assistant_appended else call.id error_call = ToolCall( id=error_id, name=call.name, arguments=call.arguments ) error_response = f"Error executing tool: {str(e)}" - # Mirror the success path: a role:"tool" message must follow - # an assistant message carrying its tool_calls, or the next - # provider completion 400s ("'tool' must be a response to a - # preceding message with 'tool_calls'"). Skip it when the - # success path already appended one for this call — a - # create_tool_message failure after that append would - # otherwise duplicate the assistant message and 400 the - # same way an orphan tool message does. + # Mirror the success path: every tool message must answer a + # call declared on the batch assistant message, or the next + # provider completion 400s. Skip re-declaring when the success + # path already did it for this call — a create_tool_message + # failure after that point would otherwise declare the call + # twice and 400 the same way an orphan tool message does. if not assistant_appended: args_str = ( json.dumps(call.arguments) @@ -1144,14 +1158,7 @@ class LLMHandler(ABC): } if call.thought_signature: tool_call_obj["thought_signature"] = call.thought_signature - assistant_msg: Dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [tool_call_obj], - } - if reasoning_content: - assistant_msg["reasoning_content"] = reasoning_content - updated_messages.append(assistant_msg) + _declare_call(tool_call_obj) error_message = self.create_tool_message(error_call, error_response) updated_messages.append(error_message) diff --git a/application/llm/openai.py b/application/llm/openai.py index 75555546..48b90726 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -158,6 +158,9 @@ class OpenAILLM(BaseLLM): # chain turns when OPENAI_RESPONSES_STORE is enabled. self._reasoning_for_calls = {} self._last_response_id = None + # call_ids the most recent response emitted; the chained-request + # coverage guard checks the next turn answers all of them. + self._last_response_call_ids = set() self._last_reasoning_items = [] self._last_usage = None # One-shot guard consumed by ``_prefer_provider_usage``: fresh @@ -222,6 +225,7 @@ class OpenAILLM(BaseLLM): self._reasoning_for_calls = {} self._last_reasoning_items = [] self._last_response_id = None + self._last_response_call_ids = set() self._imported_response_id = None self._last_finish_reason = None @@ -735,20 +739,100 @@ class OpenAILLM(BaseLLM): @staticmethod def _trim_for_previous_response(messages): """When chaining via ``previous_response_id`` the server already holds - the earlier turns, so only system context plus everything after the - last completed assistant response needs to be sent again.""" + the earlier turns, so only system context plus the tool results + answering the chained response's calls need to be sent again. + + The cut is the last assistant message. Any ``tool`` message inside the + trailing assistant/tool run *before* that cut is carried over too: a + batch written as one assistant message per call — instead of one + message carrying the whole batch — would otherwise lose every result + but the last, and the provider then rejects the request with "No tool + output found for function call ". The assistant + messages in that run are not carried: the server already holds those + calls, and chained mode accepts bare ``function_call_output`` items. + """ last_assistant = -1 for i, message in enumerate(messages): if message.get("role") == "assistant": last_assistant = i if last_assistant < 0: return messages + carried = [] + i = last_assistant - 1 + while i >= 0: + role = messages[i].get("role") + if role == "tool": + carried.append(messages[i]) + elif not (role == "assistant" and messages[i].get("tool_calls")): + break + i -= 1 + carried.reverse() head = [ m for m in messages[: last_assistant + 1] if m.get("role") == "system" ] - return head + messages[last_assistant + 1:] + return head + carried + messages[last_assistant + 1:] + + def _build_responses_input(self, messages, previous_response_id): + """Build the Responses ``input`` list, honouring store-mode chaining. + + Returns ``(input_items, previous_response_id)``. The id comes back + ``None`` when the chained payload would not carry an output for every + call the chained response holds — rather than send a request the + provider is certain to reject, the full history is sent unchained. + + A chained request answers exactly the calls of the response it chains + to, so outputs are checked in BOTH directions against + ``_last_response_call_ids``: missing ones abandon chaining, extra ones + are dropped. The extras matter because the trim's carry loop cannot + tell a split batch (whose stray results must be rescued) from a + multi-round grouped history (whose earlier rounds the server already + has) — both end in the same assistant/tool run. Without the filter, + round N re-sends every prior round's outputs back to the last user + turn and the payload grows with the tool loop. + """ + chained = bool(previous_response_id and settings.OPENAI_RESPONSES_STORE) + if not chained: + return self._to_responses_input(messages, chained=False), None + + input_items = self._to_responses_input( + self._trim_for_previous_response(messages), chained=True + ) + expected = set(self._last_response_call_ids or ()) + if expected: + kept, dropped = [], [] + for item in input_items: + if ( + item.get("type") == "function_call_output" + and item.get("call_id") not in expected + ): + dropped.append(item.get("call_id")) + else: + kept.append(item) + if dropped: + logging.debug( + "Dropping %d already-answered tool output(s) from the " + "chained Responses input (call_ids=%s)", + len(dropped), + sorted(dropped), + ) + input_items = kept + sent = { + item.get("call_id") + for item in input_items + if item.get("type") == "function_call_output" + } + missing = expected - sent + if missing: + logging.warning( + "Chained Responses request would omit tool output(s) for " + "%s; sending the full unchained input instead of a " + "request the provider would reject", + sorted(missing), + ) + return self._to_responses_input(messages, chained=False), None + return input_items, previous_response_id @staticmethod def _to_responses_tools(tools): @@ -909,10 +993,33 @@ class OpenAILLM(BaseLLM): self._last_usage = result self._last_usage_claimed = False + @staticmethod + def _function_call_ids(response): + """call_ids of every ``function_call`` item in a Responses output. + + Feeds the chained-request coverage guard in + ``_build_responses_input``: these are the calls the server will expect + a ``function_call_output`` for on the next chained turn. + """ + output = getattr(response, "output", None) + if not isinstance(output, list): + return set() + call_ids = set() + for item in output: + if isinstance(item, dict): + itype, cid = item.get("type"), item.get("call_id") + else: + itype = getattr(item, "type", None) + cid = getattr(item, "call_id", None) + if itype == "function_call" and isinstance(cid, str) and cid: + call_ids.add(cid) + return call_ids + def _record_responses_metadata(self, response): rid = getattr(response, "id", None) if rid: self._last_response_id = rid + self._last_response_call_ids = self._function_call_ids(response) usage = getattr(response, "usage", None) if usage is not None: prompt = int(getattr(usage, "input_tokens", 0) or 0) @@ -1003,13 +1110,15 @@ class OpenAILLM(BaseLLM): previous_response_id = ( self._last_response_id or previous_response_id or self._imported_response_id ) + # Built before the per-turn state is cleared: the coverage guard reads + # the call_ids the chained response emitted. + input_items, previous_response_id = self._build_responses_input( + messages, previous_response_id + ) self._last_response_id = None + self._last_response_call_ids = set() self._last_usage = None self._last_finish_reason = None - chained = bool(previous_response_id and settings.OPENAI_RESPONSES_STORE) - if chained: - messages = self._trim_for_previous_response(messages) - input_items = self._to_responses_input(messages, chained=chained) params = self._build_responses_params( model, input_items, @@ -1073,14 +1182,16 @@ class OpenAILLM(BaseLLM): previous_response_id = ( self._last_response_id or previous_response_id or self._imported_response_id ) + # Built before the per-turn state is cleared: the coverage guard reads + # the call_ids the chained response emitted. + input_items, previous_response_id = self._build_responses_input( + messages, previous_response_id + ) self._last_response_id = None + self._last_response_call_ids = set() self._last_usage = None self._last_finish_reason = None self._stream_reached_finish = False - chained = bool(previous_response_id and settings.OPENAI_RESPONSES_STORE) - if chained: - messages = self._trim_for_previous_response(messages) - input_items = self._to_responses_input(messages, chained=chained) params = self._build_responses_params( model, input_items, diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index c0a2b6fd..a480b27e 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -281,6 +281,169 @@ class TestHandleToolCallsIdPairing: assert pending is None +class _MultiCallStubAgent: + """Agent whose executor echoes each call's own id, optionally failing some.""" + + def __init__(self, fail_ids=()): + self.llm = Mock() + self.tool_executor = _StubToolExecutor() + self._fail_ids = set(fail_ids) + + def _execute_tool_action(self, tools_dict, call): + yield from () + if call.id in self._fail_ids: + raise RuntimeError(f"boom:{call.id}") + return {"ok": call.id}, call.id + + +class TestHandleToolCallsParallelBatchLayout: + """A parallel batch must produce the standard provider layout: ONE + assistant message carrying every tool_call, followed by one tool message + per call. + + Emitting a separate assistant message per call (the previous behaviour) + breaks ``_trim_for_previous_response``, which keeps only what follows the + LAST assistant message — so every output except the final one is dropped + from a chained Responses request and the provider 400s with + "No tool output found for function call ...". + """ + + def _drain(self, gen): + try: + while True: + next(gen) + except StopIteration as e: + return e.value + + def _calls(self, n): + return [ + ToolCall(id=f"call_{i}", name="search", arguments={"q": str(i)}) + for i in range(n) + ] + + def test_parallel_batch_produces_single_assistant_message(self): + handler = ConcreteHandler() + calls = self._calls(3) + + updated, pending = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), calls, tools_dict={}, messages=[] + ) + ) + + assistants = [m for m in updated if m["role"] == "assistant"] + assert len(assistants) == 1 + assert [tc["id"] for tc in assistants[0]["tool_calls"]] == [ + "call_0", + "call_1", + "call_2", + ] + assert pending is None + + def test_parallel_batch_orders_assistant_before_every_result(self): + handler = ConcreteHandler() + calls = self._calls(3) + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), calls, tools_dict={}, messages=[] + ) + ) + + roles = [m["role"] for m in updated] + assert roles == ["assistant", "tool", "tool", "tool"] + assert [m["tool_call_id"] for m in updated if m["role"] == "tool"] == [ + "call_0", + "call_1", + "call_2", + ] + + def test_every_call_in_the_batch_has_its_output(self): + """The invariant the provider enforces: every tool_call on the + assistant message is answered by a tool message.""" + handler = ConcreteHandler() + calls = self._calls(4) + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), calls, tools_dict={}, messages=[] + ) + ) + + declared = { + tc["id"] + for m in updated + if m["role"] == "assistant" + for tc in m["tool_calls"] + } + answered = {m["tool_call_id"] for m in updated if m["role"] == "tool"} + assert declared == answered + + def test_reasoning_content_attached_once_not_per_call(self): + handler = ConcreteHandler() + calls = self._calls(3) + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), + calls, + tools_dict={}, + messages=[], + reasoning_content="thinking", + ) + ) + + assistants = [m for m in updated if m["role"] == "assistant"] + assert len(assistants) == 1 + assert assistants[0]["reasoning_content"] == "thinking" + + def test_failed_call_shares_the_batch_assistant_message(self): + """A tool that raises still contributes its call to the same assistant + message and gets an error tool message — pairing must hold.""" + handler = ConcreteHandler() + calls = self._calls(3) + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(fail_ids={"call_1"}), + calls, + tools_dict={}, + messages=[], + ) + ) + + assistants = [m for m in updated if m["role"] == "assistant"] + assert len(assistants) == 1 + declared = {tc["id"] for tc in assistants[0]["tool_calls"]} + answered = {m["tool_call_id"] for m in updated if m["role"] == "tool"} + assert declared == {"call_0", "call_1", "call_2"} == answered + + def test_single_call_batch_unchanged(self): + handler = ConcreteHandler() + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), self._calls(1), tools_dict={}, messages=[] + ) + ) + + assert [m["role"] for m in updated] == ["assistant", "tool"] + assert len(updated[0]["tool_calls"]) == 1 + + def test_existing_history_is_preserved(self): + handler = ConcreteHandler() + history = [{"role": "user", "content": "hi"}] + + updated, _ = self._drain( + handler.handle_tool_calls( + _MultiCallStubAgent(), self._calls(2), tools_dict={}, messages=history + ) + ) + + assert updated[0] == {"role": "user", "content": "hi"} + assert [m["role"] for m in updated[1:]] == ["assistant", "tool", "tool"] + + # --------------------------------------------------------------------------- # _append_unsupported_attachments # --------------------------------------------------------------------------- @@ -906,18 +1069,22 @@ class TestHandleToolCalls: except StopIteration as e: messages, _pending = e.value - # Each tool message must immediately follow an assistant message that - # advertises its tool_call_id. - answered_ids = set() - for idx, msg in enumerate(messages): - if msg.get("role") == "tool": - assert idx > 0 - prev = messages[idx - 1] - assert prev["role"] == "assistant" - advertised = {tc["id"] for tc in prev["tool_calls"]} - assert msg["tool_call_id"] in advertised + # Every tool message must answer a tool_call declared by an assistant + # message EARLIER in the array, and every declared call must be + # answered. (Adjacency is deliberately not asserted: the standard + # provider layout is one assistant message carrying the whole batch + # followed by all its results, so only the first result is adjacent.) + declared_ids: set = set() + answered_ids: set = set() + for msg in messages: + if msg.get("role") == "assistant": + declared_ids.update(tc["id"] for tc in msg.get("tool_calls", [])) + elif msg.get("role") == "tool": + assert msg["tool_call_id"] in declared_ids, ( + "tool result precedes the assistant message declaring it" + ) answered_ids.add(msg["tool_call_id"]) - assert answered_ids == {"c1", "c2", "c3"} + assert answered_ids == {"c1", "c2", "c3"} == declared_ids def test_thought_signature_preserved(self): handler = ConcreteHandler() diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 967393a8..7373a4b7 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1100,6 +1100,321 @@ def test_to_responses_input_chained_keeps_bare_tool_output(monkeypatch): assert [o["call_id"] for o in outputs] == ["call_server_side"] +def _grouped_batch_messages(): + """Standard layout: one assistant message, N tool_calls, N tool results.""" + return [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{i}", + "type": "function", + "function": {"name": "read", "arguments": "{}"}, + } + for i in range(3) + ], + }, + {"role": "tool", "tool_call_id": "call_0", "content": "r0"}, + {"role": "tool", "tool_call_id": "call_1", "content": "r1"}, + {"role": "tool", "tool_call_id": "call_2", "content": "r2"}, + ] + + +def _split_batch_messages(): + """Legacy layout: one assistant message PER call, interleaved with its + result. ``_trim_for_previous_response`` must still keep every output.""" + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q"}, + ] + for i in range(3): + msgs.append({ + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": f"call_{i}", + "type": "function", + "function": {"name": "read", "arguments": "{}"}, + }], + }) + msgs.append( + {"role": "tool", "tool_call_id": f"call_{i}", "content": f"r{i}"} + ) + return msgs + + +@pytest.mark.unit +def test_trim_keeps_every_output_of_a_grouped_batch(monkeypatch): + """Chained requests must carry an output for every call the chained + response holds.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + trimmed = llm._trim_for_previous_response(_grouped_batch_messages()) + items = llm._to_responses_input(trimmed, chained=True) + outputs = [i["call_id"] for i in items if i.get("type") == "function_call_output"] + assert outputs == ["call_0", "call_1", "call_2"] + + +@pytest.mark.unit +def test_trim_keeps_every_output_of_a_split_batch(monkeypatch): + """Regression: the trim cut at the LAST assistant message, so a batch + written as one assistant message per call lost every output but the + last — the provider then 400s "No tool output found for function call + call_0" (naming the FIRST unpaired call).""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + trimmed = llm._trim_for_previous_response(_split_batch_messages()) + items = llm._to_responses_input(trimmed, chained=True) + outputs = [i["call_id"] for i in items if i.get("type") == "function_call_output"] + assert outputs == ["call_0", "call_1", "call_2"] + + +@pytest.mark.unit +def test_trim_still_drops_prior_completed_turns(monkeypatch): + """The trim must keep doing its job: a completed assistant answer ends + the replayed region, so earlier turns are not re-sent.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "old q"}, + {"role": "assistant", "content": "old answer"}, + {"role": "user", "content": "new q"}, + ] + trimmed = llm._trim_for_previous_response(messages) + assert [m.get("content") for m in trimmed] == ["sys", "new q"] + + +@pytest.mark.unit +def test_chained_input_falls_back_when_an_output_is_missing(monkeypatch): + """Coverage guard: if the outputs being sent do not cover the calls held + in the chained response, send a full non-chained input instead of a + request the provider is certain to reject.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = {"call_0", "call_1", "call_2"} + # Only the last output survives — the exact shape the old trim produced. + messages = _grouped_batch_messages()[:3] + [ + {"role": "tool", "tool_call_id": "call_2", "content": "r2"}, + ] + + items, prev = llm._build_responses_input(messages, "resp_prev") + + assert prev is None, "must drop out of chained mode" + # The unchained rebuild is internally consistent: the one call that still + # has its output is kept and paired, the two orphans are dropped by the + # pairing guard. Either way the provider gets a request it can accept. + calls = [i["call_id"] for i in items if i.get("type") == "function_call"] + outputs = [i["call_id"] for i in items if i.get("type") == "function_call_output"] + assert calls == ["call_2"] + assert outputs == ["call_2"] + + +@pytest.mark.unit +def test_chained_input_stays_chained_when_outputs_are_complete(monkeypatch): + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = {"call_0", "call_1", "call_2"} + + items, prev = llm._build_responses_input(_grouped_batch_messages(), "resp_prev") + + assert prev == "resp_prev" + outputs = [i["call_id"] for i in items if i.get("type") == "function_call_output"] + assert outputs == ["call_0", "call_1", "call_2"] + + +@pytest.mark.unit +def test_chained_input_skips_the_check_when_calls_are_unknown(monkeypatch): + """After a resume/import we may not know which calls the chained response + holds; an unknown set must not be treated as a mismatch.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = set() + + _, prev = llm._build_responses_input( + [ + {"role": "system", "content": "sys"}, + {"role": "tool", "tool_call_id": "call_server_side", "content": "r"}, + ], + "resp_prev", + ) + + assert prev == "resp_prev" + + +def _multi_round_messages(rounds, calls_per_round=2): + """A grouped history of N *consecutive* tool rounds — no user message in + between, which is what an agent tool loop actually produces.""" + out = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q"}, + ] + for r in range(rounds): + ids = [f"r{r}c{c}" for c in range(calls_per_round)] + out.append({ + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": i, + "type": "function", + "function": {"name": "t", "arguments": "{}"}, + } + for i in ids + ], + }) + out += [ + {"role": "tool", "tool_call_id": i, "content": "res"} for i in ids + ] + return out + + +@pytest.mark.unit +@pytest.mark.parametrize("rounds", [2, 3, 5]) +def test_chained_input_does_not_resend_earlier_rounds_outputs(monkeypatch, rounds): + """Only the chained response's OWN calls get answered. + + The trim's carry loop rescues stray results from a split batch, but the + trailing assistant/tool run of a multi-round grouped history has the same + shape — so the walk also picked up every earlier round's results, which + the server already has. Left unfiltered the payload grows with the round + count, defeating the point of chaining. + """ + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + last = rounds - 1 + llm._last_response_call_ids = {f"r{last}c0", f"r{last}c1"} + + items, prev = llm._build_responses_input( + _multi_round_messages(rounds), "resp_prev" + ) + + assert prev == "resp_prev" + outputs = [ + i["call_id"] for i in items if i.get("type") == "function_call_output" + ] + assert outputs == [f"r{last}c0", f"r{last}c1"], ( + f"re-sent earlier rounds: {outputs}" + ) + + +@pytest.mark.unit +def test_chained_input_still_rescues_a_split_batch(monkeypatch): + """The filter must not undo the split-batch rescue: those outputs all + answer calls the chained response emitted, so they stay.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = {"call_0", "call_1", "call_2"} + + items, prev = llm._build_responses_input(_split_batch_messages(), "resp_prev") + + assert prev == "resp_prev" + outputs = [ + i["call_id"] for i in items if i.get("type") == "function_call_output" + ] + assert outputs == ["call_0", "call_1", "call_2"] + + +@pytest.mark.unit +def test_chained_input_keeps_a_batch_split_by_compression(monkeypatch): + """Mid-batch compression starts a second assistant message for the same + round. Both halves answer the chained response's calls, so both stay.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = {"c0", "c1", "c2"} + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q"}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "c0", "type": "function", + "function": {"name": "t", "arguments": "{}"}}, + ]}, + {"role": "tool", "tool_call_id": "c0", "content": "r0"}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "c1", "type": "function", + "function": {"name": "t", "arguments": "{}"}}, + {"id": "c2", "type": "function", + "function": {"name": "t", "arguments": "{}"}}, + ]}, + {"role": "tool", "tool_call_id": "c1", "content": "r1"}, + {"role": "tool", "tool_call_id": "c2", "content": "r2"}, + ] + + items, _ = llm._build_responses_input(messages, "resp_prev") + + outputs = [ + i["call_id"] for i in items if i.get("type") == "function_call_output" + ] + assert outputs == ["c0", "c1", "c2"] + + +@pytest.mark.unit +def test_chained_input_does_not_filter_when_calls_are_unknown(monkeypatch): + """With no known call set (import/resume) an extra output cannot be told + from a needed one, so nothing is dropped.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_call_ids = set() + + items, prev = llm._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 len(outputs) == 6 + + +@pytest.mark.unit +@pytest.mark.parametrize("entrypoint", ["_responses_gen", "_responses_gen_stream"]) +def test_both_request_builders_route_through_the_coverage_guard( + monkeypatch, entrypoint +): + """Regression: the guard was wired into the non-streaming builder only, + so the streaming path — the one production actually uses — still sent + unchecked chained requests. Both must go through it.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + llm._last_response_id = "resp_prev" + seen = {} + + real = llm._build_responses_input + + def spy(messages, previous_response_id): + seen["called"] = True + return real(messages, previous_response_id) + + monkeypatch.setattr(llm, "_build_responses_input", spy) + monkeypatch.setattr( + llm, "_build_responses_params", lambda *a, **k: {"model": "m"} + ) + # The streaming path iterates the response; the non-streaming one reads + # attributes off it. + llm.client.responses.create = MagicMock( + return_value=iter([]) + if entrypoint == "_responses_gen_stream" + else types.SimpleNamespace(id="r", usage=None, output=[]) + ) + + result = getattr(llm, entrypoint)("m", [{"role": "user", "content": "hi"}]) + if entrypoint == "_responses_gen_stream": + list(result) # drain the generator so the body runs + + assert seen.get("called"), f"{entrypoint} bypassed the coverage guard" + + +@pytest.mark.unit +def test_records_call_ids_emitted_by_a_response(monkeypatch): + """The coverage guard needs to know which calls the response emitted.""" + llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True) + response = _ns( + id="resp_1", + usage=None, + output=[ + _ns(type="function_call", call_id="call_a"), + _ns(type="message", call_id=None), + _ns(type="function_call", call_id="call_b"), + ], + ) + + llm._record_responses_metadata(response) + + assert llm._last_response_id == "resp_1" + assert llm._last_response_call_ids == {"call_a", "call_b"} + + @pytest.mark.unit def test_to_responses_input_drops_out_of_order_pair_entirely(monkeypatch): """An output that precedes its call is malformed history: BOTH sides From 64b3bd33562476aa8685362d300fec7ca71fc991 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 20 Jul 2026 19:57:44 +0100 Subject: [PATCH 3/4] feat: harden webtool --- application/agents/tool_executor.py | 48 +++- application/agents/tools/read_webpage.py | 135 ++++++++++- .../answer/services/conversation_service.py | 21 +- application/api/user/tasks.py | 7 +- application/security/safe_url.py | 84 +++++++ .../storage/db/repositories/attachments.py | 15 +- .../storage/db/repositories/stack_logs.py | 34 ++- .../db/repositories/tool_call_attempts.py | 19 +- application/streaming/message_journal.py | 27 +-- application/utils.py | 24 ++ tests/agents/test_read_webpage_tool.py | 217 ++++++++++++++++-- tests/agents/test_tool_executor_sanitize.py | 60 +++++ .../agents/tools/test_read_webpage_errors.py | 14 +- .../services/test_conversation_service.py | 40 ++++ tests/api/user/test_tasks.py | 10 + tests/security/test_safe_url.py | 129 +++++++++++ .../db/repositories/test_attachments.py | 14 ++ .../db/repositories/test_stack_logs.py | 45 ++++ .../repositories/test_tool_call_attempts.py | 68 ++++++ tests/test_remaining_coverage.py | 8 +- 20 files changed, 939 insertions(+), 80 deletions(-) create mode 100644 tests/agents/test_tool_executor_sanitize.py create mode 100644 tests/storage/db/repositories/test_tool_call_attempts.py diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 66be5d1a..18472612 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -53,6 +53,47 @@ def truncate_tool_result(value: Any) -> Any: return f"{text[:PERSISTED_RESULT_MAX_LEN]}..." +# Control characters minus \t \n \r. NULs in particular are rejected by +# Postgres text/jsonb, so one binary-carrying tool result would otherwise +# kill every write lane it fans out to (conversation, activity log, +# tool_call_attempts) at once. +_CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]") + +# Ceiling on the persisted full copy of a tool result (``result_full`` in +# the conversation row and the ``tool_call_attempts`` result). Generous — +# the LLM copy is bounded separately at TOOL_RESULT_MAX_TOKENS — but a +# ceiling all the same: one uncapped result has reached 634k tokens. +RESULT_FULL_MAX_CHARS = 400_000 + + +def sanitize_tool_result(value: Any) -> Any: + """Recursively strip NUL/control characters from strings in ``value``. + + Applied once where a tool result enters the executor, so every + downstream consumer — the LLM copy, the conversation row, the + ``tool_call_attempts`` journal, the stream event — gets clean text. + """ + if isinstance(value, str): + return _CONTROL_CHARS.sub("", value) if _CONTROL_CHARS.search(value) else value + if isinstance(value, dict): + return {sanitize_tool_result(k): sanitize_tool_result(v) for k, v in value.items()} + if isinstance(value, list): + return [sanitize_tool_result(item) for item in value] + if isinstance(value, tuple): + return tuple(sanitize_tool_result(item) for item in value) + return value + + +def bound_result_full(text: str) -> str: + """Cap the persisted full-result copy at ``RESULT_FULL_MAX_CHARS``.""" + if len(text) <= RESULT_FULL_MAX_CHARS: + return text + return ( + text[:RESULT_FULL_MAX_CHARS] + + f"\n[... tool result truncated at persistence: {len(text)} chars ...]" + ) + + def result_status(result: Any) -> str: """Derive the persisted status from a tool's result payload. @@ -882,6 +923,11 @@ class ToolExecutor: ) raise + # Single fan-out point: from here ``result`` reaches the LLM copy, + # the conversation row, tool_call_attempts, and the stream event — + # sanitize once so every lane gets clean text. + result = sanitize_tool_result(result) + get_artifact_id = getattr(tool, "get_artifact_id", None) if tool_data["name"] != "api_tool" else None artifact_id = None @@ -898,7 +944,7 @@ class ToolExecutor: artifact_id = str(artifact_id).strip() if artifact_id is not None else "" if artifact_id: tool_call_data["artifact_id"] = artifact_id - result_full = str(result) + result_full = bound_result_full(str(result)) tool_call_data["resolved_arguments"] = resolved_arguments tool_call_data["result_full"] = result_full tool_call_data["result"] = truncate_tool_result(result_full) diff --git a/application/agents/tools/read_webpage.py b/application/agents/tools/read_webpage.py index ea71076f..c8b25227 100644 --- a/application/agents/tools/read_webpage.py +++ b/application/agents/tools/read_webpage.py @@ -1,6 +1,107 @@ +import codecs + from markdownify import markdownify + from application.agents.tools.base import Tool -from application.security.safe_url import UnsafeUserUrlError, pinned_request +from application.security.safe_url import ( + ResponseTooLargeError, + UnsafeUserUrlError, + pinned_fetch_bytes, +) + +# Byte ceiling for a fetched page. One uncapped response (a PDF handed to +# this tool on 07-17) produced a 634k-token "text" result that killed the +# conversation save and every other write lane downstream. +MAX_RESPONSE_BYTES = 10 * 1024 * 1024 + +# Non-``text/*`` content types markdownify can still meaningfully convert +# to text (any ``text/*`` is allowed by prefix). Anything else (PDFs, +# images, archives, ...) must be refused with a clear error so the model +# picks another tool instead of receiving binary-as-text. +_ALLOWED_CONTENT_TYPES = frozenset( + { + "application/xhtml+xml", + "application/xml", + "application/json", + "application/ld+json", + "application/rss+xml", + "application/atom+xml", + } +) + +# Magic prefixes of common binary formats, checked in case the server +# omits or lies about the Content-Type header. +_BINARY_MAGIC_PREFIXES = ( + b"%PDF-", # PDF + b"PK\x03\x04", # zip / docx / xlsx + b"\x89PNG", # PNG + b"GIF8", # GIF + b"\xff\xd8\xff", # JPEG + b"\x1f\x8b", # gzip + b"OggS", # ogg + b"\x00\x01\x00\x00", # ttf +) + +_UNSUPPORTED_CONTENT_HINT = ( + "read_webpage only handles HTML/text pages; use a tool suited to this " + "file type instead." +) + + +def _declared_charset(content_type: str) -> str | None: + """Extract the ``charset`` parameter from a Content-Type header value.""" + for part in content_type.split(";")[1:]: + key, _, value = part.strip().partition("=") + if key.strip().lower() == "charset" and value: + return value.strip("\"' ") + return None + + +def _usable_charset(content_type: str) -> str | None: + """Return the declared charset if Python can actually decode with it.""" + charset = _declared_charset(content_type) + if not charset: + return None + try: + codecs.lookup(charset) + except LookupError: + return None + return charset + + +def _is_allowed_media_type(media_type: str) -> bool: + return media_type.startswith("text/") or media_type in _ALLOWED_CONTENT_TYPES + + +def _looks_binary(content: bytes, *, trust_declared_charset: bool = False) -> bool: + """Sniff the first KB for binary magic numbers or NUL bytes. + + The magic-prefix check is unconditional — a lying charset header must + not sneak a PDF through. The NUL check only applies when no usable + charset was declared: legitimate UTF-16 text is NUL-dense, and a + declared charset means the decode path (which strips NULs) handles it. + """ + head = content[:1024] + if head.startswith(_BINARY_MAGIC_PREFIXES): + return True + if trust_declared_charset: + return False + return b"\x00" in head + + +def _decode_body(content: bytes, content_type: str) -> str: + """Decode ``content`` with the declared charset, else UTF-8. + + Never falls back to ``response.text``'s guesses: requests' RFC-2616 + ISO-8859-1 default (text/* without charset) maps every raw byte 1:1, + which is exactly how a PDF once round-tripped into NUL-laden "text". + NULs are stripped after decoding (NUL is valid UTF-8, and the sniff + only sees the first KB) so the tool never returns them regardless of + caller. + """ + charset = _usable_charset(content_type) or "utf-8" + return content.decode(charset, errors="replace").replace("\x00", "") + class ReadWebpageTool(Tool): """ @@ -31,21 +132,42 @@ class ReadWebpageTool(Tool): return "Error: URL parameter is missing." try: - response = pinned_request( - "GET", + content, response = pinned_fetch_bytes( url, + max_bytes=MAX_RESPONSE_BYTES, headers={'User-Agent': 'DocsGPT-Agent/1.0'}, timeout=10, ) response.raise_for_status() - html_content = response.text + content_type = response.headers.get("Content-Type", "") + media_type = content_type.split(";")[0].strip().lower() + if media_type and not _is_allowed_media_type(media_type): + return ( + f"Error: URL returned content type '{media_type}', which " + f"cannot be converted to text. {_UNSUPPORTED_CONTENT_HINT}" + ) + if _looks_binary( + content, + trust_declared_charset=_usable_charset(content_type) is not None, + ): + return ( + "Error: URL returned binary content, which cannot be " + f"converted to text. {_UNSUPPORTED_CONTENT_HINT}" + ) + + html_content = _decode_body(content, content_type) markdown_content = markdownify(html_content, heading_style="ATX", newline_style="BACKSLASH") return markdown_content except UnsafeUserUrlError as e: return f"Error: URL validation failed - {e}" + except ResponseTooLargeError: + return ( + f"Error: The page is too large to read (over " + f"{MAX_RESPONSE_BYTES // (1024 * 1024)} MB)." + ) except Exception as e: return f"Error fetching URL {url}: {e}" @@ -59,8 +181,9 @@ class ReadWebpageTool(Tool): "description": ( "Fetch a webpage and return its content as clean Markdown " "text. Use it whenever the user shares a URL or the answer " - "depends on a specific page. Input must be a fully " - "qualified URL." + "depends on a specific page. Only works for HTML/text " + "pages — not PDFs or other binary files. Input must be a " + "fully qualified URL." ), "parameters": { "type": "object", diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index 1ca64955..f116a274 100644 --- a/application/api/answer/services/conversation_service.py +++ b/application/api/answer/services/conversation_service.py @@ -21,6 +21,7 @@ from application.storage.db.repositories.conversations import ( MessageUpdateOutcome, ) from application.storage.db.session import db_readonly, db_session +from application.utils import strip_null_bytes logger = logging.getLogger(__name__) @@ -338,14 +339,18 @@ class ConversationService: "error", f"{type(error).__name__}: {str(error)}" ) - update_fields: Dict[str, Any] = { - "response": response, - "status": status, - "thought": thought, - "sources": sources, - "tool_calls": tool_calls or [], - "metadata": merged_metadata, - } + # Postgres rejects NUL in text and jsonb; a NUL-laden tool result + # reaching this transaction would lose the whole conversation save. + update_fields: Dict[str, Any] = strip_null_bytes( + { + "response": response, + "status": status, + "thought": thought, + "sources": sources, + "tool_calls": tool_calls or [], + "metadata": merged_metadata, + } + ) if model_id is not None: update_fields["model_id"] = model_id diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 7b45c7cc..fd577e24 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -1,6 +1,8 @@ import logging from datetime import timedelta +from sqlalchemy.exc import DataError + from application.api.user.idempotency import with_idempotency from application.celery_init import celery from application.worker import ( @@ -233,7 +235,10 @@ def _emit_attachment_poison_event(task_name, bound): ) -@celery.task(**DURABLE_TASK) +# ``dont_autoretry_for``: a DataError is deterministic (poison payload, +# e.g. NUL bytes or an over-long value) — retrying re-fails identically +# and multiplies log noise, so it goes straight to the failure path. +@celery.task(**DURABLE_TASK, dont_autoretry_for=(DataError,)) @with_idempotency( task_name="store_attachment", on_poison=_emit_attachment_poison_event, ) diff --git a/application/security/safe_url.py b/application/security/safe_url.py index a4b9eed9..974f928f 100644 --- a/application/security/safe_url.py +++ b/application/security/safe_url.py @@ -79,6 +79,12 @@ class UnsafeUserUrlError(ValueError): """ +class ResponseTooLargeError(Exception): + """Raised by :func:`pinned_fetch_bytes` when a response body exceeds + the caller's byte ceiling (declared via ``Content-Length`` or found + during the bounded read).""" + + def _strip_ipv6_brackets(host: str) -> str: """Return ``host`` with surrounding ``[`` / ``]`` removed if present.""" @@ -348,6 +354,84 @@ def pinned_request( session.close() +def pinned_fetch_bytes( + url: str, + *, + max_bytes: int, + headers: dict[str, str] | None = None, + timeout: float = 10.0, + allow_redirects: bool = False, +) -> tuple[bytes, requests.Response]: + """GET ``url`` with SSRF pinning, streaming at most ``max_bytes`` of body. + + The size limit is enforced twice: a declared ``Content-Length`` above + the ceiling is rejected before any body is read, and the streamed + read itself stops the moment the ceiling is crossed — a server that + lies about (or omits) ``Content-Length`` cannot bypass the cap. + + Returns: + Tuple of the body bytes and the (already-closed) ``Response``, + kept for its status code and headers. + + Raises: + UnsafeUserUrlError: If the URL fails the SSRF guard. + ResponseTooLargeError: If the body exceeds ``max_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)) + try: + response = session.request( + method="GET", + url=pinned_url, + headers=request_headers, + timeout=timeout, + allow_redirects=allow_redirects, + stream=True, + ) + try: + # An error response's body is never used (callers + # raise_for_status) — skip the transfer and let the caller + # surface the more informative HTTP-status error. + if response.status_code >= 400: + return b"", response + declared = response.headers.get("Content-Length") + if declared and declared.isdigit() and int(declared) > max_bytes: + raise ResponseTooLargeError( + f"response declares {declared} bytes, over the " + f"{max_bytes}-byte limit" + ) + chunks: list[bytes] = [] + received = 0 + for chunk in response.iter_content(chunk_size=65536): + received += len(chunk) + if received > max_bytes: + raise ResponseTooLargeError( + f"response body exceeds the {max_bytes}-byte limit" + ) + chunks.append(chunk) + return b"".join(chunks), response + finally: + response.close() + finally: + session.close() + + def pinned_post( url: str, *, diff --git a/application/storage/db/repositories/attachments.py b/application/storage/db/repositories/attachments.py index 08662920..eeddd5e9 100644 --- a/application/storage/db/repositories/attachments.py +++ b/application/storage/db/repositories/attachments.py @@ -8,6 +8,7 @@ from typing import Any, Optional from sqlalchemy import Connection, text from application.storage.db.base_repository import looks_like_uuid, row_to_dict +from application.utils import strip_null_bytes _UPDATABLE_SCALARS = { @@ -69,15 +70,21 @@ class AttachmentsRepository: ), { "user_id": user_id, - "filename": filename, + "filename": strip_null_bytes(filename), "upload_path": upload_path, "mime_type": mime_type, "size": size, - "content": content, + # A parsed document with embedded NULs (mislabeled binary) + # would otherwise fail the INSERT — Postgres rejects \x00. + "content": strip_null_bytes(content), "token_count": token_count, "openai_file_id": openai_file_id, "google_file_uri": google_file_uri, - "metadata": json.dumps(metadata) if metadata is not None else None, + "metadata": ( + json.dumps(strip_null_bytes(metadata)) + if metadata is not None + else None + ), "legacy_mongo_id": legacy_mongo_id, }, ) @@ -193,6 +200,7 @@ class AttachmentsRepository: set_clauses: list[str] = [] params: dict = {"id": attachment_id, "user_id": user_id} for col, val in filtered.items(): + val = strip_null_bytes(val) if col in _UPDATABLE_JSONB: set_clauses.append(f"{col} = CAST(:{col} AS jsonb)") params[col] = json.dumps(val) if val is not None else None @@ -247,6 +255,7 @@ class AttachmentsRepository: set_clauses: list[str] = [] params: dict = {"legacy_id": legacy_mongo_id, "user_id": user_id} for col, val in filtered.items(): + val = strip_null_bytes(val) if col in _UPDATABLE_JSONB: set_clauses.append(f"{col} = CAST(:{col} AS jsonb)") params[col] = json.dumps(val) if val is not None else None diff --git a/application/storage/db/repositories/stack_logs.py b/application/storage/db/repositories/stack_logs.py index fc1a6387..a5ac17fc 100644 --- a/application/storage/db/repositories/stack_logs.py +++ b/application/storage/db/repositories/stack_logs.py @@ -15,9 +15,28 @@ from typing import Optional from application.storage.db.redaction import redact_secrets from application.storage.db.serialization import PGNativeJSONEncoder +from application.utils import strip_null_bytes from sqlalchemy import Connection, text +# Longest string kept inside ``stacks``. The scalar columns are truncated +# by the caller, but stacks used to go in whole — one uncapped tool +# result (634k tokens, 07-17) rode into the activity log through here. +_STACKS_STRING_MAX_LEN = 10000 + + +def _bound_strings(value): + """Recursively truncate strings in ``value`` to ``_STACKS_STRING_MAX_LEN``.""" + if isinstance(value, str): + if len(value) <= _STACKS_STRING_MAX_LEN: + return value + return value[:_STACKS_STRING_MAX_LEN] + "...[truncated]" + if isinstance(value, dict): + return {k: _bound_strings(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_bound_strings(item) for item in value] + return value + class StackLogsRepository: """Postgres-backed replacement for Mongo ``stack_logs`` collection.""" @@ -50,13 +69,16 @@ class StackLogsRepository: ), { "activity_id": activity_id, - "endpoint": endpoint, - "level": level, - "user_id": user_id, - "api_key": api_key, - "query": query, + "endpoint": strip_null_bytes(endpoint), + "level": strip_null_bytes(level), + "user_id": strip_null_bytes(user_id), + "api_key": strip_null_bytes(api_key), + "query": strip_null_bytes(query), "stacks": json.dumps( - redact_secrets(stacks or []), cls=PGNativeJSONEncoder + redact_secrets( + _bound_strings(strip_null_bytes(stacks or [])) + ), + cls=PGNativeJSONEncoder, ), "timestamp": timestamp, }, diff --git a/application/storage/db/repositories/tool_call_attempts.py b/application/storage/db/repositories/tool_call_attempts.py index 0954d733..3ead533b 100644 --- a/application/storage/db/repositories/tool_call_attempts.py +++ b/application/storage/db/repositories/tool_call_attempts.py @@ -8,6 +8,12 @@ from typing import Any, Optional from sqlalchemy import Connection, text from application.storage.db.serialization import PGNativeJSONEncoder +from application.utils import strip_null_bytes + + +def _dump_jsonb(value: Any) -> str: + """Serialize for a jsonb CAST, stripping NULs Postgres would reject.""" + return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder) class ToolCallAttemptsRepository: @@ -56,7 +62,7 @@ class ToolCallAttemptsRepository: "tool_id": tool_id, "tool_name": tool_name, "action_name": action_name, - "arguments": json.dumps(arguments if arguments is not None else {}, cls=PGNativeJSONEncoder), + "arguments": _dump_jsonb(arguments if arguments is not None else {}), "message_id": message_id, "user_id": user_id, "agent_id": agent_id, @@ -120,8 +126,8 @@ class ToolCallAttemptsRepository: "tool_id": tool_id, "tool_name": tool_name, "action_name": action_name, - "arguments": json.dumps(arguments if arguments is not None else {}, cls=PGNativeJSONEncoder), - "result": json.dumps(result_payload, cls=PGNativeJSONEncoder), + "arguments": _dump_jsonb(arguments if arguments is not None else {}), + "result": _dump_jsonb(result_payload), "message_id": message_id, "user_id": user_id, "agent_id": agent_id, @@ -162,7 +168,7 @@ class ToolCallAttemptsRepository: params: dict[str, Any] = { "call_id": call_id, "status": status, - "result": json.dumps(result_payload, cls=PGNativeJSONEncoder), + "result": _dump_jsonb(result_payload), } if message_id is not None: sql += ", message_id = CAST(:message_id AS uuid)" @@ -191,7 +197,10 @@ class ToolCallAttemptsRepository: "UPDATE tool_call_attempts SET status = 'failed', error = :error " "WHERE call_id = :call_id AND status = 'proposed'" ) - params: dict[str, Any] = {"call_id": call_id, "error": error} + params: dict[str, Any] = { + "call_id": call_id, + "error": strip_null_bytes(error), + } if user_id is not None: sql += " AND user_id IS NOT DISTINCT FROM :user_id" params["user_id"] = user_id diff --git a/application/streaming/message_journal.py b/application/streaming/message_journal.py index 30503d03..ef0e94b4 100644 --- a/application/streaming/message_journal.py +++ b/application/streaming/message_journal.py @@ -21,6 +21,7 @@ from application.storage.db.session import db_readonly, db_session from application.streaming.broadcast_channel import Topic from application.streaming.event_replay import encode_pubsub_message from application.streaming.keys import message_topic_name +from application.utils import strip_null_bytes as _strip_null_bytes logger = logging.getLogger(__name__) @@ -34,29 +35,9 @@ DEFAULT_BATCH_SIZE = 16 DEFAULT_BATCH_INTERVAL_MS = 100 -def _strip_null_bytes(value: Any) -> Any: - """Recursively strip ``\\x00`` from string keys/values in ``value``. - - Postgres JSONB rejects the NUL escape; an LLM emitting a stray NUL - in a chunk would otherwise raise ``DataError`` at INSERT and the row - would be lost from the journal (live stream proceeds, reconnect - snapshot misses the chunk). Mirrors the strip already done in - ``parser/embedding_pipeline.py`` and - ``api/user/attachments/routes.py``. - """ - if isinstance(value, str): - return value.replace("\x00", "") if "\x00" in value else value - if isinstance(value, dict): - return { - (k.replace("\x00", "") if isinstance(k, str) and "\x00" in k else k): - _strip_null_bytes(v) - for k, v in value.items() - } - if isinstance(value, list): - return [_strip_null_bytes(item) for item in value] - if isinstance(value, tuple): - return tuple(_strip_null_bytes(item) for item in value) - return value +# NUL stripping lives in ``application.utils.strip_null_bytes`` (shared +# with the conversation-finalize, activity-log, tool_call_attempts, and +# attachments write lanes); imported above as ``_strip_null_bytes``. # Postgres SQLSTATE for a foreign-key violation. ``message_events`` has diff --git a/application/utils.py b/application/utils.py index 176f7720..eea3569c 100644 --- a/application/utils.py +++ b/application/utils.py @@ -55,6 +55,30 @@ def safe_filename(filename): return safe_name +def strip_null_bytes(value): + """Recursively strip ``\\x00`` from string keys/values in ``value``. + + Postgres rejects NUL in both text and jsonb; one NUL-laden payload + (e.g. a binary response mis-decoded to text) would otherwise raise + ``DataError`` and lose the whole row. Shared by the message journal, + conversation finalize, activity log, tool_call_attempts, and + attachments write lanes. + """ + if isinstance(value, str): + return value.replace("\x00", "") if "\x00" in value else value + if isinstance(value, dict): + return { + (k.replace("\x00", "") if isinstance(k, str) and "\x00" in k else k): + strip_null_bytes(v) + for k, v in value.items() + } + if isinstance(value, list): + return [strip_null_bytes(item) for item in value] + if isinstance(value, tuple): + return tuple(strip_null_bytes(item) for item in value) + return value + + def num_tokens_from_string(string: str) -> int: encoding = get_encoding() if isinstance(string, str): diff --git a/tests/agents/test_read_webpage_tool.py b/tests/agents/test_read_webpage_tool.py index ae6e7ab4..f0560f00 100644 --- a/tests/agents/test_read_webpage_tool.py +++ b/tests/agents/test_read_webpage_tool.py @@ -1,11 +1,29 @@ """Tests for application/agents/tools/read_webpage.py""" -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest import requests from application.agents.tools.read_webpage import ReadWebpageTool +from application.security.safe_url import ResponseTooLargeError + + +class _FakeResponse: + def __init__(self, status_code: int = 200, headers=None): + self.status_code = status_code + self.headers = headers if headers is not None else {} + + def raise_for_status(self): + if self.status_code >= 400: + raise requests.exceptions.HTTPError(str(self.status_code)) + + +def _fetch_result(body: bytes, content_type=None, status_code: int = 200): + headers = {} + if content_type is not None: + headers["Content-Type"] = content_type + return body, _FakeResponse(status_code=status_code, headers=headers) @pytest.fixture @@ -25,48 +43,211 @@ class TestReadWebpageExecuteAction: assert "Error" in result assert "URL parameter is missing" in result - @patch("application.agents.tools.read_webpage.pinned_request") - def test_successful_fetch(self, mock_pinned_request, tool): - mock_resp = MagicMock() - mock_resp.status_code = 200 - mock_resp.text = "

Title

Content

" - mock_pinned_request.return_value = mock_resp + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_successful_fetch(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"

Title

Content

", + content_type="text/html; charset=utf-8", + ) result = tool.execute_action("read_webpage", url="https://example.com") assert "Title" in result assert "Content" in result - @patch("application.agents.tools.read_webpage.pinned_request") - def test_request_error(self, mock_pinned_request, tool): - mock_pinned_request.side_effect = requests.exceptions.ConnectionError("refused") + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_request_error(self, mock_fetch, tool): + mock_fetch.side_effect = requests.exceptions.ConnectionError("refused") result = tool.execute_action("read_webpage", url="https://example.com") assert "Error fetching URL" in result - @patch("application.agents.tools.read_webpage.pinned_request") - def test_ssrf_blocked(self, mock_pinned_request, tool): + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_ssrf_blocked(self, mock_fetch, tool): from application.security.safe_url import UnsafeUserUrlError - mock_pinned_request.side_effect = UnsafeUserUrlError("blocked") + mock_fetch.side_effect = UnsafeUserUrlError("blocked") result = tool.execute_action("read_webpage", url="http://169.254.169.254/") assert "Error" in result assert "validation failed" in result - @patch("application.agents.tools.read_webpage.pinned_request") - def test_http_error(self, mock_pinned_request, tool): - mock_resp = MagicMock() - mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError("404") - mock_pinned_request.return_value = mock_resp + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_http_error(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result(b"", status_code=404) result = tool.execute_action("read_webpage", url="https://example.com/404") assert "Error fetching URL" in result +@pytest.mark.unit +class TestReadWebpageContentGuards: + """Regression tests for the PDF-as-text incident: a binary body must + never come back as NUL-laden 'markdown'.""" + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_rejects_pdf_content_type(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"%PDF-1.7 binary...", content_type="application/pdf", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/doc.pdf") + + assert result.startswith("Error") + assert "application/pdf" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_rejects_pdf_magic_without_content_type(self, mock_fetch, tool): + # No Content-Type header at all — the 07-17 incident shape. + mock_fetch.return_value = _fetch_result(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3...") + + result = tool.execute_action("read_webpage", url="https://example.com/doc.pdf") + + assert result.startswith("Error") + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_rejects_nul_laden_body_despite_html_content_type(self, mock_fetch, tool): + # Mislabeled binary: content-type lies, the NUL sniff must catch it. + mock_fetch.return_value = _fetch_result( + b"\x00\x00\x00binary\x00garbage\x00\x00", + content_type="text/html", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/fake") + + assert result.startswith("Error") + assert "\x00" not in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_rejects_octet_stream(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"anything", content_type="application/octet-stream", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/blob") + + assert result.startswith("Error") + assert "application/octet-stream" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_response_too_large(self, mock_fetch, tool): + mock_fetch.side_effect = ResponseTooLargeError("body exceeds 10485760 bytes") + + result = tool.execute_action("read_webpage", url="https://example.com/huge") + + assert result.startswith("Error") + assert "large" in result.lower() + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_utf8_decoded_without_declared_charset(self, mock_fetch, tool): + # text/html with no charset param: requests' RFC-2616 ISO-8859-1 + # fallback would mojibake this; we must default to UTF-8. + mock_fetch.return_value = _fetch_result( + "

café über

".encode("utf-8"), + content_type="text/html", + ) + + result = tool.execute_action("read_webpage", url="https://example.com") + + assert "café über" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_declared_charset_respected(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + "

café

".encode("latin-1"), + content_type="text/html; charset=ISO-8859-1", + ) + + result = tool.execute_action("read_webpage", url="https://example.com") + + assert "café" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_unknown_declared_charset_falls_back_to_utf8(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + "

ok

".encode("utf-8"), + content_type="text/html; charset=not-a-real-charset", + ) + + result = tool.execute_action("read_webpage", url="https://example.com") + + assert "ok" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_text_plain_allowed(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"plain text document", content_type="text/plain", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/readme") + + assert "plain text document" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_utf16_page_with_declared_charset_allowed(self, mock_fetch, tool): + # UTF-16 text is NUL-dense; a correct charset declaration must + # win over the NUL sniff (which is for undeclared/mislabeled bodies). + mock_fetch.return_value = _fetch_result( + "Hello world".encode("utf-16"), + content_type="text/html; charset=utf-16", + ) + + result = tool.execute_action("read_webpage", url="https://example.com") + + assert "Hello world" in result + assert "\x00" not in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_pdf_magic_beats_declared_charset(self, mock_fetch, tool): + # The magic-prefix check stays unconditional: a lying + # ``text/html; charset=utf-8`` header must not sneak a PDF through. + mock_fetch.return_value = _fetch_result( + b"%PDF-1.7 binary...", + content_type="text/html; charset=utf-8", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/fake") + + assert result.startswith("Error") + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_nuls_past_sniff_window_are_stripped(self, mock_fetch, tool): + # The sniff only sees the first KB; NULs beyond it must still + # never leave the tool (self-contained, not reliant on the + # executor's sanitizer). + body = b"

" + b"x" * 1100 + b"\x00\x00 tail

" + mock_fetch.return_value = _fetch_result(body, content_type="text/html") + + result = tool.execute_action("read_webpage", url="https://example.com") + + assert "\x00" not in result + assert "tail" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_text_csv_allowed(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"name,qty\nwidget,2", content_type="text/csv", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/data.csv") + + assert "widget" in result + + @patch("application.agents.tools.read_webpage.pinned_fetch_bytes") + def test_rss_feed_allowed(self, mock_fetch, tool): + mock_fetch.return_value = _fetch_result( + b"Feed title", + content_type="application/rss+xml", + ) + + result = tool.execute_action("read_webpage", url="https://example.com/feed") + + assert "Feed title" in result + + @pytest.mark.unit class TestReadWebpageMetadata: def test_actions_metadata(self, tool): diff --git a/tests/agents/test_tool_executor_sanitize.py b/tests/agents/test_tool_executor_sanitize.py new file mode 100644 index 00000000..2cc87ec8 --- /dev/null +++ b/tests/agents/test_tool_executor_sanitize.py @@ -0,0 +1,60 @@ +"""Tests for tool-result sanitization at the executor fan-out point. + +``result_full`` fans out to the conversation row, ``tool_call_attempts``, +and (via the executor's return value) the LLM copy — sanitizing at the +source protects every lane at once. +""" + +import pytest + +from application.agents.tool_executor import ( + RESULT_FULL_MAX_CHARS, + bound_result_full, + sanitize_tool_result, +) + + +@pytest.mark.unit +class TestSanitizeToolResult: + def test_strips_nul_bytes(self): + assert sanitize_tool_result("a\x00b\x00c") == "abc" + + def test_strips_control_chars_keeps_whitespace(self): + assert ( + sanitize_tool_result("line1\nline2\ttab\rret\x01\x08\x0b\x1f\x7fend") + == "line1\nline2\ttab\rret" + "end" + ) + + def test_clean_string_returned_unchanged(self): + s = "perfectly normal résumé text\nwith newlines" + assert sanitize_tool_result(s) is s + + def test_recurses_into_dicts_including_keys(self): + got = sanitize_tool_result({"k\x00ey": {"inner": "v\x00al"}}) + assert got == {"key": {"inner": "val"}} + + def test_recurses_into_lists_and_tuples(self): + assert sanitize_tool_result(["a\x00", ("b\x00",)]) == ["a", ("b",)] + + def test_non_string_scalars_pass_through(self): + for value in (42, 3.14, True, None, b"\x00raw"): + assert sanitize_tool_result(value) == value + + +@pytest.mark.unit +class TestBoundResultFull: + def test_short_text_unchanged(self): + assert bound_result_full("short") == "short" + + def test_text_at_limit_unchanged(self): + s = "x" * RESULT_FULL_MAX_CHARS + assert bound_result_full(s) == s + + def test_oversized_text_truncated_with_marker(self): + s = "x" * (RESULT_FULL_MAX_CHARS + 1000) + got = bound_result_full(s) + assert len(got) < len(s) + assert "truncated" in got + # The marker names the original size so the audit copy is honest + # about what was dropped. + assert str(len(s)) in got diff --git a/tests/agents/tools/test_read_webpage_errors.py b/tests/agents/tools/test_read_webpage_errors.py index 05848b17..5f708f97 100644 --- a/tests/agents/tools/test_read_webpage_errors.py +++ b/tests/agents/tools/test_read_webpage_errors.py @@ -11,7 +11,7 @@ class TestReadWebpageErrors: tool = ReadWebpageTool(config={}) with patch( - "application.agents.tools.read_webpage.pinned_request", + "application.agents.tools.read_webpage.pinned_fetch_bytes", side_effect=requests.exceptions.RequestException("bad url"), ): got = tool.execute_action( @@ -20,17 +20,21 @@ class TestReadWebpageErrors: assert "Error fetching URL" in got def test_generic_exception_returns_error_string(self): + from unittest.mock import MagicMock + from application.agents.tools.read_webpage import ReadWebpageTool tool = ReadWebpageTool(config={}) + response = MagicMock() + response.headers = {"Content-Type": "text/html"} + response.raise_for_status.return_value = None with patch( "application.agents.tools.read_webpage.markdownify", side_effect=RuntimeError("boom"), ), patch( - "application.agents.tools.read_webpage.pinned_request", - ) as mock_pinned_request: - mock_pinned_request.return_value.text = "

hi

" - mock_pinned_request.return_value.raise_for_status.return_value = None + "application.agents.tools.read_webpage.pinned_fetch_bytes", + return_value=(b"

hi

", response), + ): got = tool.execute_action( "read_webpage", url="https://example.com/", ) diff --git a/tests/api/answer/services/test_conversation_service.py b/tests/api/answer/services/test_conversation_service.py index e016bcb5..70c05e41 100644 --- a/tests/api/answer/services/test_conversation_service.py +++ b/tests/api/answer/services/test_conversation_service.py @@ -474,6 +474,46 @@ class TestFinalizeMessage: assert len(msgs[0]["sources"][0]["text"]) == 1000 assert msgs[0]["metadata"]["foo"] == "bar" + def test_finalize_strips_null_bytes(self, pg_conn): + """A NUL-laden tool result (07-17 PDF incident) must not kill the + conversation save — Postgres rejects \\x00 in text and jsonb.""" + from application.api.answer.services.conversation_service import ( + ConversationService, + ) + from application.storage.db.repositories.conversations import ( + ConversationsRepository, + MessageUpdateOutcome, + ) + + user = "u-fin-nul" + with _patch_db(pg_conn): + svc = ConversationService() + res = svc.save_user_question( + conversation_id=None, + question="q", + decoded_token={"sub": user}, + ) + assert svc.finalize_message( + res["message_id"], + "answer\x00with\x00nuls", + thought="th\x00ought", + sources=[{"text": "s\x00rc", "title": "t\x00itle"}], + tool_calls=[ + {"name": "read_webpage", "result_full": "pdf\x00garbage\x00"} + ], + metadata={"k\x00ey": "v\x00al"}, + status="complete", + ) is MessageUpdateOutcome.UPDATED + + msgs = ConversationsRepository(pg_conn).get_messages( + res["conversation_id"], + ) + assert msgs[0]["response"] == "answerwithnuls" + assert msgs[0]["thought"] == "thought" + assert msgs[0]["sources"][0]["text"] == "src" + assert msgs[0]["tool_calls"][0]["result_full"] == "pdfgarbage" + assert msgs[0]["metadata"] == {"key": "val"} + def test_finalizes_failed_records_error(self, pg_conn): from application.api.answer.services.conversation_service import ( ConversationService, diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index adb13638..e9094577 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -182,6 +182,16 @@ class TestStoreAttachmentTask: mock_worker.assert_called_once_with(ANY, {"file": "info"}, "user1") assert result == {"status": "ok"} + @pytest.mark.unit + def test_data_errors_are_not_autoretried(self): + # A DataError is deterministic (poison payload) — retrying it five + # times just multiplies log noise for the same terminal failure. + from sqlalchemy.exc import DataError + + from application.api.user.tasks import store_attachment + + assert DataError in getattr(store_attachment, "dont_autoretry_for", ()) + class TestProcessAgentWebhookTask: @pytest.mark.unit diff --git a/tests/security/test_safe_url.py b/tests/security/test_safe_url.py index 32e6eaa1..301f6f6c 100644 --- a/tests/security/test_safe_url.py +++ b/tests/security/test_safe_url.py @@ -15,8 +15,10 @@ import pytest import requests from application.security.safe_url import ( + ResponseTooLargeError, UnsafeUserUrlError, _PinnedHTTPSTransport, + pinned_fetch_bytes, pinned_httpx_client, pinned_post, validate_user_base_url, @@ -725,3 +727,130 @@ def test_pinned_httpx_transport_refuses_unexpected_host(monkeypatch): pass finally: client.close() + + +# pinned_fetch_bytes: streamed, size-bounded GET + + +class _StreamStubResponse: + """Streamed drop-in for ``requests.Response`` (``iter_content``-based).""" + + def __init__(self, body: bytes = b"", status_code: int = 200, headers=None): + self._body = body + self.status_code = status_code + self.headers = headers if headers is not None else {} + self.closed = False + + def iter_content(self, chunk_size=1): + for i in range(0, len(self._body), chunk_size): + yield self._body[i:i + chunk_size] + + def close(self): + self.closed = True + + def raise_for_status(self): + if self.status_code >= 400: + raise requests.exceptions.HTTPError(str(self.status_code)) + + +def _capture_stream_send(monkeypatch, response: _StreamStubResponse): + """Replace ``Session.send`` with a stub returning ``response``.""" + + captured: dict = {} + + def _send(self, prepared, **kwargs): + captured["prepared"] = prepared + captured["send_kwargs"] = kwargs + return response + + monkeypatch.setattr(requests.Session, "send", _send) + return captured + + +@pytest.mark.unit +def test_pinned_fetch_bytes_returns_body_and_response(monkeypatch): + body = b"hello" + stub = _StreamStubResponse(body, headers={"Content-Type": "text/html"}) + captured = _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + content, response = pinned_fetch_bytes( + "http://example.com/page", max_bytes=1024, + ) + assert content == body + assert response.headers["Content-Type"] == "text/html" + # The request must stream (no full-body buffering by requests) and + # still carry the pinned-IP URL rewrite + original Host header. + assert captured["send_kwargs"].get("stream") is True + assert captured["prepared"].url == "http://93.184.216.34/page" + assert captured["prepared"].headers["Host"] == "example.com" + # Response closed before the session is torn down. + assert stub.closed is True + + +@pytest.mark.unit +def test_pinned_fetch_bytes_rejects_declared_content_length_over_limit(monkeypatch): + stub = _StreamStubResponse( + b"x" * 100, headers={"Content-Length": "100"}, + ) + _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + with pytest.raises(ResponseTooLargeError): + pinned_fetch_bytes("http://example.com/big", max_bytes=10) + assert stub.closed is True + + +@pytest.mark.unit +def test_pinned_fetch_bytes_rejects_oversized_body_without_content_length(monkeypatch): + # No Content-Length header — the limit must be enforced during the + # bounded read, not just from the declared size. + stub = _StreamStubResponse(b"x" * 100) + _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + with pytest.raises(ResponseTooLargeError): + pinned_fetch_bytes("http://example.com/big", max_bytes=10) + assert stub.closed is True + + +@pytest.mark.unit +def test_pinned_fetch_bytes_body_exactly_at_limit_passes(monkeypatch): + stub = _StreamStubResponse(b"x" * 10, headers={"Content-Length": "10"}) + _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + content, _ = pinned_fetch_bytes("http://example.com/ok", max_bytes=10) + assert content == b"x" * 10 + + +@pytest.mark.unit +def test_pinned_fetch_bytes_rejects_unsafe_url(): + with pytest.raises(UnsafeUserUrlError): + pinned_fetch_bytes("http://127.0.0.1/latest/meta-data", max_bytes=10) + + +@pytest.mark.unit +def test_pinned_fetch_bytes_skips_body_read_on_error_status(monkeypatch): + # An error response's body is dead weight — don't spend up to + # max_bytes of transfer on it, and don't let its size trump the + # more informative HTTP-status error in the caller. + stub = _StreamStubResponse( + b"x" * 100, status_code=404, headers={"Content-Length": "100"}, + ) + _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + content, response = pinned_fetch_bytes( + "http://example.com/missing", max_bytes=10, + ) + assert content == b"" + assert response.status_code == 404 + assert stub.closed is True + + +@pytest.mark.unit +def test_pinned_fetch_bytes_passes_headers(monkeypatch): + stub = _StreamStubResponse(b"ok") + captured = _capture_stream_send(monkeypatch, stub) + with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")): + pinned_fetch_bytes( + "http://example.com/", max_bytes=10, + headers={"User-Agent": "DocsGPT-Agent/1.0"}, + ) + assert captured["prepared"].headers["User-Agent"] == "DocsGPT-Agent/1.0" diff --git a/tests/storage/db/repositories/test_attachments.py b/tests/storage/db/repositories/test_attachments.py index 6420fb26..ad2a2626 100644 --- a/tests/storage/db/repositories/test_attachments.py +++ b/tests/storage/db/repositories/test_attachments.py @@ -58,6 +58,20 @@ class TestCreate: ) assert doc["legacy_mongo_id"] == "507f1f77bcf86cd799439011" + def test_create_strips_null_bytes_from_content_and_metadata(self, pg_conn): + # A parsed document with embedded NULs (mislabeled binary) must + # not kill the INSERT — Postgres rejects \x00 in text and jsonb. + repo = _repo(pg_conn) + doc = repo.create( + "u", + "weird.pdf", + "/uploads/weird.pdf", + content="parsed\x00text\x00here", + metadata={"transcript_lang\x00": "en\x00"}, + ) + assert doc["content"] == "parsedtexthere" + assert doc["metadata"] == {"transcript_lang": "en"} + class TestGet: def test_get_existing(self, pg_conn): diff --git a/tests/storage/db/repositories/test_stack_logs.py b/tests/storage/db/repositories/test_stack_logs.py index 44325e7d..300ad5e1 100644 --- a/tests/storage/db/repositories/test_stack_logs.py +++ b/tests/storage/db/repositories/test_stack_logs.py @@ -126,3 +126,48 @@ class TestInsert: assert data["prompt_tokens"] == 10 assert data["generated_tokens"] == 5 assert data["token_budget"] == 1000 + + def test_strips_null_bytes_from_stacks(self, pg_conn): + # Postgres jsonb rejects the NUL escape — a NUL-laden tool result in the + # stacks would otherwise kill the whole activity-log INSERT (and + # with it the error visibility for the incident that caused it). + repo = _repo(pg_conn) + repo.insert( + activity_id="act-nul", + level="error", + query="q\x00uery", + stacks=[{"component": "tool", "data": {"result": "pdf\x00junk\x00"}}], + ) + row = pg_conn.execute( + text("SELECT query, stacks FROM stack_logs WHERE activity_id = 'act-nul'") + ).fetchone() + mapping = dict(row._mapping) + assert mapping["query"] == "query" + assert mapping["stacks"][0]["data"]["result"] == "pdfjunk" + + def test_bounds_oversized_strings_in_stacks(self, pg_conn): + # The 07-17 incident put a 634k-token tool result into stacks + # whole; long strings must be truncated before insert. + repo = _repo(pg_conn) + repo.insert( + activity_id="act-big", + level="info", + stacks=[{"component": "tool", "data": {"result": "x" * 50000}}], + ) + row = pg_conn.execute( + text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-big'") + ).fetchone() + stored = dict(row._mapping)["stacks"][0]["data"]["result"] + assert len(stored) <= 10100 # 10k cap plus a short truncation marker + + def test_short_strings_in_stacks_untouched(self, pg_conn): + repo = _repo(pg_conn) + repo.insert( + activity_id="act-small", + level="info", + stacks=[{"component": "tool", "data": {"result": "ok"}}], + ) + row = pg_conn.execute( + text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-small'") + ).fetchone() + assert dict(row._mapping)["stacks"][0]["data"]["result"] == "ok" diff --git a/tests/storage/db/repositories/test_tool_call_attempts.py b/tests/storage/db/repositories/test_tool_call_attempts.py new file mode 100644 index 00000000..b536067d --- /dev/null +++ b/tests/storage/db/repositories/test_tool_call_attempts.py @@ -0,0 +1,68 @@ +"""Tests for ToolCallAttemptsRepository against a real Postgres instance.""" + +from __future__ import annotations + +from sqlalchemy import text + +from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, +) + + +def _repo(conn) -> ToolCallAttemptsRepository: + return ToolCallAttemptsRepository(conn) + + +class TestNulSafety: + """Postgres jsonb rejects \\u0000; a NUL-laden tool result must not + lose the attempt row (07-17 PDF incident).""" + + def test_record_proposed_strips_nul_from_arguments(self, pg_conn): + repo = _repo(pg_conn) + assert repo.record_proposed( + "call-nul-1", "read_webpage", "read_webpage", + {"url": "https://x\x00.example"}, + user_id="u1", + ) + row = pg_conn.execute( + text( + "SELECT arguments FROM tool_call_attempts " + "WHERE call_id = 'call-nul-1'" + ) + ).fetchone() + assert dict(row._mapping)["arguments"] == {"url": "https://x.example"} + + def test_mark_executed_strips_nul_from_result(self, pg_conn): + repo = _repo(pg_conn) + assert repo.record_proposed( + "call-nul-2", "read_webpage", "read_webpage", {}, user_id="u1", + ) + assert repo.mark_executed( + "call-nul-2", "pdf\x00garbage\x00text", user_id="u1", + ) + row = pg_conn.execute( + text( + "SELECT status, result FROM tool_call_attempts " + "WHERE call_id = 'call-nul-2'" + ) + ).fetchone() + mapping = dict(row._mapping) + assert mapping["status"] == "confirmed" # no message_id → confirmed + assert mapping["result"]["result"] == "pdfgarbagetext" + + def test_upsert_executed_strips_nul_from_result(self, pg_conn): + repo = _repo(pg_conn) + repo.upsert_executed( + "call-nul-3", "read_webpage", "read_webpage", + {"q": "a\x00b"}, "res\x00ult", + user_id="u1", + ) + row = pg_conn.execute( + text( + "SELECT arguments, result FROM tool_call_attempts " + "WHERE call_id = 'call-nul-3'" + ) + ).fetchone() + mapping = dict(row._mapping) + assert mapping["arguments"] == {"q": "ab"} + assert mapping["result"]["result"] == "result" diff --git a/tests/test_remaining_coverage.py b/tests/test_remaining_coverage.py index 91c0a7a4..6899cb81 100644 --- a/tests/test_remaining_coverage.py +++ b/tests/test_remaining_coverage.py @@ -582,12 +582,12 @@ class TestReadWebpageErrors: tool = ReadWebpageTool({}) with patch( - "application.agents.tools.read_webpage.pinned_request" - ) as mock_pinned_request: + "application.agents.tools.read_webpage.pinned_fetch_bytes" + ) as mock_fetch: mock_response = MagicMock() mock_response.raise_for_status.return_value = None - mock_response.text = "test" - mock_pinned_request.return_value = mock_response + mock_response.headers = {"Content-Type": "text/html"} + mock_fetch.return_value = (b"test", mock_response) with patch( "application.agents.tools.read_webpage.markdownify", side_effect=Exception("parse error"), From e4265ce38a7910b4819942657a64d7d337d0c7b4 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 20 Jul 2026 20:35:18 +0100 Subject: [PATCH 4/4] 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())