From 923140ea345f24ee1b06f839d83f71a280870461 Mon Sep 17 00:00:00 2001 From: Patrick Erichsen Date: Mon, 2 Nov 2026 10:34:21 -0800 Subject: [PATCH 1/4] fix(llm): bound serialized requests with configured tokenizers Signed-off-by: Patrick Erichsen --- src/skillspector/llm_analyzer_base.py | 188 +++++++++++++++--- .../test_semantic_security_discovery.py | 3 +- tests/nodes/test_llm_analyzer_base.py | 25 ++- tests/nodes/test_llm_prompt_budget.py | 179 +++++++++++++++++ 4 files changed, 365 insertions(+), 30 deletions(-) create mode 100644 tests/nodes/test_llm_prompt_budget.py diff --git a/src/skillspector/llm_analyzer_base.py b/src/skillspector/llm_analyzer_base.py index 386b27a08..e63a957ee 100644 --- a/src/skillspector/llm_analyzer_base.py +++ b/src/skillspector/llm_analyzer_base.py @@ -39,6 +39,7 @@ from random import uniform from typing import Any, Literal, cast +import tiktoken from langchain_anthropic import ChatAnthropic from langchain_core.messages import BaseMessage from langchain_openai.chat_models.base import BaseChatOpenAI @@ -164,6 +165,55 @@ def _require_structured_response(response: object) -> None: ) +REQUEST_FRAMING_TOKENS = 256 + + +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: + encoding = tiktoken.encoding_for_model(label) + except KeyError: + 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.""" @@ -899,6 +949,7 @@ 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*. @@ -906,10 +957,14 @@ def chunk_file_by_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: @@ -923,25 +978,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 @@ -1061,6 +1128,7 @@ def __init__( self._require_time_remaining() raise self._require_time_remaining() + self._token_counter, self._token_counter_source = _input_token_counter(self._llm) # 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. @@ -1180,13 +1248,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) @@ -1196,27 +1264,44 @@ 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 + fixed_bytes = self._fixed_prompt_bytes(path, file_findings, 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, @@ -1225,13 +1310,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( @@ -1317,6 +1427,33 @@ 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: + counter = self._token_counter if chat_model is None else _input_token_counter(chat_model)[0] + 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 _invoke_batch(self, batch: Batch, prompt: str) -> tuple[Batch, list]: """Invoke and parse one batch synchronously.""" logger.debug( @@ -1326,6 +1463,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 = _invoke_with_usage(structured_llm, prompt, self._usage_collector) @@ -1356,6 +1494,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, @@ -1472,6 +1611,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 _ainvoke_with_usage(structured_llm, prompt, self._usage_collector) diff --git a/tests/nodes/analyzers/test_semantic_security_discovery.py b/tests/nodes/analyzers/test_semantic_security_discovery.py index ad74002f8..ba8a6f0c3 100644 --- a/tests/nodes/analyzers/test_semantic_security_discovery.py +++ b/tests/nodes/analyzers/test_semantic_security_discovery.py @@ -314,7 +314,8 @@ def test_oversized_file_multiple_batches(self, base_state) -> None: with patch( "skillspector.llm_analyzer_base.get_max_input_tokens", - return_value=50, + # Leave room for the fixed prompt/schema; only source needs splitting. + return_value=8192, ): with patch.object( LLMAnalyzerBase, "arun_batches", new_callable=AsyncMock, return_value=[] diff --git a/tests/nodes/test_llm_analyzer_base.py b/tests/nodes/test_llm_analyzer_base.py index b854cc8c7..61b1beccc 100644 --- a/tests/nodes/test_llm_analyzer_base.py +++ b/tests/nodes/test_llm_analyzer_base.py @@ -1141,13 +1141,22 @@ def test_cli_provider_splits_multibyte_file_under_token_budget( assert len(analyzer.build_prompt(batch).encode("utf-8")) <= MAX_INPUT_BYTES @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) - def test_http_provider_batching_unchanged(self, http_provider: None) -> None: + def test_http_provider_serialized_prompts_fit_and_keep_all_lines( + self, http_provider: None + ) -> None: analyzer = self._analyzer() content = "".join("z" * 10 + "\n" for _ in range(150_000)) - budget = max(analyzer._input_budget - estimate_tokens("test"), 1024) - expected = [(s, e) for _t, s, e in chunk_file_by_lines(content, budget)] - batches = analyzer.get_batches(["big.js"], {"big.js": content}) - assert [(b.start_line, b.end_line) for b in batches] == expected + covered: set[int] = set() + for batch in analyzer.get_batches(["big.js"], {"big.js": content}): + prompt = analyzer.build_prompt(batch) + if analyzer.response_schema is not None: + from skillspector.llm_utils import structured_output_instructions + + prompt += structured_output_instructions(analyzer.response_schema) + assert (len(prompt.encode("utf-8")) + 3) // 4 <= analyzer._input_budget + assert batch.end_line is not None + covered.update(range(batch.start_line, batch.end_line + 1)) + assert covered == set(range(1, 150_001)) @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) def test_cli_payload_with_schema_instructions_fits_without_margin( @@ -3208,6 +3217,12 @@ def test_cli_batches_fit_with_large_metadata_and_findings( content = "".join("m" * 40 + "\n" for _ in range(30_000)) findings = [self._make_finding("big.js", line=i) for i in range(1, 30_000, 100)] metadata = "d" * 400_000 + # Exercise the CLI byte cap with an explicit known local counting fixture; + # unknown codecs instead use the conservative byte-bound refusal test. + import tiktoken + + codec = tiktoken.get_encoding("cl100k_base") + analyzer._token_counter = lambda text: len(codec.encode(text, disallowed_special=())) batches = analyzer.get_batches( ["big.js"], {"big.js": content}, findings, metadata_text=metadata ) diff --git a/tests/nodes/test_llm_prompt_budget.py b/tests/nodes/test_llm_prompt_budget.py new file mode 100644 index 000000000..d4de8dc6a --- /dev/null +++ b/tests/nodes/test_llm_prompt_budget.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serialized prompt budgets must include every line and the real prompt wrapper.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +import skillspector.llm_analyzer_base as lab +from skillspector.models import Finding + + +@pytest.fixture +def analyzer(monkeypatch: pytest.MonkeyPatch) -> lab.LLMAnalyzerBase: + """Construct the real planner with an inert provider; no model credentials or calls.""" + monkeypatch.setattr(lab, "get_max_input_tokens", lambda model: 8192) + monkeypatch.setattr(lab, "get_chat_model", lambda **kwargs: MagicMock()) + monkeypatch.setattr(lab, "_uses_native_connection_retries", lambda *args, **kwargs: False) + monkeypatch.setattr(lab, "new_inference_usage_collector", lambda **kwargs: MagicMock()) + monkeypatch.setattr(lab.LLMAnalyzerBase, "_bind_structured_output", lambda *args: MagicMock()) + monkeypatch.setattr(lab, "_cli_max_input_bytes", lambda: None) + return lab.LLMAnalyzerBase("Review every supplied line.", "fixture") + + +def test_short_line_rounding_does_not_exceed_chunk_budget() -> None: + chunks = lab.chunk_file_by_lines("ab\nab\nab", 1) + assert all(lab.estimate_tokens(text) <= 1 for text, _, _ in chunks) + assert {line for _, start, end in chunks for line in range(start, end + 1)} == {1, 2, 3} + + +@pytest.mark.parametrize("line", ["ab\n", "ab\r\n", "漢字🙂\n"]) +@pytest.mark.parametrize("compact", [False, True]) +def test_real_serialized_prompt_fits_and_keeps_every_line( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch, line: str, compact: bool +) -> None: + monkeypatch.setenv("SKILLSPECTOR_COMPACT_PROMPTS", "1" if compact else "0") + content = line * 5000 + batches = analyzer.get_batches(["source.js"], {"source.js": content}) + source_lines = content.splitlines(keepends=True) + covered: set[int] = set() + for batch in batches: + prompt = analyzer.build_prompt(batch) + if analyzer.response_schema is not None: + prompt += lab.structured_output_instructions(analyzer.response_schema) + assert len(prompt.encode("utf-8")) + lab.REQUEST_FRAMING_TOKENS <= analyzer._input_budget + end = batch.end_line or len(source_lines) + assert batch.content == "".join(source_lines[batch.start_line - 1 : end]) + covered.update(range(batch.start_line, end + 1)) + assert covered == set(range(1, len(source_lines) + 1)) + assert len(batches) < len(source_lines) + + +def test_real_custom_findings_and_metadata_overhead_is_measured( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original = analyzer.build_prompt + + def custom_prompt(batch: lab.Batch, **kwargs: object) -> str: + return ( + original(batch) + "\nMetadata: " + "m" * 1000 + "\n" + "detail " * len(batch.findings) + ) + + monkeypatch.setattr(analyzer, "build_prompt", custom_prompt) + finding = Finding( + rule_id="fixture", message="fixture", severity="LOW", file="source.js", start_line=1 + ) + batches = analyzer.get_batches(["source.js"], {"source.js": "line\n" * 4000}, [finding] * 100) + for batch in batches: + prompt = analyzer.build_prompt(batch) + lab.structured_output_instructions( + analyzer.response_schema + ) + assert len(prompt.encode("utf-8")) + lab.REQUEST_FRAMING_TOKENS <= analyzer._input_budget + + +def test_empty_file_stays_empty(analyzer: lab.LLMAnalyzerBase) -> None: + assert lab.chunk_file_by_lines("", 1) == [("", 1, 1)] + batches = analyzer.get_batches(["empty"], {"empty": ""}) + assert len(batches) == 1 + assert batches[0].content == "" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_unsplittable_line_is_an_explicit_failure_without_provider_call( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + provider = MagicMock(side_effect=AssertionError("provider must not be reached")) + monkeypatch.setattr(lab, "_invoke_with_usage", provider) + monkeypatch.setattr(lab, "_ainvoke_with_usage", provider) + content = "x" * 10000 + batches = analyzer.get_batches(["long"], {"long": content}) + assert len(batches) == 1 + assert batches[0].content == content + if asynchronous: + outcome = await analyzer.arun_batches_detailed(batches, max_concurrency=1) + else: + outcome = analyzer.run_batches_detailed(batches) + assert not outcome.successful + assert len(outcome.failures) == 1 + assert outcome.failures[0].error_class == "LLMInputBudgetError" + assert outcome.failures[0].reason is lab.LedgerReason.LLM_BATCH_FAILED + provider.assert_not_called() + + +def test_unknown_openai_codec_uses_conservative_bytes() -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = None + model.model_name = "unknown-future-model" + counter, source = lab._input_token_counter(model) + assert source == "conservative_utf8_bytes" + assert counter("漢🙂") == len("漢🙂".encode()) + + +def test_recognized_and_explicitly_configured_codec_is_respected() -> None: + import tiktoken + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "custom-provider-deployment" + counter, source = lab._input_token_counter(model) + text = "All source text stays data: <|endoftext|>" + assert source == "recognized_tiktoken_codec" + assert counter(text) == len( + tiktoken.encoding_for_model("gpt-4o").encode(text, disallowed_special=()) + ) + + +def test_explicit_custom_codec_is_respected_and_failure_is_conservative() -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = lambda text: [1, 2, 3] + counter, source = lab._input_token_counter(model) + assert source == "configured_chat_model_tokenizer" + assert counter("source") == 3 + model.custom_get_token_ids = MagicMock(side_effect=ValueError("codec unavailable")) + counter, _ = lab._input_token_counter(model) + assert counter("漢🙂") == len("漢🙂".encode()) + + +def test_actual_request_client_is_rechecked_when_its_codec_differs( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = None + model.model_name = "unknown-runtime-model" + monkeypatch.setattr(analyzer, "_model_for_call", lambda: (model, None)) + monkeypatch.setattr(analyzer, "_token_counter", lambda text: 1) + provider = MagicMock(side_effect=AssertionError("provider must not be reached")) + monkeypatch.setattr(lab, "_invoke_with_usage", provider) + batch = lab.Batch(file_path="source", content="x" * 10000) + with pytest.raises(lab.LLMInputBudgetError): + analyzer._invoke_batch(batch, analyzer.build_prompt(batch)) + provider.assert_not_called() + + +def test_unfit_fixed_wrapper_retains_one_batch_without_per_line_refinement( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original = analyzer.build_prompt + rendered = MagicMock(side_effect=lambda batch, **kwargs: "metadata" * 10000 + original(batch)) + monkeypatch.setattr(analyzer, "build_prompt", rendered) + content = "source line\n" * 10000 + batches = analyzer.get_batches(["source"], {"source": content}) + assert len(batches) == 1 + assert batches[0].content == content + assert rendered.call_count <= 3 + with pytest.raises(lab.LLMInputBudgetError): + analyzer._check_prompt_budget(analyzer.build_prompt(batches[0])) From 952c842d7a5f620810b357f9af780e3b09c12073 Mon Sep 17 00:00:00 2001 From: Patrick Erichsen Date: Mon, 2 Nov 2026 10:46:02 -0800 Subject: [PATCH 2/4] fix(llm): fall back when tokenizer initialization fails Signed-off-by: Patrick Erichsen --- src/skillspector/llm_analyzer_base.py | 4 +++- tests/nodes/test_llm_prompt_budget.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/skillspector/llm_analyzer_base.py b/src/skillspector/llm_analyzer_base.py index e63a957ee..e16f82ebd 100644 --- a/src/skillspector/llm_analyzer_base.py +++ b/src/skillspector/llm_analyzer_base.py @@ -198,7 +198,9 @@ def configured(text: str) -> int: return byte_bound, "conservative_utf8_bytes" try: encoding = tiktoken.encoding_for_model(label) - except KeyError: + except Exception: + # Recognized codecs can still fail to initialize when the cache is + # empty and its loader cannot reach or validate the vocabulary. return byte_bound, "conservative_utf8_bytes" def recognized(text: str) -> int: diff --git a/tests/nodes/test_llm_prompt_budget.py b/tests/nodes/test_llm_prompt_budget.py index d4de8dc6a..8abae04ce 100644 --- a/tests/nodes/test_llm_prompt_budget.py +++ b/tests/nodes/test_llm_prompt_budget.py @@ -177,3 +177,22 @@ def test_unfit_fixed_wrapper_retains_one_batch_without_per_line_refinement( assert rendered.call_count <= 3 with pytest.raises(lab.LLMInputBudgetError): analyzer._check_prompt_budget(analyzer.build_prompt(batches[0])) + + +@pytest.mark.parametrize("failure", [OSError, ValueError, RuntimeError]) +def test_codec_initialization_failure_uses_conservative_bytes( + monkeypatch: pytest.MonkeyPatch, failure: type[Exception] +) -> None: + import tiktoken + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + loader = MagicMock(side_effect=failure("offline or invalid tokenizer vocabulary")) + monkeypatch.setattr(tiktoken, "encoding_for_model", loader) + counter, source = lab._input_token_counter(model) + assert source == "conservative_utf8_bytes" + assert counter("漢🙂") == len("漢🙂".encode()) + loader.assert_called_once_with("gpt-4o") From 09552075e50f1f49ae7c5fcbcaac8ea31ca8e939 Mon Sep 17 00:00:00 2001 From: Patrick Erichsen Date: Mon, 2 Nov 2026 11:02:47 -0800 Subject: [PATCH 3/4] fix(llm): keep tokenizer selection offline and chunk planning bounded Signed-off-by: Patrick Erichsen --- src/skillspector/llm_analyzer_base.py | 45 ++++++++- tests/nodes/test_llm_prompt_budget.py | 139 ++++++++++++++++++++++++-- 2 files changed, 172 insertions(+), 12 deletions(-) diff --git a/src/skillspector/llm_analyzer_base.py b/src/skillspector/llm_analyzer_base.py index e16f82ebd..852fa636a 100644 --- a/src/skillspector/llm_analyzer_base.py +++ b/src/skillspector/llm_analyzer_base.py @@ -44,6 +44,7 @@ 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, @@ -168,6 +169,21 @@ 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. @@ -197,10 +213,14 @@ def configured(text: str) -> int: if not isinstance(label, str): return byte_bound, "conservative_utf8_bytes" try: - encoding = tiktoken.encoding_for_model(label) + # 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: - # Recognized codecs can still fail to initialize when the cache is - # empty and its loader cannot reach or validate the vocabulary. + return byte_bound, "conservative_utf8_bytes" + if encoding is None: return byte_bound, "conservative_utf8_bytes" def recognized(text: str) -> int: @@ -1131,6 +1151,7 @@ def __init__( 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. @@ -1272,7 +1293,9 @@ def get_batches( file_findings = findings_by_file.get(path, []) last_line = max(len(content.splitlines()), 1) - fixed_bytes = self._fixed_prompt_bytes(path, file_findings, last_line, **prompt_kwargs) + # 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 @@ -1435,7 +1458,19 @@ def _serialized_input(self, prompt: str) -> str: return prompt def _count_prompt_tokens(self, prompt: str, *, chat_model: object | None = None) -> int: - counter = self._token_counter if chat_model is None else _input_token_counter(chat_model)[0] + 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: diff --git a/tests/nodes/test_llm_prompt_budget.py b/tests/nodes/test_llm_prompt_budget.py index 8abae04ce..86f667bbb 100644 --- a/tests/nodes/test_llm_prompt_budget.py +++ b/tests/nodes/test_llm_prompt_budget.py @@ -116,20 +116,24 @@ def test_unknown_openai_codec_uses_conservative_bytes() -> None: assert counter("漢🙂") == len("漢🙂".encode()) -def test_recognized_and_explicitly_configured_codec_is_respected() -> None: - import tiktoken +def test_recognized_and_explicitly_configured_codec_is_respected( + monkeypatch: pytest.MonkeyPatch, +) -> None: from langchain_openai import ChatOpenAI + from tiktoken import registry model = MagicMock(spec=ChatOpenAI) model.custom_get_token_ids = None model.tiktoken_model_name = "gpt-4o" model.model_name = "custom-provider-deployment" + codec = MagicMock() + codec.encode.return_value = [1, 2, 3] + monkeypatch.setattr(registry, "ENCODINGS", {"o200k_base": codec}) counter, source = lab._input_token_counter(model) text = "All source text stays data: <|endoftext|>" assert source == "recognized_tiktoken_codec" - assert counter(text) == len( - tiktoken.encoding_for_model("gpt-4o").encode(text, disallowed_special=()) - ) + assert counter(text) == 3 + codec.encode.assert_called_once_with(text, disallowed_special=()) def test_explicit_custom_codec_is_respected_and_failure_is_conservative() -> None: @@ -180,7 +184,7 @@ def test_unfit_fixed_wrapper_retains_one_batch_without_per_line_refinement( @pytest.mark.parametrize("failure", [OSError, ValueError, RuntimeError]) -def test_codec_initialization_failure_uses_conservative_bytes( +def test_codec_name_lookup_failure_uses_conservative_bytes( monkeypatch: pytest.MonkeyPatch, failure: type[Exception] ) -> None: import tiktoken @@ -191,8 +195,129 @@ def test_codec_initialization_failure_uses_conservative_bytes( model.tiktoken_model_name = "gpt-4o" model.model_name = "recognized-model" loader = MagicMock(side_effect=failure("offline or invalid tokenizer vocabulary")) - monkeypatch.setattr(tiktoken, "encoding_for_model", loader) + monkeypatch.setattr(tiktoken, "encoding_name_for_model", loader) counter, source = lab._input_token_counter(model) assert source == "conservative_utf8_bytes" assert counter("漢🙂") == len("漢🙂".encode()) loader.assert_called_once_with("gpt-4o") + + +@pytest.mark.parametrize("cached", [False, True]) +def test_codec_selection_never_loads_or_downloads_vocabulary( + monkeypatch: pytest.MonkeyPatch, cached: bool +) -> None: + import requests + import tiktoken + from langchain_openai import ChatOpenAI + from tiktoken import registry + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + codec = MagicMock() + codec.encode.return_value = [1, 2, 3] + monkeypatch.setattr(registry, "ENCODINGS", {"o200k_base": codec} if cached else {}) + loader = MagicMock(side_effect=AssertionError("Vocabulary loading/network is forbidden")) + for owner, name in [ + (tiktoken, "encoding_for_model"), + (tiktoken, "get_encoding"), + (requests, "get"), + ]: + monkeypatch.setattr(owner, name, loader) + counter, source = lab._input_token_counter(model) + text = "漢🙂 source" + assert counter(text) == (3 if cached else len(text.encode())) + assert source == ("recognized_tiktoken_codec" if cached else "conservative_utf8_bytes") + loader.assert_not_called() + + +def test_cached_only_constructor_cannot_consume_shared_deadline( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + import tiktoken + from langchain_openai import ChatOpenAI + from tiktoken import registry + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + monkeypatch.setattr(registry, "ENCODINGS", {}) + clock = [0.0] + + def unbounded_loader(_label): + clock[0] += 1000.0 + return MagicMock() + + loader = MagicMock(side_effect=unbounded_loader) + monkeypatch.setattr(tiktoken, "encoding_for_model", loader) + monkeypatch.setattr(lab, "get_chat_model", lambda **kwargs: model) + created = lab.LLMAnalyzerBase("Review all lines", "gpt-4o", timeout=lambda: 1.0 - clock[0]) + assert created._remaining_timeout() == 1.0 + assert created._token_counter_source == "conservative_utf8_bytes" + loader.assert_not_called() + + +def test_selected_fallback_is_reused_for_equivalent_request_clients( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + from langchain_openai import ChatOpenAI + from tiktoken import registry + + def client(): + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + return model + + model = client() + monkeypatch.setattr(registry, "ENCODINGS", {}) + before = analyzer._count_prompt_tokens("漢🙂 source", chat_model=model) + codec = MagicMock() + codec.encode.return_value = [1] + registry.ENCODINGS["o200k_base"] = codec + assert analyzer._count_prompt_tokens("漢🙂 source", chat_model=model) == before + assert analyzer._count_prompt_tokens("漢🙂 source", chat_model=client()) == before + codec.encode.assert_not_called() + + +def test_large_aggregate_findings_do_not_create_one_request_per_line( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original_prompt = analyzer.build_prompt + + def prompt(batch: lab.Batch, **kwargs: object) -> str: + return ( + original_prompt(batch) + "\nFindings:\n" + "\n".join(f.message for f in batch.findings) + ) + + monkeypatch.setattr(analyzer, "build_prompt", prompt) + content = "original source line\n" * 5000 + findings = [ + Finding( + rule_id="fixture", + message="Retained finding detail " * 60, + severity="HIGH", + file="source", + start_line=line, + ) + for line in range(1, 5001, 80) + ] + assert sum(len(f.message) for f in findings) > 4 * analyzer._input_budget + batches = analyzer.get_batches(["source"], {"source": content}, findings) + assert len(batches) < 150, ( + "Range-specific findings fit; aggregate metadata must not force one-line requests" + ) + covered = set() + retained = set() + lines = content.splitlines(keepends=True) + for batch in batches: + end = batch.end_line or len(lines) + assert batch.content == "".join(lines[batch.start_line - 1 : end]) + covered.update(range(batch.start_line, end + 1)) + retained.update(f.finding_id for f in batch.findings) + analyzer._check_prompt_budget(analyzer.build_prompt(batch)) + assert covered == set(range(1, len(lines) + 1)) + assert retained == {f.finding_id for f in findings} From 8e1c1c3a1c8f89f119e3c78e8cb685d25d5fe4a3 Mon Sep 17 00:00:00 2001 From: Patrick Erichsen Date: Mon, 2 Nov 2026 17:34:49 -0800 Subject: [PATCH 4/4] test(llm): initialize budget-aware redaction and formatting fixtures The redaction test bypassed constructor state required by serialized token accounting. Construct it with the existing inert provider fixture while retaining credential redaction assertions. The formatting-limit test set the input budget before the initial request, so the new submission guard correctly stopped before formatting. Reduce the budget after the successful mocked response and prove formatting refuses a second request. Signed-off-by: Patrick Erichsen --- tests/nodes/analyzers/test_dependency_sources.py | 8 ++++---- tests/unit/test_structured_response.py | 8 ++++++-- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/nodes/analyzers/test_dependency_sources.py b/tests/nodes/analyzers/test_dependency_sources.py index 7d67440b0..46999ec3a 100644 --- a/tests/nodes/analyzers/test_dependency_sources.py +++ b/tests/nodes/analyzers/test_dependency_sources.py @@ -14,7 +14,6 @@ from skillspector.llm_analyzer_base import Batch from skillspector.models import Finding from skillspector.nodes.meta_analyzer import ( - PER_FILE_ANALYSIS_PROMPT, LLMMetaAnalyzer, _fallback_filtered, _passthrough_with_defaults, @@ -1234,14 +1233,13 @@ def test_cargo_url_credentials_are_redacted_for_supported_schemes(destination: s assert "packages.example.invalid" in serialized -def test_sc10_credentials_are_redacted_before_provider_prompt_construction() -> None: +def test_sc10_credentials_are_redacted_before_provider_prompt_construction(mock_llm) -> None: username = "provider-user-sentinel" password = "provider-password-sentinel" token = "provider-token-sentinel" content = f"registry=ssh://{username}:{password}@packages.example.invalid/index?token={token}\n" finding = _analyze({".npmrc": content})[0] - analyzer = LLMMetaAnalyzer.__new__(LLMMetaAnalyzer) - analyzer.base_prompt = PER_FILE_ANALYSIS_PROMPT + analyzer = LLMMetaAnalyzer(model="fixture") analyzer._input_budget = 100_000 batch = analyzer.get_batches([".npmrc"], {".npmrc": content}, [finding])[0] @@ -1251,6 +1249,8 @@ def test_sc10_credentials_are_redacted_before_provider_prompt_construction() -> assert secret not in batch.content assert secret not in prompt assert "packages.example.invalid" in prompt + mock_llm.invoke.assert_not_called() + mock_llm.ainvoke.assert_not_called() def test_hidden_source_finding_is_marked_local_only() -> None: diff --git a/tests/unit/test_structured_response.py b/tests/unit/test_structured_response.py index d0613bf33..36e68b385 100644 --- a/tests/unit/test_structured_response.py +++ b/tests/unit/test_structured_response.py @@ -479,10 +479,14 @@ def test_missing_or_extra_assessment_data_is_not_sent_for_formatting(http_model, def test_formatting_rechecks_the_input_allowance(http_model): - requests = http_model([('{"findings": [],}', "stop")]) + def reduce_input_allowance(_request_number): + analyzer._input_budget = 1 + + requests = http_model([('{"findings": [],}', "stop")], after_response=reduce_input_allowance) analyzer = make_analyzer("discovery") - analyzer._input_budget = 1 outcome = analyzer.run_batches_detailed([Batch(file_path="a.py", content="x")]) + assert not outcome.successful + assert len(outcome.failures) == 1 assert outcome.failures[0].reason == LedgerReason.LLM_RESPONSE_LIMIT assert len(requests) == 1