fix(llm): bound complete serialized inputs before provider calls

src/skillspector/llm_analyzer_base.py

@@ -40,10 +40,12 @@
 from random import uniform
 from typing import Any, Literal, cast, get_args
 
+import tiktoken
 from langchain_anthropic import ChatAnthropic
 from langchain_core.messages import BaseMessage
 from langchain_openai.chat_models.base import BaseChatOpenAI
 from pydantic import BaseModel, Field, ValidationError, field_validator
+from tiktoken import registry as tiktoken_registry
 
 from skillspector.inference_usage import (
     InferenceUsageRecord,
@@ -177,6 +179,76 @@ def _require_structured_response(response: object) -> None:
         )
 
 
+REQUEST_FRAMING_TOKENS = 256
+
+
+def _input_token_counter_key(chat_model: object) -> tuple[object, ...]:
+    """Retain the selected counter for equivalent client configurations."""
+    if not isinstance(chat_model, BaseChatOpenAI):
+        return ("conservative_utf8_bytes",)
+    custom = getattr(chat_model, "custom_get_token_ids", None)
+    if callable(custom):
+        # The selected counter retains this callable, so its identity cannot be
+        # recycled while this key is cached. Do not retain whole HTTP clients.
+        return ("configured_chat_model_tokenizer", id(custom))
+    label = getattr(chat_model, "tiktoken_model_name", None) or getattr(
+        chat_model, "model_name", None
+    )
+    return ("openai_codec", label if isinstance(label, str) else None)
+
+
+def _input_token_counter(chat_model: object) -> tuple[Callable[[str], int], str]:
+    """Use a configured/recognized OpenAI codec; unknown codecs use UTF-8 bytes.
+
+    A model label must not silently inherit another model's tokenizer. Byte
+    length is a conservative bound for text when no verified codec is known.
+    """
+
+    def byte_bound(text: str) -> int:
+        return len(text.encode("utf-8", errors="replace"))
+
+    if not isinstance(chat_model, BaseChatOpenAI):
+        return byte_bound, "conservative_utf8_bytes"
+    custom = getattr(chat_model, "custom_get_token_ids", None)
+    if callable(custom):
+
+        def configured(text: str) -> int:
+            try:
+                count = len(custom(text))
+                return count if count > 0 or not text else byte_bound(text)
+            except Exception:
+                return byte_bound(text)
+
+        return configured, "configured_chat_model_tokenizer"
+    label = getattr(chat_model, "tiktoken_model_name", None) or getattr(
+        chat_model, "model_name", None
+    )
+    if not isinstance(label, str):
+        return byte_bound, "conservative_utf8_bytes"
+    try:
+        # Name lookup is a pure table lookup. Never initialize a codec here:
+        # get_encoding/encoding_for_model can download its vocabulary without a
+        # timeout, outside the shared scan deadline. Use only an existing codec.
+        encoding_name = tiktoken.encoding_name_for_model(label)
+        encoding = tiktoken_registry.ENCODINGS.get(encoding_name)
+    except Exception:
+        return byte_bound, "conservative_utf8_bytes"
+    if encoding is None:
+        return byte_bound, "conservative_utf8_bytes"
+
+    def recognized(text: str) -> int:
+        try:
+            return len(encoding.encode(text, disallowed_special=()))
+        except Exception:
+            return byte_bound(text)
+
+    return recognized, "recognized_tiktoken_codec"
+
+
+class LLMInputBudgetError(RuntimeError):
+    """Retain an unsubmitted batch that cannot fit the serialized input budget."""
+
+
 class LLMRuntimeLimitError(RuntimeError):
     """Signal that no shared scan time remains for an LLM operation."""
 
@@ -1074,17 +1146,22 @@ def chunk_file_by_lines(
     overlap_lines: int = CHUNK_OVERLAP_LINES,
     *,
     max_bytes: int | None = None,
+    fits_chunk: Callable[[str, int, int], bool] | None = None,
 ) -> list[tuple[str, int, int]]:
     """Split *content* into line-range chunks that each fit within *max_tokens*.
 
     Returns a list of ``(chunk_text, start_line, end_line)`` tuples where lines
     are 1-indexed.  Consecutive chunks share *overlap_lines* lines of context so
     findings near chunk boundaries still have surrounding code.
 
+    The token estimate includes numbered line prefixes and rounds cumulative
+    serialized size, rather than rounding each line independently.
+
     When *max_bytes* is set, each chunk must also fit within that many UTF-8
     bytes *after* :func:`number_lines` has prefixed it, which is what a CLI
     provider's input cap measures.  A single line that alone exceeds the budget
-    still becomes its own chunk: content is never dropped or truncated.
+    still becomes its own chunk: content is never dropped or truncated. The
+    invocation guard records an explicit failure instead of submitting it.
     """
     lines = content.splitlines(keepends=True)
     if not lines:
@@ -1098,25 +1175,37 @@ def chunk_file_by_lines(
     start_idx = 0
 
     while start_idx < len(lines):
-        token_count = 0
         byte_count = 0
         end_idx = start_idx
 
         while end_idx < len(lines):
-            line_tokens = estimate_tokens(lines[end_idx])
-            line_bytes = (
-                len(lines[end_idx].encode("utf-8", errors="replace")) + prefix_bytes
-                if max_bytes is not None
-                else 0
-            )
-            over_tokens = token_count + line_tokens > max_tokens
+            line_bytes = len(lines[end_idx].encode("utf-8", errors="replace")) + prefix_bytes
+            # Round the cumulative serialized size, not each short line. Prefixes
+            # and UTF-8 bytes conservatively bound the chars/4 token estimate.
+            over_tokens = (
+                byte_count + line_bytes + CHARS_PER_TOKEN - 1
+            ) // CHARS_PER_TOKEN > max_tokens
             over_bytes = max_bytes is not None and byte_count + line_bytes > max_bytes
             if (over_tokens or over_bytes) and end_idx > start_idx:
                 break
-            token_count += line_tokens
             byte_count += line_bytes
             end_idx += 1
 
+        if fits_chunk is not None and not fits_chunk(
+            "".join(lines[start_idx:end_idx]), start_idx + 1, end_idx
+        ):
+            low, high = start_idx + 1, end_idx
+            best = None
+            while low <= high:
+                middle = (low + high) // 2
+                if fits_chunk("".join(lines[start_idx:middle]), start_idx + 1, middle):
+                    best = middle
+                    low = middle + 1
+                else:
+                    high = middle - 1
+            # An unfit single line stays represented; the call guard refuses it.
+            end_idx = best if best is not None else start_idx + 1
+
         chunk_text = "".join(lines[start_idx:end_idx])
         chunks.append((chunk_text, start_idx + 1, end_idx))  # 1-indexed
 
@@ -1241,6 +1330,8 @@ def __init__(
             self._require_time_remaining()
             raise
         self._require_time_remaining()
+        self._token_counter, self._token_counter_source = _input_token_counter(self._llm)
+        self._call_token_counters = {_input_token_counter_key(self._llm): self._token_counter}
         # Native SDK retries cannot re-read a workflow-wide deadline between
         # attempts.  A dynamic deadline therefore uses our explicit retry loop,
         # which checks and caps every retry/backoff against remaining time.
@@ -1360,13 +1451,13 @@ def get_batches(
     ) -> list[Batch]:
         """Create one :class:`Batch` per file, splitting oversized files into chunks.
 
-        For CLI providers every batch must also keep the built prompt under the
+        Every batch accounts for the real prompt wrapper, findings and schema
+        as well as line prefixes. For CLI providers it must also stay under the
         provider's byte cap (see :func:`_cli_max_input_bytes`).  The fixed part of
         the prompt is measured with :meth:`build_prompt` itself, called with
         *prompt_kwargs* (the same keyword arguments later passed to the run), so
         subclass prompt layouts are accounted for rather than estimated.
         """
-        base_overhead = estimate_tokens(self.base_prompt)
         max_input_bytes = _cli_max_input_bytes()
 
         findings_by_file: dict[str, list[Finding]] = defaultdict(list)
@@ -1376,27 +1467,46 @@ def get_batches(
 
         batches: list[Batch] = []
         for path in file_paths:
-            content = file_cache.get(path) or "No content available for this file."
+            content = file_cache.get(path)
+            if content is None:
+                content = "No content available for this file."
             file_findings = findings_by_file.get(path, [])
 
-            extra = self._estimate_extra_overhead(file_findings)
-            content_budget = max(self._input_budget - base_overhead - extra, 1024)
-
             last_line = max(len(content.splitlines()), 1)
-            fixed_bytes = 0
+            # Only the irreducible wrapper belongs in the coarse budget. The
+            # range predicate measures findings that actually occur in a chunk.
+            fixed_bytes = self._fixed_prompt_bytes(path, [], last_line, **prompt_kwargs)
+            minimum = Batch(file_path=path, content="", findings=[])
+            if not self._prompt_fits_budget(self.build_prompt(minimum, **prompt_kwargs)):
+                # No line split can shrink a fixed wrapper. Keep all content as
+                # one explicit unsubmitted failure instead of refining every line.
+                retained = Batch(file_path=path, content=content, findings=file_findings)
+                batches.append(retained)
+                logger.warning("Fixed LLM prompt cannot fit input budget; batch content retained")
+                if max_input_bytes is not None and fixed_bytes > max_input_bytes:
+                    logger.warning(
+                        "LLM prompt exceeds the provider input limit; batch content retained"
+                    )
+                continue
+            fixed_tokens = (fixed_bytes + CHARS_PER_TOKEN - 1) // CHARS_PER_TOKEN
+            content_budget = max(self._input_budget - fixed_tokens, 1)
             byte_budget: int | None = None
             if max_input_bytes is not None:
-                fixed_bytes = self._fixed_prompt_bytes(
-                    path, file_findings, last_line, **prompt_kwargs
-                )
                 byte_budget = max_input_bytes - fixed_bytes - PROMPT_BYTE_SAFETY_MARGIN
 
             first_batch = len(batches)
-            content_tokens = estimate_tokens(content)
+            content_tokens = (
+                _numbered_size_bound(content, last_line) + CHARS_PER_TOKEN - 1
+            ) // CHARS_PER_TOKEN
             fits_bytes = byte_budget is None or (
                 _numbered_size_bound(content, last_line) <= byte_budget
             )
-            if content_tokens <= content_budget and fits_bytes:
+            whole = Batch(file_path=path, content=content, findings=file_findings)
+            if (
+                content_tokens <= content_budget
+                and fits_bytes
+                and self._prompt_fits_budget(self.build_prompt(whole, **prompt_kwargs))
+            ):
                 batches.append(
                     Batch(
                         file_path=path,
@@ -1405,13 +1515,38 @@ def get_batches(
                     )
                 )
             else:
-                chunk_budget = max(int(content_budget), 1024)
+                chunk_budget = content_budget
                 # A non-positive byte budget means the fixed prompt alone is over the
                 # cap; splitting the content further cannot help, and every batch is
                 # reported below and fails closed at the provider.
                 chunk_bytes = byte_budget if byte_budget is not None and byte_budget > 0 else None
+
+                def fits_chunk(
+                    text: str,
+                    start: int,
+                    end: int,
+                    *,
+                    path: str = path,
+                    file_findings: list[Finding] = file_findings,
+                ) -> bool:
+                    return self._prompt_fits_budget(
+                        self.build_prompt(
+                            Batch(
+                                file_path=path,
+                                content=text,
+                                start_line=start,
+                                end_line=end,
+                                findings=findings_in_range(file_findings, start, end),
+                            ),
+                            **prompt_kwargs,
+                        )
+                    )
+
                 for chunk_text, s_line, e_line in chunk_file_by_lines(
-                    content, chunk_budget, max_bytes=chunk_bytes
+                    content,
+                    chunk_budget,
+                    max_bytes=chunk_bytes,
+                    fits_chunk=fits_chunk,
                 ):
                     chunk_findings = findings_in_range(file_findings, s_line, e_line)
                     batches.append(
@@ -1497,6 +1632,45 @@ def parse_response(self, response: object, batch: Batch) -> list[Finding]:
 
     # -- Run loop -----------------------------------------------------------
 
+    def _serialized_input(self, prompt: str) -> str:
+        if self.response_schema is not None:
+            return prompt + structured_output_instructions(self.response_schema)
+        return prompt
+
+    def _count_prompt_tokens(self, prompt: str, *, chat_model: object | None = None) -> int:
+        if chat_model is None:
+            counter = self._token_counter
+        else:
+            # Reuse a chosen fallback even if another request later populates
+            # the process-wide codec registry. Equivalent refreshed HTTP clients
+            # share a counter without keeping their connection pools alive.
+            counters = getattr(self, "_call_token_counters", None)
+            if counters is None:
+                counters = self._call_token_counters = {}
+            key = _input_token_counter_key(chat_model)
+            if key not in counters:
+                counters[key] = _input_token_counter(chat_model)[0]
+            counter = counters[key]
+        return counter(self._serialized_input(prompt)) + REQUEST_FRAMING_TOKENS
+
+    def _prompt_fits_budget(self, prompt: str) -> bool:
+        try:
+            self._check_prompt_budget(prompt)
+        except LLMInputBudgetError:
+            return False
+        return True
+
+    def _check_prompt_budget(self, prompt: str, *, chat_model: object | None = None) -> None:
+        """Reject an unfit serialized batch before it reaches a provider."""
+        size = len(self._serialized_input(prompt).encode("utf-8", errors="replace"))
+        count = self._count_prompt_tokens(prompt, chat_model=chat_model)
+        byte_cap = _cli_max_input_bytes()
+        if count > self._input_budget or (byte_cap is not None and size > byte_cap):
+            raise LLMInputBudgetError(
+                "serialized LLM input exceeds its token or byte budget; "
+                "batch content retained but not submitted"
+            )
+
     def _record_failed_request(self, exc: BaseException, prompt: str, ordinal: int) -> None:
         # Diagnostics must never replace the original provider exception or report usage.
         try:
@@ -1536,6 +1710,7 @@ def _invoke_batch(self, batch: Batch, prompt: str) -> tuple[Batch, list]:
             len(batch.findings),
         )
         llm, structured_llm = self._model_for_call()
+        self._check_prompt_budget(prompt, chat_model=llm)
         if structured_llm:
             try:
                 response = self._invoke_with_failure_diagnostic(structured_llm, prompt)
@@ -1566,6 +1741,7 @@ def _formatting_call(self, source: str) -> tuple[object, str]:
                 "formatting input exceeds model budget", issue=ResponseIssue.LIMIT
             )
         llm, _ = self._model_for_call()
+        self._check_prompt_budget(prompt, chat_model=llm)
         return bind_structured_output(
             llm,
             schema,
@@ -1874,6 +2050,9 @@ def _recover_batch(self, batch: Batch, prompt: str, source: str) -> _RecoveredBa
                 return self._recovered_result(batch, items, False, exc)
             try:
                 llm, _ = self._model_for_call()
+                # Recovery repeats source plus quoted prior answers; measure the
+                # complete prompt against the actual refreshed request client too.
+                self._check_prompt_budget(question.prompt, chat_model=llm)
                 response = self._invoke_with_failure_diagnostic(llm, question.prompt)
                 if (getattr(response, "response_metadata", {}) or {}).get(
                     "completion_source"
@@ -1901,6 +2080,9 @@ async def _arecover_batch(self, batch: Batch, prompt: str, source: str) -> _Reco
                 return self._recovered_result(batch, items, False, exc)
             try:
                 llm, _ = self._model_for_call()
+                # Keep the async raw-recovery path under the same serialized
+                # token/byte limits without dropping already recovered findings.
+                self._check_prompt_budget(question.prompt, chat_model=llm)
                 response = await self._ainvoke_with_failure_diagnostic(llm, question.prompt)
                 if (getattr(response, "response_metadata", {}) or {}).get(
                     "completion_source"
@@ -2066,6 +2248,7 @@ async def _ainvoke_batch(self, batch: Batch, prompt: str) -> tuple[Batch, list]:
             len(batch.findings),
         )
         llm, structured_llm = self._model_for_call()
+        self._check_prompt_budget(prompt, chat_model=llm)
         if structured_llm:
             try:
                 response = await self._ainvoke_with_failure_diagnostic(structured_llm, prompt)