From ca5f80995da0e6af31877ffb54a45d0eafad35b2 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 15 Jul 2026 21:20:39 +0100 Subject: [PATCH] fix: accurate per-call token usage and oversized-context guards Token accounting: - Drain each tool round's provider stream to exhaustion before running tools and recursing, so the usage decorator persists exactly one token_usage row per LLM call, at call end. Previously every round's generator was abandoned mid-iteration and flushed together at request teardown, writing N near-identical rows stamped with the final round's provider counts (duplicate billing). - Consume the Chat Completions include_usage terminal chunk (it arrives after finish_reason and was never read) so streamed calls record provider-exact token counts instead of tiktoken estimates. - Claim provider-reported usage once per call (_last_usage_claimed) so a late-finalized generator can never adopt another call's counts. Oversized-context guards: - Enforce Responses API function_call/function_call_output pairing in the input builder (drop unpaired items; bypassed for store-mode previous_response_id chaining where calls are matched server-side). - Hard pre-send context gate: shrink oversized tool results and refuse payloads that cannot fit the model's window before dispatch, so a hopeless request is never sent or billed. - Cap a single tool result entering the LLM context (TOOL_RESULT_MAX_TOKENS, default 20000); the tool journal and persistence keep the full result. Applied on the resume/continuation path too. - Skip the fallback attempt when the payload cannot fit the fallback model's context window (10% estimation slack). - Compression: never save a compression point that does not reduce tokens; bound oversized verbatim fields kept after a compression point (COMPRESSION_RECENT_FIELD_MAX_TOKENS, default 8000). Robustness fixes from review: - Google parallel function calls: complete index-less ToolCalls are no longer merged into one another (dict arguments raised TypeError on +=; second call could execute with the first call's arguments). - Trailing-frame failures after a delivered answer no longer error the stream or restream the whole answer from the fallback (_stream_reached_finish). - In-memory compression falls back to minimal pruning when the summary is not smaller than the original. - keep<=0 guard in the middle-truncation helpers (a tiny cap returned marker + full text). Frontend: tooltip on the Analytics tokens stat card explaining that agent tool loops re-send conversation context on every step. --- application/agents/base.py | 68 ++++- .../answer/services/compression/service.py | 98 +++++- application/core/settings.py | 2 + application/llm/base.py | 54 ++++ application/llm/handlers/base.py | 285 +++++++++++++----- application/llm/openai.py | 113 ++++++- application/usage.py | 13 + frontend/src/locale/en.json | 1 + frontend/src/settings/Analytics.tsx | 6 +- .../services/compression/test_service.py | 125 ++++++++ tests/llm/handlers/test_context_gate.py | 206 +++++++++++++ tests/llm/handlers/test_review_regressions.py | 263 ++++++++++++++++ .../handlers/test_round_usage_persistence.py | 220 ++++++++++++++ tests/llm/test_fallback.py | 121 ++++++++ tests/llm/test_openai_responses.py | 159 ++++++++-- tests/test_compression_service.py | 11 + 16 files changed, 1626 insertions(+), 119 deletions(-) create mode 100644 tests/llm/handlers/test_context_gate.py create mode 100644 tests/llm/handlers/test_review_regressions.py create mode 100644 tests/llm/handlers/test_round_usage_persistence.py diff --git a/application/agents/base.py b/application/agents/base.py index 061b726c..2137982c 100644 --- a/application/agents/base.py +++ b/application/agents/base.py @@ -14,7 +14,10 @@ from application.core.json_schema_utils import ( normalize_json_schema_payload, ) from application.core.settings import settings -from application.llm.handlers.base import ToolCall +from application.llm.handlers.base import ( + ToolCall, + _bound_tool_response_for_llm, +) from application.llm.handlers.handler_creator import LLMHandlerCreator from application.llm.llm_creator import LLMCreator from application.logging import build_stack_data, log_activity, LogContext @@ -316,6 +319,9 @@ class BaseAgent(ABC): except StopIteration as e: tool_response, _ = e.value break + # Same per-result cap as the in-loop path + # (handle_tool_calls); the journal keeps the full result. + tool_response = _bound_tool_response_for_llm(tool_response) messages.append( self.llm_handler.create_tool_message(tc, tool_response) ) @@ -355,7 +361,11 @@ class BaseAgent(ABC): id=call_id, name=pending["name"], arguments=args ) messages.append( - self.llm_handler.create_tool_message(tc, result_str) + self.llm_handler.create_tool_message( + # Client-supplied results get the same per-result + # cap as server-side tool executions. + tc, _bound_tool_response_for_llm(result_str) + ) ) yield { "type": "tool_call", @@ -483,6 +493,10 @@ class BaseAgent(ABC): end_chars = int(target_chars * 0.4) truncation_marker = "\n\n[... content truncated to fit context limit ...]\n\n" + if end_chars <= 0: + # ``text[-0:]`` returns the WHOLE string — a "truncation" that + # grows the text by the marker length. + return truncation_marker.strip() truncated = text[:start_chars] + truncation_marker + text[-end_chars:] logger.info( @@ -491,6 +505,53 @@ class BaseAgent(ABC): ) return truncated + def _enforce_context_window(self, messages: List[Dict]) -> List[Dict]: + """Hard pre-send gate: never dispatch a payload that cannot fit. + + ``_validate_context_size`` only logs; an over-window payload used to + go straight to the provider, get rejected (context-length 400 / + capacity cap), take the fallback down with it, and still record its + full estimated prompt as usage. Called immediately before an LLM + dispatch: progressively middle-truncates the largest tool results + (the usual culprit) and raises when even that cannot fit — BEFORE + the usage decorators run, so a hopeless payload costs nothing. + """ + from application.core.model_utils import get_token_limit + from application.utils import num_tokens_from_string + + context_limit = get_token_limit( + self.model_id, user_id=self.model_user_id or self.user + ) + current_tokens = self._calculate_current_context_tokens(messages) + if current_tokens < context_limit: + return messages + + logger.warning( + f"Context ({current_tokens:,} tokens) exceeds the model's window " + f"({context_limit:,}). Shrinking tool results before dispatch." + ) + for per_message_cap in (8000, 2000, 500): + for message in messages: + content = message.get("content") + if ( + message.get("role") == "tool" + and isinstance(content, str) + and num_tokens_from_string(content) > per_message_cap + ): + message["content"] = self._truncate_text_middle( + content, per_message_cap + ) + current_tokens = self._calculate_current_context_tokens(messages) + if current_tokens < context_limit: + return messages + + raise ValueError( + f"Conversation context ({current_tokens:,} tokens) exceeds the " + f"model's context window ({context_limit:,} tokens) even after " + f"shrinking tool results. Start a new conversation or remove " + f"large attachments." + ) + # ---- Message building ---- def _build_messages( @@ -695,6 +756,9 @@ class BaseAgent(ABC): preserve_responses_state: bool = False, ): self._validate_context_size(messages) + # Hard gate: refuse/shrink instead of dispatching a payload the + # provider is guaranteed to reject (see _enforce_context_window). + messages = self._enforce_context_window(messages) if not preserve_responses_state: starter = getattr(self.llm, "start_responses_turn", None) diff --git a/application/api/answer/services/compression/service.py b/application/api/answer/services/compression/service.py index 3af3ac63..2934d79e 100644 --- a/application/api/answer/services/compression/service.py +++ b/application/api/answer/services/compression/service.py @@ -131,6 +131,18 @@ class CompressionService: original_tokens / compressed_tokens if compressed_tokens > 0 else 0 ) + # Port of the in-memory path's guard: a "successful" summary + # that isn't smaller than what it replaces must never become a + # compression point — it would make every later rebuild WORSE + # while reporting success (observed in prod: "successful" + # compressions with negative savings). + if compressed_tokens >= original_tokens: + raise ValueError( + f"Compression did not reduce token count " + f"({original_tokens} → {compressed_tokens}); " + f"keeping original history" + ) + logger.info( f"Compression complete: {original_tokens} → {compressed_tokens} tokens " f"({compression_ratio:.1f}x compression)" @@ -244,7 +256,7 @@ class CompressionService: f"(messages {last_compressed_index + 1}-{total_queries - 1})" ) - return compressed_summary, recent_queries + return compressed_summary, self._bound_recent_queries(recent_queries) except Exception as e: logger.error( @@ -255,6 +267,90 @@ class CompressionService: return None, [] return None, queries + def _truncate_middle_tokens(self, text: str, max_tokens: int) -> str: + """Middle-truncate ``text`` to roughly ``max_tokens`` tokens.""" + from application.utils import num_tokens_from_string + + current = num_tokens_from_string(text) + if current <= max_tokens: + return text + chars_per_token = len(text) / current if current > 0 else 4 + target_chars = int(max_tokens * chars_per_token * 0.95) + keep = int(target_chars * 0.4) + marker = "\n\n[... trimmed to fit context after compression ...]\n\n" + if keep <= 0: + # ``text[-0:]`` would return the WHOLE string, not nothing. + return marker.strip() + return text[:keep] + marker + text[-keep:] + + def _bound_recent_queries( + self, queries: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Cap oversized verbatim fields in the post-compression-point tail. + + Compression summarizes everything up to the compression point, but + the tail rides along verbatim — a single giant prompt / response / + tool result there can defeat the whole compression (measured in + prod: half of compressions produced no reduction in the next + call's prompt). Returns copies; the caller's conversation dict is + never mutated. + """ + max_tokens = int( + getattr(settings, "COMPRESSION_RECENT_FIELD_MAX_TOKENS", 8000) or 0 + ) + if max_tokens <= 0: + return queries + from application.utils import num_tokens_from_string + + bounded: List[Dict[str, Any]] = [] + trimmed = 0 + for query in queries: + if not isinstance(query, dict): + bounded.append(query) + continue + out = query + for field in ("prompt", "response"): + value = query.get(field) + if ( + isinstance(value, str) + and num_tokens_from_string(value) > max_tokens + ): + if out is query: + out = dict(query) + out[field] = self._truncate_middle_tokens(value, max_tokens) + trimmed += 1 + tool_calls = query.get("tool_calls") + if isinstance(tool_calls, list): + new_calls = None + for idx, tc in enumerate(tool_calls): + if not isinstance(tc, dict): + continue + result = tc.get("result") + if ( + isinstance(result, str) + and num_tokens_from_string(result) > max_tokens + ): + if new_calls is None: + new_calls = [ + dict(c) if isinstance(c, dict) else c + for c in tool_calls + ] + new_calls[idx]["result"] = self._truncate_middle_tokens( + result, max_tokens + ) + trimmed += 1 + if new_calls is not None: + if out is query: + out = dict(query) + out["tool_calls"] = new_calls + bounded.append(out) + if trimmed: + logger.info( + f"Bounded {trimmed} oversized field(s) in recent " + f"uncompressed queries (cap: {max_tokens} tokens each)" + ) + return bounded + def _extract_summary(self, llm_response: str) -> str: """ Extract clean summary from LLM response. diff --git a/application/core/settings.py b/application/core/settings.py index ce0230af..5836497a 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -272,6 +272,8 @@ class Settings(BaseSettings): COMPRESSION_MODEL_OVERRIDE: Optional[str] = None # Use different model for compression COMPRESSION_PROMPT_VERSION: str = "v1.0" # Track prompt iterations COMPRESSION_MAX_HISTORY_POINTS: int = 3 # Keep only last N compression points to prevent DB bloat + COMPRESSION_RECENT_FIELD_MAX_TOKENS: int = 8000 # Per-field cap on the verbatim tail kept after a compression point (0 disables) + TOOL_RESULT_MAX_TOKENS: int = 20000 # Cap on a single tool result entering the LLM context (0 disables); journal/DB keep the full result # Internal SSE push channel (notifications + durable replay journal) # Master switch — when False, /api/events emits a "push_disabled" comment diff --git a/application/llm/base.py b/application/llm/base.py index 79c8397b..7399320b 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -142,6 +142,44 @@ class BaseLLM(ABC): return args_dict return {k: v for k, v in args_dict.items() if v is not None} + @staticmethod + def _fallback_payload_fits(fallback, kwargs) -> bool: + """Whether the failed request's payload can fit the fallback model. + + A primary rejected for size (context-length 400, capacity cap) hands + the *same* oversized payload to the fallback, which then rejects it + too — one guaranteed-failed provider call plus one estimated-prompt + ``token_usage`` row for nothing. Skip the attempt when the estimated + prompt already exceeds the fallback's context window. Estimation + errors never block the attempt. + """ + messages = kwargs.get("messages") + if not messages: + return True + try: + from application.core.model_utils import get_token_limit + from application.usage import _count_prompt_tokens + + estimated = _count_prompt_tokens(messages, tools=kwargs.get("tools")) + limit = get_token_limit( + fallback.model_id, + user_id=getattr(fallback, "model_user_id", None), + ) + # 10% slack: the tiktoken estimate over-counts vs provider + # tokenizers (and ``get_token_limit`` returns a conservative + # default for unregistered models) — only skip when the payload + # is decisively over, never on a borderline estimate. + if estimated > int(limit * 1.1): + logger.warning( + f"Skipping fallback to {fallback.model_id}: estimated " + f"prompt (~{estimated} tokens) cannot fit its context " + f"window ({limit} tokens)." + ) + return False + except Exception: + logger.debug("Fallback payload size estimation failed", exc_info=True) + return True + def _execute_with_fallback( self, method_name: str, decorators: list, *args, **kwargs ): @@ -186,6 +224,8 @@ class BaseLLM(ABC): logger.error(f"Primary LLM failed and no fallback configured: {str(e)}") raise fallback = self.fallback_llm + if not self._fallback_payload_fits(fallback, kwargs): + raise self._responding_provider = fallback.provider_name logger.warning( f"Primary LLM failed. Falling back to " @@ -232,12 +272,26 @@ class BaseLLM(ABC): try: yield from decorated_method() except Exception as e: + if getattr(self, "_stream_reached_finish", False): + # The primary already delivered a finish signal — only + # trailing frames (usage chunk, [DONE]) failed. Restreaming + # from the fallback would duplicate the entire answer the + # user already received (and re-run tool calls). Re-raise; + # the streaming handler treats post-finish failures as + # non-fatal. + logger.warning( + f"Primary LLM failed after delivering its finish signal; " + f"not engaging fallback. Error: {str(e)}" + ) + raise if not self.fallback_llm: logger.error( f"Primary LLM failed and no fallback configured: {str(e)}" ) raise fallback = self.fallback_llm + if not self._fallback_payload_fits(fallback, kwargs): + raise self._responding_provider = fallback.provider_name logger.warning( f"Primary LLM failed mid-stream. Falling back to " diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index aa00177b..44cb558d 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -22,6 +22,44 @@ _FINALIZE_INSTRUCTION = ( ) +def _bound_tool_response_for_llm(tool_response: Any) -> Any: + """Cap a single tool result before it enters the LLM message array. + + One uncapped web page / API response has produced multi-hundred-KB tool + results that get re-sent in every subsequent round of the tool loop — + blowing conversations past every model's context window (observed in + prod up to a 14M-token estimated prompt). The durability journal and + the persisted conversation keep the FULL result (bounded separately at + persistence); only the copy handed to the model is truncated here. + """ + from application.core.settings import settings + from application.utils import num_tokens_from_string + + max_tokens = int(getattr(settings, "TOOL_RESULT_MAX_TOKENS", 20000) or 0) + if max_tokens <= 0: + return tool_response + text = tool_response if isinstance(tool_response, str) else str(tool_response) + tokens = num_tokens_from_string(text) + if tokens <= max_tokens: + return tool_response + chars_per_token = len(text) / tokens if tokens > 0 else 4 + target_chars = int(max_tokens * chars_per_token * 0.95) + keep = int(target_chars * 0.4) + marker = ( + f"\n\n[... tool result truncated: {tokens:,} tokens exceeded the " + f"{max_tokens:,}-token per-result limit ...]\n\n" + ) + if keep <= 0: + # ``text[-0:]`` would return the WHOLE string, not nothing. + return marker.strip() + logger.warning( + "Tool result truncated from %s to ~%s tokens before LLM handoff", + tokens, + max_tokens, + ) + return text[:keep] + marker + text[-keep:] + + @dataclass class ToolCall: """Represents a tool/function call from the LLM.""" @@ -697,13 +735,20 @@ class LLMHandler(ABC): logger.warning("Not enough queries to compress in-memory context") return False, None - metadata = compression_service.compress_conversation( - conversation, - compress_up_to_index=compress_up_to, - ) + try: + metadata = compression_service.compress_conversation( + conversation, + compress_up_to_index=compress_up_to, + ) + except ValueError: + # compress_conversation raises when the summary is not + # smaller than the original (negative-savings guard). For + # the in-memory path that is not fatal — fall back to + # minimal pruning and keep the tool loop running. + metadata = None # If compression doesn't reduce tokens, fall back to minimal pruning - if ( + if metadata is None or ( metadata.compressed_token_count >= metadata.original_token_count ): @@ -1014,6 +1059,9 @@ class LLMHandler(ABC): except StopIteration as e: tool_response, call_id = e.value break + # The journal / persisted conversation received the full + # result inside the executor; the model gets a bounded copy. + tool_response = _bound_tool_response_for_llm(tool_response) # Standard internal format: assistant message with tool_calls array args_str = ( @@ -1236,8 +1284,58 @@ class LLMHandler(ABC): buffer = "" tool_calls = {} reasoning_buffer = "" + finish_reason = None - for chunk in self._iterate_stream(response): + # Consume the provider stream to exhaustion before acting on its + # finish reason. Acting mid-iteration (as this loop used to) left + # every tool round's generator abandoned until request teardown, + # where all their usage-decorator ``finally`` blocks fired at once + # and mis-billed each round with the shared ``_last_usage`` of the + # final round (duplicate token_usage rows). Exhausting the stream + # also delivers the terminal usage-only chunk (Chat Completions + # ``stream_options.include_usage``), which arrives *after* the + # finish_reason chunk — so provider-exact counts land before the + # decorator persists this round's row. + stream_iter = self._iterate_stream(response) + while True: + try: + chunk = next(stream_iter) + except StopIteration: + break + except Exception: + # A failure in the trailing frames (usage chunk, [DONE]) + # after the round already finished must not fail — or + # fallback-restream — an answer the user has fully + # received. GeneratorExit/KeyboardInterrupt are + # BaseException and still propagate. + # + # The handler-local ``finish_reason`` only ever sees + # tool-call finishes (providers don't surface a parseable + # "stop" chunk — content arrives as bare strings), so the + # final answer round is covered by the LLM-level + # ``_stream_reached_finish`` flag instead — checked on the + # primary AND its fallback (whichever actually served the + # stream). + llm = getattr(agent, "llm", None) + provider_finished = bool( + getattr(llm, "_stream_reached_finish", False) + or getattr( + getattr(llm, "_fallback_llm", None), + "_stream_reached_finish", + False, + ) + ) + if finish_reason is not None or provider_finished: + logger.warning( + "Provider stream failed after finish " + "(finish_reason=%s, provider_finished=%s); " + "ignoring trailing-frame failure", + finish_reason, + provider_finished, + exc_info=True, + ) + break + raise if isinstance(chunk, dict) and chunk.get("type") == "thought": reasoning_buffer += chunk.get("thought") or "" yield chunk @@ -1251,6 +1349,13 @@ class LLMHandler(ABC): if parsed.tool_calls: for call in parsed.tool_calls: + if call.index is None: + # Providers like Google emit each parallel call as + # a COMPLETE, index-less ToolCall per chunk. They + # must never be merged into one another (dict + # arguments would even raise on ``+=``). + tool_calls[("complete", len(tool_calls))] = call + continue if call.index not in tool_calls: tool_calls[call.index] = call else: @@ -1262,88 +1367,104 @@ class LLMHandler(ABC): if call.arguments: if existing.arguments is None: existing.arguments = call.arguments - else: + elif isinstance(existing.arguments, str) and isinstance( + call.arguments, str + ): existing.arguments += call.arguments + else: + # Complete (non-delta) payloads: latest wins. + existing.arguments = call.arguments # Preserve thought_signature for Google Gemini 3 models if call.thought_signature: existing.thought_signature = call.thought_signature if parsed.finish_reason == "tool_calls": - tool_handler_gen = self.handle_tool_calls( - agent, - list(tool_calls.values()), - tools_dict, - messages, - reasoning_content=reasoning_buffer, - ) - while True: - try: - yield next(tool_handler_gen) - except StopIteration as e: - messages, pending_actions = e.value - break - tool_calls = {} - pause_reasoning = reasoning_buffer - reasoning_buffer = "" - - # If tools need approval or client execution, pause the loop - if pending_actions: - agent._pending_continuation = { - "messages": messages, - "pending_tool_calls": pending_actions, - "tools_dict": tools_dict, - "reasoning_content": pause_reasoning, - } - yield { - "type": "tool_calls_pending", - "data": {"pending_tool_calls": pending_actions}, - } - return - - next_iteration = _iteration + 1 - cap_reached = next_iteration >= MAX_TOOL_ITERATIONS - - # Check if context limit was reached during tool execution - if hasattr(agent, 'context_limit_reached') and agent.context_limit_reached: - # Add system message warning about context limit - messages.append({ - "role": "system", - "content": ( - "WARNING: Context window limit has been reached. " - "Please provide a final response to the user without making additional tool calls. " - "Summarize the work completed so far." - ) - }) - logger.info("Context limit reached - instructing agent to wrap up") - elif cap_reached: - logger.warning( - "agent tool loop hit cap (%d); forcing finalize", - MAX_TOOL_ITERATIONS, - ) - messages.append( - {"role": "system", "content": _FINALIZE_INSTRUCTION}, - ) - - # See note above on agent.model_id vs llm.model_id. - response = agent.llm.gen_stream( - model=getattr(agent.llm, "model_id", None) or agent.model_id, - messages=messages, - tools=( - None - if cap_reached - or getattr(agent, "context_limit_reached", False) - else agent.tools - ), - ) - self.llm_calls.append(build_stack_data(agent.llm)) - - yield from self.handle_streaming( - agent, response, tools_dict, messages, - _iteration=next_iteration, - ) - return + finish_reason = "tool_calls" + continue if parsed.content: buffer += parsed.content yield buffer buffer = "" - if parsed.finish_reason == "stop": - return + if parsed.finish_reason == "stop" and finish_reason is None: + finish_reason = "stop" + + if finish_reason != "tool_calls": + return + + tool_handler_gen = self.handle_tool_calls( + agent, + list(tool_calls.values()), + tools_dict, + messages, + reasoning_content=reasoning_buffer, + ) + while True: + try: + yield next(tool_handler_gen) + except StopIteration as e: + messages, pending_actions = e.value + break + pause_reasoning = reasoning_buffer + + # If tools need approval or client execution, pause the loop + if pending_actions: + agent._pending_continuation = { + "messages": messages, + "pending_tool_calls": pending_actions, + "tools_dict": tools_dict, + "reasoning_content": pause_reasoning, + } + yield { + "type": "tool_calls_pending", + "data": {"pending_tool_calls": pending_actions}, + } + return + + next_iteration = _iteration + 1 + cap_reached = next_iteration >= MAX_TOOL_ITERATIONS + + # Check if context limit was reached during tool execution + if hasattr(agent, 'context_limit_reached') and agent.context_limit_reached: + # Add system message warning about context limit + messages.append({ + "role": "system", + "content": ( + "WARNING: Context window limit has been reached. " + "Please provide a final response to the user without making additional tool calls. " + "Summarize the work completed so far." + ) + }) + logger.info("Context limit reached - instructing agent to wrap up") + elif cap_reached: + logger.warning( + "agent tool loop hit cap (%d); forcing finalize", + MAX_TOOL_ITERATIONS, + ) + messages.append( + {"role": "system", "content": _FINALIZE_INSTRUCTION}, + ) + + # Hard pre-send gate: tool results appended this round may have + # pushed the payload past the model's window — shrink or refuse + # BEFORE the usage decorators run (a rejected dispatch would still + # bill its full estimated prompt). + enforce = getattr(agent, "_enforce_context_window", None) + if callable(enforce): + messages = enforce(messages) + + # See note above on agent.model_id vs llm.model_id. + response = agent.llm.gen_stream( + model=getattr(agent.llm, "model_id", None) or agent.model_id, + messages=messages, + tools=( + None + if cap_reached + or getattr(agent, "context_limit_reached", False) + else agent.tools + ), + ) + self.llm_calls.append(build_stack_data(agent.llm)) + + yield from self.handle_streaming( + agent, response, tools_dict, messages, + _iteration=next_iteration, + ) diff --git a/application/llm/openai.py b/application/llm/openai.py index d2d2ec9a..166288a9 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -160,6 +160,14 @@ class OpenAILLM(BaseLLM): self._last_response_id = None self._last_reasoning_items = [] self._last_usage = None + # One-shot guard consumed by ``_prefer_provider_usage``: fresh + # provider usage flips it False, the first billing read flips it + # True, so no two token_usage rows share one reported usage. + self._last_usage_claimed = False + # True once the current stream delivered a finish signal — only + # trailing frames remain, so a failure there must not trigger a + # fallback restream of the already-delivered answer. + self._stream_reached_finish = False self._imported_response_id = None def responses_chain_key(self) -> str: @@ -443,6 +451,7 @@ class OpenAILLM(BaseLLM): if response_format: request_params["response_format"] = response_format self._last_usage = None + self._stream_reached_finish = False response = self.client.chat.completions.create(**request_params) logging.debug("OpenAI request completed") self._record_chat_usage(getattr(response, "usage", None)) @@ -515,6 +524,7 @@ class OpenAILLM(BaseLLM): stream_options.setdefault("include_usage", True) request_params["stream_options"] = stream_options self._last_usage = None + self._stream_reached_finish = False response = self.client.chat.completions.create(**request_params) try: @@ -537,6 +547,12 @@ class OpenAILLM(BaseLLM): has_tool_calls = bool(getattr(delta, "tool_calls", None)) finish_reason = getattr(choice, "finish_reason", None) + if finish_reason: + # The answer is complete; only trailing frames (usage + # chunk, [DONE]) remain. ``_stream_with_fallback`` reads + # this to refuse restreaming a delivered answer when a + # trailing frame fails. + self._stream_reached_finish = True # Yield non-content chunks only when needed for tool-call handling. if has_tool_calls or finish_reason == "tool_calls": @@ -596,29 +612,77 @@ class OpenAILLM(BaseLLM): parts.append(file_part) return parts - def _to_responses_input(self, messages): + def _to_responses_input(self, messages, chained=False): """Translate cleaned Chat-Completions messages into a Responses ``input`` item list. Reasoning items captured during the in-turn tool loop are re-injected ahead of the function calls they belong to (deduped by id) so the model keeps its chain-of-thought across the round-trip. + + Pairing invariant: the Responses API 400s on a ``function_call`` + without a matching ``function_call_output`` ("No tool output found + for function call …") and on an output without its call. History + rebuilds (compression, pause/resume) can produce either orphan, so + both directions are dropped here rather than sent to certain + rejection — along with any reasoning items that would then precede + a dropped call with no following item. A call is only "paired" when + its output appears LATER in the list (an output *before* its call + is malformed history; both sides are dropped). + + ``chained=True`` (store-mode ``previous_response_id`` requests) + disables the guard entirely: ``_trim_for_previous_response`` + deliberately sends bare ``function_call_output`` items whose calls + live server-side in the previous response — dropping those breaks + every tool round. """ input_items = [] emitted_reasoning = set() - for message in messages: + # First position (message index) at which each call_id's output + # appears — used to require call-before-output ordering. + output_positions: dict = {} + for position, m in enumerate(messages): + if m.get("role") == "tool" and m.get("tool_call_id") is not None: + output_positions.setdefault(m["tool_call_id"], position) + emitted_call_ids = set() + for position, message in enumerate(messages): role = message.get("role") message_reasoning = message.get("responses_reasoning_items") or [] - for item in message_reasoning: - item_id = item.get("id") if isinstance(item, dict) else None - if item_id and item_id in emitted_reasoning: - continue - if item_id: - emitted_reasoning.add(item_id) - input_items.append(item) tool_calls = message.get("tool_calls") if tool_calls and role == "assistant": - for tc in tool_calls: + kept_calls = [ + tc + for tc in tool_calls + if chained + or output_positions.get(tc.get("id", ""), -1) > position + ] + if len(kept_calls) < len(tool_calls): + kept_ids = {id(tc) for tc in kept_calls} + dropped = [ + tc.get("id", "") + for tc in tool_calls + if id(tc) not in kept_ids + ] + logging.warning( + "Dropping %d function_call item(s) without a matching " + "later function_call_output (call_ids=%s) from " + "Responses input", + len(dropped), + dropped, + ) + if not kept_calls: + # Nothing left to emit for this message; its reasoning + # items must not be emitted either (a trailing reasoning + # item with no following item is itself rejected). + continue + for item in message_reasoning: + item_id = item.get("id") if isinstance(item, dict) else None + if item_id and item_id in emitted_reasoning: + continue + if item_id: + emitted_reasoning.add(item_id) + input_items.append(item) + for tc in kept_calls: call_id = tc.get("id", "") for item in self._reasoning_for_calls.get(call_id, []): item_id = item.get("id") @@ -634,9 +698,24 @@ class OpenAILLM(BaseLLM): "name": func.get("name", ""), "arguments": func.get("arguments", "") or "{}", }) + emitted_call_ids.add(call_id) continue + for item in message_reasoning: + item_id = item.get("id") if isinstance(item, dict) else None + if item_id and item_id in emitted_reasoning: + continue + if item_id: + emitted_reasoning.add(item_id) + input_items.append(item) tool_call_id = message.get("tool_call_id") if role == "tool" and tool_call_id is not None: + if not chained and tool_call_id not in emitted_call_ids: + logging.warning( + "Dropping orphaned function_call_output (call_id=%s) " + "with no preceding function_call from Responses input", + tool_call_id, + ) + continue tool_content = message.get("content") input_items.append({ "type": "function_call_output", @@ -827,6 +906,7 @@ class OpenAILLM(BaseLLM): if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result + self._last_usage_claimed = False def _record_responses_metadata(self, response): rid = getattr(response, "id", None) @@ -850,6 +930,7 @@ class OpenAILLM(BaseLLM): if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result + self._last_usage_claimed = False @staticmethod def _responses_status_error(response) -> str | None: @@ -924,9 +1005,10 @@ class OpenAILLM(BaseLLM): self._last_response_id = None self._last_usage = None self._last_finish_reason = None - if previous_response_id and settings.OPENAI_RESPONSES_STORE: + 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) + input_items = self._to_responses_input(messages, chained=chained) params = self._build_responses_params( model, input_items, @@ -993,9 +1075,11 @@ class OpenAILLM(BaseLLM): self._last_response_id = None self._last_usage = None self._last_finish_reason = None - if previous_response_id and settings.OPENAI_RESPONSES_STORE: + 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) + input_items = self._to_responses_input(messages, chained=chained) params = self._build_responses_params( model, input_items, @@ -1069,6 +1153,7 @@ class OpenAILLM(BaseLLM): status_error = self._responses_status_error(completed_response) if status_error: raise RuntimeError(status_error) + self._stream_reached_finish = True self._record_responses_metadata(completed_response) self._last_reasoning_items = reasoning_items self._last_finish_reason = "tool_calls" if func_calls else "stop" diff --git a/application/usage.py b/application/usage.py index 5f7c55c5..11842ffc 100644 --- a/application/usage.py +++ b/application/usage.py @@ -155,10 +155,23 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in reported = getattr(llm, "_last_usage", None) if not isinstance(reported, dict): return call_usage + # ``_last_usage`` is shared instance state overwritten by every call on + # this LLM. Each reported usage may be billed to exactly ONE call: the + # provider clears ``_last_usage_claimed`` when it records fresh usage, + # and the first decorator ``finally`` to read it claims it. Without + # this, a generator finalized late (abandoned round, GC) would adopt a + # *different* call's provider counts. ``_last_usage`` itself is left in + # place for read-only consumers (client-facing usage metadata). + if getattr(llm, "_last_usage_claimed", False): + return call_usage prompt = reported.get("prompt_tokens") completion = reported.get("completion_tokens") if prompt is None or completion is None: return call_usage + try: + llm._last_usage_claimed = True + except AttributeError: + pass return { "prompt_tokens": int(prompt or 0), "generated_tokens": int(completion or 0), diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index d474b801..f25b028b 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -344,6 +344,7 @@ "stats": { "messages": "Messages", "tokens": "Tokens", + "tokensHint": "All LLM-billed tokens. Agents re-send the conversation context on every tool step, so tool-heavy chats use many times more tokens than the visible messages suggest.", "toolCalls": "Tool Calls", "runSuccess": "Run Success", "feedback": "Feedback" diff --git a/frontend/src/settings/Analytics.tsx b/frontend/src/settings/Analytics.tsx index 47b4e135..1183f2af 100644 --- a/frontend/src/settings/Analytics.tsx +++ b/frontend/src/settings/Analytics.tsx @@ -336,7 +336,7 @@ export default function Analytics({ agentId }: AnalyticsProps) { ? `${Math.round((scheduleTotals.completed / scheduleRunsTotal) * 100)}%` : '—'; - const statCards = [ + const statCards: { label: string; value: string; hint?: string }[] = [ { label: t('settings.analytics.stats.messages'), value: totalMessages.toLocaleString(), @@ -344,6 +344,7 @@ export default function Analytics({ agentId }: AnalyticsProps) { { label: t('settings.analytics.stats.tokens'), value: totalTokens.toLocaleString(), + hint: t('settings.analytics.stats.tokensHint'), }, { label: t('settings.analytics.stats.toolCalls'), @@ -410,7 +411,8 @@ export default function Analytics({ agentId }: AnalyticsProps) { {statCards.map((card) => (

{card.label}

diff --git a/tests/api/answer/services/compression/test_service.py b/tests/api/answer/services/compression/test_service.py index 99042780..6145a168 100644 --- a/tests/api/answer/services/compression/test_service.py +++ b/tests/api/answer/services/compression/test_service.py @@ -421,3 +421,128 @@ class TestLogToolCallStats: } ] svc._log_tool_call_stats(queries) + + +@pytest.mark.unit +class TestNegativeSavingsGuard: + """A summary that isn't smaller than the original must never be saved + as a compression point.""" + + def test_growing_compression_raises(self, mock_llm, sample_conversation): + mock_llm.gen.return_value = ( + "

" + ("padding words that inflate the summary " * 200) + "" + ) + service = CompressionService(llm=mock_llm, model_id="m") + with pytest.raises(ValueError, match="did not reduce"): + service.compress_conversation(sample_conversation, 2) + + def test_growing_compression_never_saved( + self, mock_llm, mock_conversation_service, sample_conversation + ): + mock_llm.gen.return_value = ( + "" + ("padding words that inflate the summary " * 200) + "" + ) + service = CompressionService( + llm=mock_llm, + model_id="m", + conversation_service=mock_conversation_service, + ) + with pytest.raises(ValueError, match="did not reduce"): + service.compress_and_save("conv-1", sample_conversation, 2) + mock_conversation_service.update_compression_metadata.assert_not_called() + + def test_shrinking_compression_still_succeeds(self, mock_llm, sample_conversation): + # Default fixture summary is tiny relative to the conversation. + service = CompressionService(llm=mock_llm, model_id="m") + metadata = service.compress_conversation(sample_conversation, 2) + assert metadata.compressed_token_count < metadata.original_token_count + + +@pytest.mark.unit +class TestBoundRecentQueries: + """Oversized verbatim fields in the post-compression tail are capped.""" + + def _compressed_conversation(self, recent_query): + return { + "queries": [ + {"prompt": "old", "response": "old answer"}, + recent_query, + ], + "compression_metadata": { + "is_compressed": True, + "compression_points": [ + { + "compressed_summary": "summary", + "query_index": 0, + "compressed_token_count": 10, + "original_token_count": 100, + } + ], + }, + } + + def test_oversized_response_is_trimmed_without_mutating_original( + self, mock_llm, monkeypatch + ): + monkeypatch.setattr( + "application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS", + 50, + raising=False, + ) + big = "word " * 2000 + recent = {"prompt": "q", "response": big} + conversation = self._compressed_conversation(recent) + service = CompressionService(llm=mock_llm, model_id="m") + + summary, queries = service.get_compressed_context(conversation) + + assert summary == "summary" + assert len(queries) == 1 + assert "trimmed to fit context" in queries[0]["response"] + assert len(queries[0]["response"]) < len(big) + # Original conversation dict untouched. + assert recent["response"] == big + + def test_oversized_tool_result_is_trimmed(self, mock_llm, monkeypatch): + monkeypatch.setattr( + "application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS", + 50, + raising=False, + ) + big = "data " * 2000 + recent = { + "prompt": "q", + "response": "short", + "tool_calls": [{"tool_name": "t", "result": big}], + } + conversation = self._compressed_conversation(recent) + service = CompressionService(llm=mock_llm, model_id="m") + + _summary, queries = service.get_compressed_context(conversation) + + assert "trimmed to fit context" in queries[0]["tool_calls"][0]["result"] + assert recent["tool_calls"][0]["result"] == big + + def test_small_fields_pass_through_unchanged(self, mock_llm): + recent = {"prompt": "q", "response": "short answer"} + conversation = self._compressed_conversation(recent) + service = CompressionService(llm=mock_llm, model_id="m") + + _summary, queries = service.get_compressed_context(conversation) + + assert queries[0] is recent + + def test_zero_cap_disables_bounding(self, mock_llm, monkeypatch): + monkeypatch.setattr( + "application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS", + 0, + raising=False, + ) + big = "word " * 2000 + recent = {"prompt": "q", "response": big} + conversation = self._compressed_conversation(recent) + service = CompressionService(llm=mock_llm, model_id="m") + + _summary, queries = service.get_compressed_context(conversation) + + assert queries[0] is recent diff --git a/tests/llm/handlers/test_context_gate.py b/tests/llm/handlers/test_context_gate.py new file mode 100644 index 00000000..14ecc453 --- /dev/null +++ b/tests/llm/handlers/test_context_gate.py @@ -0,0 +1,206 @@ +"""Pre-send context gate and per-tool-result cap. + +Covers the two guards that keep tool-loop payloads inside the model's +context window: ``_bound_tool_response_for_llm`` (one giant tool result +must not enter the message array uncapped) and +``BaseAgent._enforce_context_window`` (an over-window payload is shrunk +or refused *before* the usage decorators run). +""" + +from unittest.mock import Mock, patch + +import pytest + +from application.agents.base import BaseAgent +from application.llm.handlers.base import ( + LLMHandler, + LLMResponse, + ToolCall, + _bound_tool_response_for_llm, +) + + +class GateHandler(LLMHandler): + def parse_response(self, response): + kind = response.get("kind") + if kind == "tool_call": + return LLMResponse( + content="", + tool_calls=[ToolCall(id="1", name="t", arguments="{}", index=0)], + finish_reason="tool_calls", + raw_response=None, + ) + return LLMResponse( + content="", tool_calls=[], finish_reason="stop", raw_response=None + ) + + def create_tool_message(self, tool_call, result): + return { + "role": "tool", + "tool_call_id": tool_call.id, + "content": result if isinstance(result, str) else str(result), + } + + def _iterate_stream(self, response): + yield from response + + +class MockAgent(BaseAgent): + def _gen_inner(self, query, log_context=None): + yield {"answer": "test"} + + +@pytest.fixture +def agent(): + a = MockAgent( + endpoint="test", + llm_name="openai", + model_id="gpt-4o", + api_key="test-key", + ) + a.llm = Mock() + return a + + +@pytest.mark.unit +class TestBoundToolResponse: + def test_small_result_passes_through_unchanged(self): + result = {"status": "ok", "data": "small"} + assert _bound_tool_response_for_llm(result) is result + + def test_oversized_result_is_truncated(self, monkeypatch): + monkeypatch.setattr( + "application.core.settings.settings.TOOL_RESULT_MAX_TOKENS", + 30, + raising=False, + ) + big = "word " * 2000 + bounded = _bound_tool_response_for_llm(big) + assert "tool result truncated" in bounded + assert len(bounded) < len(big) + + def test_zero_cap_disables_truncation(self, monkeypatch): + monkeypatch.setattr( + "application.core.settings.settings.TOOL_RESULT_MAX_TOKENS", + 0, + raising=False, + ) + big = "word " * 2000 + assert _bound_tool_response_for_llm(big) is big + + def test_handle_tool_calls_bounds_the_llm_copy(self, monkeypatch): + """The message handed to the LLM is capped even though the executor + returned the full result (journal/persistence keep the original).""" + monkeypatch.setattr( + "application.core.settings.settings.TOOL_RESULT_MAX_TOKENS", + 30, + raising=False, + ) + handler = GateHandler() + big = "word " * 2000 + + def fake_executor(tools_dict, call): + yield {"type": "tool_call", "data": {}} + return big, "call_1" + + mock_agent = Mock() + mock_agent._check_context_limit = Mock(return_value=False) + mock_agent._execute_tool_action = Mock(side_effect=fake_executor) + mock_agent.tool_executor.check_pause = Mock(return_value=None) + + gen = handler.handle_tool_calls( + mock_agent, + [ToolCall(id="call_1", name="t", arguments="{}", index=0)], + {}, + [], + ) + while True: + try: + next(gen) + except StopIteration as e: + messages, pending = e.value + break + + tool_messages = [m for m in messages if m.get("role") == "tool"] + assert len(tool_messages) == 1 + assert "tool result truncated" in tool_messages[0]["content"] + assert len(tool_messages[0]["content"]) < len(big) + assert pending is None + + +@pytest.mark.unit +class TestEnforceContextWindow: + def test_within_window_returns_messages_untouched(self, agent): + messages = [{"role": "user", "content": "hello"}] + with patch( + "application.core.model_utils.get_token_limit", return_value=1000 + ): + assert agent._enforce_context_window(messages) is messages + + def test_over_window_shrinks_tool_messages(self, agent): + messages = [ + {"role": "user", "content": "question"}, + {"role": "tool", "tool_call_id": "1", "content": "word " * 2000}, + ] + with patch( + "application.core.model_utils.get_token_limit", return_value=1000 + ): + shrunk = agent._enforce_context_window(messages) + tool_content = shrunk[1]["content"] + assert "truncated to fit context limit" in tool_content + assert agent._calculate_current_context_tokens(shrunk) < 1000 + + def test_impossible_payload_raises_before_dispatch(self, agent): + # The bulk is NOT in tool messages, so shrinking cannot help. + messages = [{"role": "user", "content": "word " * 3000}] + with patch( + "application.core.model_utils.get_token_limit", return_value=100 + ): + with pytest.raises(ValueError, match="exceeds the model's context window"): + agent._enforce_context_window(messages) + + def test_streaming_loop_gates_before_next_round(self): + """handle_streaming must run the gate before dispatching the next + round's gen_stream.""" + handler = GateHandler() + order = [] + + mock_agent = Mock() + mock_agent.context_limit_reached = False + mock_agent.tools = [] + mock_agent.model_id = "m" + mock_agent._enforce_context_window = Mock( + side_effect=lambda msgs: (order.append("gate"), msgs)[1] + ) + mock_agent.llm.gen_stream = Mock( + side_effect=lambda **kw: (order.append("dispatch"), [{"kind": "stop"}])[1] + ) + + def fake_tool_calls(a, calls, tools_dict, messages, **kwargs): + yield {"type": "tool_call", "data": {}} + return messages, None + + with patch.object(handler, "handle_tool_calls", fake_tool_calls): + list( + handler.handle_streaming( + mock_agent, [{"kind": "tool_call"}], {}, [] + ) + ) + + assert order == ["gate", "dispatch"] + + +@pytest.mark.unit +class TestTinyCapTruncation: + def test_tiny_cap_never_returns_more_than_original(self, monkeypatch): + """keep==0 used to produce marker + FULL text (text[-0:] is the + whole string) — a 'truncation' that grows the payload.""" + monkeypatch.setattr( + "application.core.settings.settings.TOOL_RESULT_MAX_TOKENS", + 1, + raising=False, + ) + big = "word " * 500 + bounded = _bound_tool_response_for_llm(big) + assert len(bounded) < len(big) + assert "truncated" in bounded diff --git a/tests/llm/handlers/test_review_regressions.py b/tests/llm/handlers/test_review_regressions.py new file mode 100644 index 00000000..fceddae4 --- /dev/null +++ b/tests/llm/handlers/test_review_regressions.py @@ -0,0 +1,263 @@ +"""Regression tests from the adversarial review of the stream-drain change. + +Covers: Google-style index-less parallel tool calls (complete payloads must +never be merged), trailing-frame failures after finish_reason (must not fail +or fallback-restream a delivered answer), and the in-memory compression +path surviving the negative-savings ValueError via minimal pruning. +""" + +from unittest.mock import Mock, patch + +import pytest + +from application.llm.handlers.base import LLMHandler, LLMResponse, ToolCall + + +class ScriptedHandler(LLMHandler): + """Chunks are dicts carrying a prebuilt LLMResponse under 'resp'.""" + + def parse_response(self, response): + return response["resp"] + + def create_tool_message(self, tool_call, result): + return { + "role": "tool", + "tool_call_id": tool_call.id, + "content": str(result), + } + + def _iterate_stream(self, response): + yield from response + + +def _agent(llm=None): + agent = Mock() + agent.llm = llm or Mock() + agent.tools = [] + agent.model_id = "test-model" + agent.context_limit_reached = False + agent._check_context_limit = Mock(return_value=False) + agent._enforce_context_window = Mock(side_effect=lambda msgs: msgs) + # Explicit falsy values: auto-created Mock attributes are truthy and + # would trip the provider-finished check in the drain loop. + agent.llm._stream_reached_finish = False + agent.llm._fallback_llm = None + return agent + + +@pytest.mark.unit +class TestIndexlessParallelToolCalls: + def test_google_style_parallel_calls_stay_separate(self): + """Two complete index-less calls (dict arguments) must reach + handle_tool_calls as two distinct calls — the merge branch would + raise ``dict += dict`` or execute tool_B with tool_A's args.""" + handler = ScriptedHandler() + chunk_a = { + "resp": LLMResponse( + content="", + tool_calls=[ + ToolCall(id="a", name="tool_A", arguments={"x": 1}, index=None) + ], + finish_reason="tool_calls", + raw_response=None, + ) + } + chunk_b = { + "resp": LLMResponse( + content="", + tool_calls=[ + ToolCall(id="b", name="tool_B", arguments={"y": 2}, index=None) + ], + finish_reason="tool_calls", + raw_response=None, + ) + } + stop_round = [ + { + "resp": LLMResponse( + content="done", + tool_calls=[], + finish_reason="stop", + raw_response=None, + ) + } + ] + + received_calls = [] + + def fake_tool_calls(agent, calls, tools_dict, messages, **kwargs): + received_calls.extend(calls) + yield {"type": "tool_call", "data": {}} + return messages, None + + llm = Mock() + llm.gen_stream = Mock(return_value=stop_round) + agent = _agent(llm) + + with patch.object(handler, "handle_tool_calls", fake_tool_calls): + list(handler.handle_streaming(agent, [chunk_a, chunk_b], {}, [])) + + assert [(c.name, c.arguments) for c in received_calls] == [ + ("tool_A", {"x": 1}), + ("tool_B", {"y": 2}), + ] + + def test_indexed_string_deltas_still_merge(self): + """OpenAI-style indexed argument fragments keep concatenating.""" + handler = ScriptedHandler() + frag1 = { + "resp": LLMResponse( + content="", + tool_calls=[ToolCall(id="c1", name="t", arguments='{"q":', index=0)], + finish_reason=None, + raw_response=None, + ) + } + frag2 = { + "resp": LLMResponse( + content="", + tool_calls=[ToolCall(id="", name="", arguments='"x"}', index=0)], + finish_reason="tool_calls", + raw_response=None, + ) + } + stop_round = [ + { + "resp": LLMResponse( + content="done", + tool_calls=[], + finish_reason="stop", + raw_response=None, + ) + } + ] + + received_calls = [] + + def fake_tool_calls(agent, calls, tools_dict, messages, **kwargs): + received_calls.extend(calls) + yield {"type": "tool_call", "data": {}} + return messages, None + + llm = Mock() + llm.gen_stream = Mock(return_value=stop_round) + agent = _agent(llm) + + with patch.object(handler, "handle_tool_calls", fake_tool_calls): + list(handler.handle_streaming(agent, [frag1, frag2], {}, [])) + + assert len(received_calls) == 1 + assert received_calls[0].arguments == '{"q":"x"}' + + +class FailingTailStream: + """Yields chunks, then raises — models a reset in the trailing frames.""" + + def __init__(self, chunks, error): + self._chunks = chunks + self._error = error + + def __iter__(self): + yield from self._chunks + raise self._error + + +@pytest.mark.unit +class TestTrailingFrameFailures: + def _stop_chunk(self, text="hello"): + return { + "resp": LLMResponse( + content=text, + tool_calls=[], + finish_reason="stop", + raw_response=None, + ) + } + + def test_failure_after_finish_is_swallowed(self): + handler = ScriptedHandler() + stream = FailingTailStream( + [self._stop_chunk()], RuntimeError("reset in trailing frame") + ) + out = list(handler.handle_streaming(_agent(), stream, {}, [])) + assert "hello" in out + + def test_failure_before_finish_still_propagates(self): + handler = ScriptedHandler() + stream = FailingTailStream([], RuntimeError("early failure")) + with pytest.raises(RuntimeError, match="early failure"): + list(handler.handle_streaming(_agent(), stream, {}, [])) + + def test_final_round_failure_after_provider_finish_is_swallowed(self): + """The final answer round never surfaces a parseable 'stop' chunk + (content arrives as bare strings), so the swallow must key off the + LLM-level _stream_reached_finish flag.""" + handler = ScriptedHandler() + agent = _agent() + agent.llm._stream_reached_finish = True + # Bare string content, no parseable finish chunk, then a trailing + # frame failure — the OpenAI chat final-round shape. + stream = FailingTailStream( + ["the full answer"], RuntimeError("reset in trailing frame") + ) + out = list(handler.handle_streaming(agent, stream, {}, [])) + assert "the full answer" in out + + def test_fallback_served_stream_failure_after_finish_is_swallowed(self): + """When the fallback delivered the answer, its finish flag lives on + the fallback instance — the swallow must check it there too.""" + handler = ScriptedHandler() + agent = _agent() + agent.llm._stream_reached_finish = False + agent.llm._fallback_llm = Mock() + agent.llm._fallback_llm._stream_reached_finish = True + stream = FailingTailStream( + ["fallback answer"], RuntimeError("reset in trailing frame") + ) + out = list(handler.handle_streaming(agent, stream, {}, [])) + assert "fallback answer" in out + + +@pytest.mark.unit +class TestInMemoryCompressionNegativeSavings: + def test_value_error_falls_back_to_minimal_pruning(self): + """The negative-savings ValueError from compress_conversation must + route to _prune_messages_minimal, not kill the tool loop.""" + handler = ScriptedHandler() + agent = _agent() + agent.model_id = "gpt-4o" + agent.decoded_token = {"sub": "u1"} + messages = [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + ] + pruned = [{"role": "system", "content": "pruned"}] + + with patch.object( + handler, + "_build_conversation_from_messages", + return_value={"queries": [{"prompt": "q1", "response": "a1"}]}, + ), patch( + "application.core.model_utils.get_provider_from_model_id", + return_value="openai", + ), patch( + "application.core.model_utils.get_api_key_for_provider", + return_value="k", + ), patch( + "application.llm.llm_creator.LLMCreator.create_llm", + return_value=Mock(), + ), patch( + "application.api.answer.services.compression.service." + "CompressionService.compress_conversation", + side_effect=ValueError( + "Compression did not reduce token count (10 → 20); " + "keeping original history" + ), + ), patch.object( + handler, "_prune_messages_minimal", return_value=pruned + ): + ok, rebuilt = handler._perform_in_memory_compression(agent, messages) + + assert ok is True + assert rebuilt == pruned + assert agent.context_limit_reached is False diff --git a/tests/llm/handlers/test_round_usage_persistence.py b/tests/llm/handlers/test_round_usage_persistence.py new file mode 100644 index 00000000..2ada0b9e --- /dev/null +++ b/tests/llm/handlers/test_round_usage_persistence.py @@ -0,0 +1,220 @@ +"""Per-round token-usage persistence in the streaming tool loop. + +Regression tests for the duplicate ``token_usage`` rows bug: each tool +round's provider stream must be consumed to exhaustion *before* the next +round starts, so the ``stream_token_usage`` ``finally`` persists one row +per round with that round's own counts — never a late flush at request +teardown that adopts the final round's shared ``_last_usage``. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from application import usage as usage_mod +from application.llm.handlers.base import LLMHandler, LLMResponse, ToolCall + + +class RecordingStream: + """Iterable stream that records whether it was fully consumed.""" + + def __init__(self, chunks): + self._chunks = list(chunks) + self.exhausted = False + + def __iter__(self): + for chunk in self._chunks: + yield chunk + self.exhausted = True + + +class RoundHandler(LLMHandler): + """Minimal concrete handler: chunk dicts drive the parse result.""" + + def parse_response(self, response): + kind = response.get("kind") + if kind == "tool_call": + return LLMResponse( + content="", + tool_calls=[ToolCall(id="1", name="t", arguments="{}", index=0)], + finish_reason="tool_calls", + raw_response=None, + ) + if kind == "content": + return LLMResponse( + content="hello", + tool_calls=[], + finish_reason=None, + raw_response=None, + ) + if kind == "stop": + return LLMResponse( + content="", + tool_calls=[], + finish_reason="stop", + raw_response=None, + ) + # Terminal usage-only chunk: no content, no finish_reason. + return LLMResponse( + content="", tool_calls=[], finish_reason=None, raw_response=None + ) + + def create_tool_message(self, tool_call, result): + return { + "role": "tool", + "tool_call_id": tool_call.id, + "content": str(result), + } + + def _iterate_stream(self, response): + yield from response + + +def _agent(llm): + agent = Mock() + agent.llm = llm + agent.tools = [] + agent.model_id = "test-model" + agent.context_limit_reached = False + agent._check_context_limit = Mock(return_value=False) + agent._enforce_context_window = Mock(side_effect=lambda msgs: msgs) + return agent + + +class TestStreamDrainedBeforeNextRound: + def test_round_stream_exhausted_before_tool_handling(self): + """The round's generator must be fully consumed (usage chunk and + all) before handle_tool_calls runs — not abandoned mid-iteration.""" + handler = RoundHandler() + round1 = RecordingStream([{"kind": "tool_call"}, {"kind": "usage"}]) + round2 = RecordingStream([{"kind": "stop"}]) + + exhausted_at_tool_time = [] + + def fake_tool_calls(agent, calls, tools_dict, messages, **kwargs): + exhausted_at_tool_time.append(round1.exhausted) + yield {"type": "tool_call", "data": {}} + return messages, None + + llm = Mock() + llm.gen_stream = Mock(return_value=round2) + agent = _agent(llm) + + with patch.object(handler, "handle_tool_calls", fake_tool_calls): + list(handler.handle_streaming(agent, round1, {}, [])) + + assert exhausted_at_tool_time == [True] + assert round2.exhausted is True + + def test_stop_round_drains_trailing_usage_chunk(self): + """A round that finishes with ``stop`` still consumes the trailing + usage-only chunk (Chat Completions include_usage arrives after the + finish_reason chunk).""" + handler = RoundHandler() + stream = RecordingStream( + [{"kind": "content"}, {"kind": "stop"}, {"kind": "usage"}] + ) + agent = _agent(Mock()) + + out = list(handler.handle_streaming(agent, stream, {}, [])) + + assert stream.exhausted is True + assert "hello" in out + + +class TestPerRoundUsageRows: + def test_one_row_per_round_with_own_counts(self): + """Two rounds → two persisted rows, each with that round's own + provider counts (not two copies of the final round's).""" + handler = RoundHandler() + + class FakeLLM: + def __init__(self): + self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0} + self.decoded_token = {"sub": "u1"} + self.user_api_key = None + self.agent_id = None + self._last_usage = None + self._last_usage_claimed = False + self.round = 0 + + def gen_stream(self, model, messages, tools=None, **kwargs): + wrapped = usage_mod.stream_token_usage(FakeLLM._raw_gen_stream) + return wrapped(self, model, messages, True, tools, **kwargs) + + def _raw_gen_stream(self, model, messages, stream, tools, **kwargs): + self._last_usage = None + self._last_usage_claimed = False + self.round += 1 + current = self.round + if current == 1: + yield {"kind": "tool_call"} + else: + yield {"kind": "content"} + yield {"kind": "stop"} + # Provider reports usage in the terminal chunk, after the + # finish_reason — only a drained stream ever reaches this. + self._last_usage = { + "prompt_tokens": 100 * current, + "completion_tokens": 10 * current, + } + self._last_usage_claimed = False + + llm = FakeLLM() + agent = _agent(llm) + + def fake_tool_calls(a, calls, tools_dict, messages, **kwargs): + yield {"type": "tool_call", "data": {}} + return messages, None + + rows = [] + with patch.object( + usage_mod, + "_persist_call_usage", + side_effect=lambda llm_, cu: rows.append(dict(cu)), + ): + with patch.object(handler, "handle_tool_calls", fake_tool_calls): + first = llm.gen_stream( + model="m", messages=[{"role": "user", "content": "hi"}] + ) + list(handler.handle_streaming(agent, first, {}, [])) + + assert len(rows) == 2 + assert rows[0] == {"prompt_tokens": 100, "generated_tokens": 10} + assert rows[1] == {"prompt_tokens": 200, "generated_tokens": 20} + + +class TestPreferProviderUsageClaim: + def test_usage_claimed_only_once(self): + llm = SimpleNamespace( + _last_usage={"prompt_tokens": 7, "completion_tokens": 9}, + _last_usage_claimed=False, + ) + estimate = {"prompt_tokens": 42, "generated_tokens": 1} + + first = usage_mod._prefer_provider_usage(llm, dict(estimate)) + assert first == {"prompt_tokens": 7, "generated_tokens": 9} + assert llm._last_usage_claimed is True + + second = usage_mod._prefer_provider_usage(llm, dict(estimate)) + assert second == dict(estimate) + + def test_fresh_usage_resets_claim(self): + llm = SimpleNamespace( + _last_usage={"prompt_tokens": 7, "completion_tokens": 9}, + _last_usage_claimed=False, + ) + usage_mod._prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + + # Provider records a new call's usage → unclaimed again. + llm._last_usage = {"prompt_tokens": 70, "completion_tokens": 90} + llm._last_usage_claimed = False + + result = usage_mod._prefer_provider_usage( + llm, {"prompt_tokens": 2, "generated_tokens": 2} + ) + assert result == {"prompt_tokens": 70, "generated_tokens": 90} + + def test_missing_usage_keeps_estimate(self): + llm = SimpleNamespace(_last_usage=None, _last_usage_claimed=False) + estimate = {"prompt_tokens": 5, "generated_tokens": 3} + assert usage_mod._prefer_provider_usage(llm, dict(estimate)) == estimate diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index f42cb040..d5e76652 100644 --- a/tests/llm/test_fallback.py +++ b/tests/llm/test_fallback.py @@ -811,3 +811,124 @@ class TestRespondingProviderTracking: primary = _Google(fail_at=0, backup_models=["backup-model"]) primary.gen(**CALL_ARGS) assert primary._responding_provider == "openai" + + +# Tests — fallback payload size gate + + +@pytest.mark.integration +class TestFallbackPayloadSizeGate: + """A payload that cannot fit the fallback's context window must skip the + fallback attempt (it would be a guaranteed second rejection) and + propagate the primary's error instead.""" + + def _primary_with_backup(self, patch_model_utils, backup): + patch_model_utils( + get_provider=lambda mid, **_kw: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + return FakeLLM(fail_at=0, backup_models=["backup-model"]) + + BIG_ARGS = dict( + model="test-model", + messages=[{"role": "user", "content": "word " * 300}], + ) + + def test_gen_skips_fallback_when_payload_cannot_fit( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(responses=["backup ok"]) + primary = self._primary_with_backup(patch_model_utils, backup) + monkeypatch.setattr( + "application.core.model_utils.get_token_limit", + lambda mid, user_id=None: 10, + ) + with pytest.raises(RuntimeError, match="primary model unavailable"): + primary.gen(**self.BIG_ARGS) + assert backup.gen_called is False + + def test_stream_skips_fallback_when_payload_cannot_fit( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(stream_chunks=["backup chunk"]) + primary = self._primary_with_backup(patch_model_utils, backup) + monkeypatch.setattr( + "application.core.model_utils.get_token_limit", + lambda mid, user_id=None: 10, + ) + with pytest.raises(RuntimeError, match="mid-stream failure"): + list(primary.gen_stream(**self.BIG_ARGS)) + assert backup.gen_stream_called is False + + def test_fallback_proceeds_when_payload_fits( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(responses=["backup ok"]) + primary = self._primary_with_backup(patch_model_utils, backup) + monkeypatch.setattr( + "application.core.model_utils.get_token_limit", + lambda mid, user_id=None: 100000, + ) + assert primary.gen(**self.BIG_ARGS) == "backup ok" + assert backup.gen_called is True + + def test_estimation_failure_never_blocks_fallback( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(responses=["backup ok"]) + primary = self._primary_with_backup(patch_model_utils, backup) + + def boom(*a, **kw): + raise ValueError("estimator broken") + + monkeypatch.setattr("application.usage._count_prompt_tokens", boom) + assert primary.gen(**self.BIG_ARGS) == "backup ok" + assert backup.gen_called is True + + +# Tests — no fallback restream after the primary delivered its finish signal + + +@pytest.mark.integration +class TestNoRestreamAfterFinish: + """A trailing-frame failure (between the finish chunk and stream end) + must NOT restream the already-delivered answer from the fallback.""" + + class FinishThenFailLLM(FakeLLM): + def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs): + self.gen_stream_called = True + yield "the full answer" + # Provider marked the stream finished (finish_reason arrived)... + self._stream_reached_finish = True + # ...then the trailing usage/[DONE] frame dies. + raise RuntimeError("connection reset in trailing frame") + + def test_trailing_frame_failure_skips_fallback(self, patch_model_utils): + backup = FakeLLM(stream_chunks=["fallback chunk"]) + patch_model_utils( + get_provider=lambda mid, **_kw: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + primary = self.FinishThenFailLLM(backup_models=["backup-model"]) + + received = [] + with pytest.raises(RuntimeError, match="trailing frame"): + for chunk in primary.gen_stream(**CALL_ARGS): + received.append(chunk) + + assert received == ["the full answer"] + assert backup.gen_stream_called is False + + def test_pre_finish_failure_still_falls_back(self, patch_model_utils): + backup = FakeLLM(stream_chunks=["fallback chunk"]) + patch_model_utils( + get_provider=lambda mid, **_kw: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + primary = FakeLLM(fail_at=0, backup_models=["backup-model"]) + out = list(primary.gen_stream(**CALL_ARGS)) + assert out == ["fallback chunk"] + assert backup.gen_stream_called is True diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 1a9e52d2..1fa9aa78 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -126,15 +126,20 @@ def test_to_responses_input_reinjects_reasoning(monkeypatch): "encrypted_content": "enc", "summary": [], } llm._reasoning_for_calls = {"call_1": [reasoning_item]} - messages = [{ - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": "call_1", - "type": "function", - "function": {"name": "t", "arguments": "{}"}, - }], - }] + # The call must be paired with its output — unpaired calls are dropped + # (the API rejects them with "No tool output found for function call"). + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "t", "arguments": "{}"}, + }], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ok"}, + ] items = llm._to_responses_input(messages) # Reasoning item is emitted immediately before its function call. assert items[0] == reasoning_item @@ -254,15 +259,19 @@ def test_export_import_replays_encrypted_reasoning_on_fresh_instance(monkeypatch second = _make_llm(monkeypatch, _responses_caps()) assert second.import_responses_state(first.export_responses_state()) is True - items = second._to_responses_input([{ - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": "call_1", - "type": "function", - "function": {"name": "run", "arguments": "{}"}, - }], - }]) + # Paired with its output: unpaired calls are dropped by the builder. + items = second._to_responses_input([ + { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "run", "arguments": "{}"}, + }], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ok"}, + ]) assert items[0] == reasoning assert items[1]["type"] == "function_call" @@ -982,3 +991,117 @@ def test_builtin_default_models_stay_chat_completions(): catalogs = load_model_yamls([BUILTIN_MODELS_DIR]) models = {m.id: m for c in catalogs for m in c.models} assert models["gpt-5.4-mini"].capabilities.api_flavor == "chat_completions" + + +@pytest.mark.unit +def test_to_responses_input_drops_unpaired_function_call(monkeypatch): + """A function_call with no matching output must not reach the provider + (it 400s with "No tool output found for function call ...").""" + llm = _make_llm(monkeypatch, _responses_caps()) + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_orphan", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + }, + { + "id": "call_ok", + "type": "function", + "function": {"name": "read", "arguments": "{}"}, + }, + ], + }, + {"role": "tool", "tool_call_id": "call_ok", "content": "result"}, + ] + items = llm._to_responses_input(messages) + calls = [i for i in items if i.get("type") == "function_call"] + outputs = [i for i in items if i.get("type") == "function_call_output"] + assert [c["call_id"] for c in calls] == ["call_ok"] + assert [o["call_id"] for o in outputs] == ["call_ok"] + + +@pytest.mark.unit +def test_to_responses_input_drops_orphaned_output(monkeypatch): + """A function_call_output whose call was never emitted is dropped.""" + llm = _make_llm(monkeypatch, _responses_caps()) + messages = [ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "call_ghost", "content": "result"}, + {"role": "assistant", "content": "final"}, + ] + items = llm._to_responses_input(messages) + assert not [i for i in items if i.get("type") == "function_call_output"] + assert items[-1]["role"] == "assistant" + + +@pytest.mark.unit +def test_to_responses_input_skips_reasoning_of_dropped_calls(monkeypatch): + """When every call on an assistant message is unpaired, its reasoning + items are suppressed too — a trailing reasoning item with no following + item is itself rejected by the Responses API.""" + llm = _make_llm(monkeypatch, _responses_caps()) + llm._reasoning_for_calls = { + "call_orphan": [{"type": "reasoning", "id": "rs_1", "summary": []}] + } + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "responses_reasoning_items": [ + {"type": "reasoning", "id": "rs_msg", "summary": []} + ], + "tool_calls": [{ + "id": "call_orphan", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + }], + }, + ] + items = llm._to_responses_input(messages) + assert not [i for i in items if i.get("type") == "reasoning"] + assert not [i for i in items if i.get("type") == "function_call"] + + +@pytest.mark.unit +def test_to_responses_input_chained_keeps_bare_tool_output(monkeypatch): + """Store-mode chaining (previous_response_id) deliberately sends a bare + function_call_output whose call lives server-side — the pairing guard + must not drop it (that would 400 every chained tool round).""" + llm = _make_llm(monkeypatch, _responses_caps()) + messages = [ + {"role": "system", "content": "sys"}, + {"role": "tool", "tool_call_id": "call_server_side", "content": "result"}, + ] + items = llm._to_responses_input(messages, chained=True) + outputs = [i for i in items if i.get("type") == "function_call_output"] + assert [o["call_id"] for o in outputs] == ["call_server_side"] + + +@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 + are dropped (keeping the call alone would produce the unpaired-call + 400 the guard exists to prevent).""" + llm = _make_llm(monkeypatch, _responses_caps()) + messages = [ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "call_x", "content": "early result"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_x", + "type": "function", + "function": {"name": "t", "arguments": "{}"}, + }], + }, + ] + items = llm._to_responses_input(messages) + assert not [i for i in items if i.get("type") == "function_call"] + assert not [i for i in items if i.get("type") == "function_call_output"] diff --git a/tests/test_compression_service.py b/tests/test_compression_service.py index 0eaeed5f..22527665 100644 --- a/tests/test_compression_service.py +++ b/tests/test_compression_service.py @@ -210,6 +210,17 @@ class TestCompressionService: """ compression_service.llm.gen.return_value = mock_summary + # Real compressions run near a context-window threshold, so the + # original is always far larger than the summary. Pad the fixture + # accordingly — a summary bigger than the original is rejected by + # the negative-savings guard. + sample_conversation["queries"][0]["response"] = ( + "Python is a high-level programming language. " * 40 + ) + sample_conversation["queries"][1]["response"] = ( + "You can install Python from python.org. " * 40 + ) + # Compress first 2 queries result = compression_service.compress_conversation( conversation=sample_conversation, compress_up_to_index=1