From d8d08facdedccaab159951c1b516bef0f3ff2245 Mon Sep 17 00:00:00 2001 From: "aditya.chhabra" Date: Mon, 23 Mar 2026 14:05:55 +0000 Subject: [PATCH] feat(acms): budget enforcement with max_file_size and max_total_size constraints Implements byte-size budget enforcement for the ACMS context assembly pipeline. enforce_size_budget() filters context fragments against max_file_size (per-fragment) and max_total_size (cumulative) limits defined in a ContextView. New domain models BudgetViolation and BudgetEnforcementResult provide structured violation reporting. Pipeline integration in ACMSPipeline.assemble() applies enforcement as a pre-filter when a context_view is provided. Review feedback addressed: - Fixed PR milestone to v3.4.0 (matching ticket #847) - Rebased branch onto latest master - Added CHANGELOG.md entry - Extracted duplicated enforcement logic into _apply_budget_enforcement() shared method on ACMSPipeline, called by both parent and subclass - Fixed _make_fragment return type from object to ContextFragment - Moved all imports to top of files (acms_pipeline.py, steps file) - Added errors="replace" to encode("utf-8") for surrogate safety - Documented thread-safety caveat on last_enforcement_result property - Added short-circuit early return when both limits are None - Added multi-byte Unicode content test scenario - Added total_size assertions to key scenarios - Added violation type verification to mixed limits scenario - Merged duplicate singular/plural step definitions - Added edge case scenarios (empty list, all exceed, exact boundary) - Fixed misleading docstring in ContextAssemblyPipeline.assemble() - Re-exported VALID_PHASES through core __init__.py public API ISSUES CLOSED: #847 --- CHANGELOG.md | 8 + features/m5_acms_smoke.feature | 93 ++++++++++ features/steps/m5_acms_smoke_steps.py | 164 ++++++++++++++++++ robot/acms_pipeline.robot | 8 + robot/helper_acms_pipeline.py | 82 ++++++++- .../application/services/acms_pipeline.py | 9 +- .../application/services/acms_service.py | 63 ++++++- .../domain/models/core/__init__.py | 8 + .../domain/models/core/context_policy.py | 162 +++++++++++++++++ 9 files changed, 592 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 53c68594e..060c49280 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -77,6 +77,14 @@ that demoted fragments must accumulate fresh accesses before re-promotion, preventing staleness enforcement from being immediately undone by a single access. (#821) +- Added byte-size budget enforcement for the ACMS context assembly + pipeline. `enforce_size_budget()` filters context fragments against + `max_file_size` (per-fragment) and `max_total_size` (cumulative) + limits defined in a `ContextView`. New domain models + `BudgetViolation` and `BudgetEnforcementResult` provide structured + violation reporting. Pipeline integration in `ACMSPipeline.assemble()` + applies enforcement as a pre-filter when a `context_view` is + provided. (#847) - Aligned plan lifecycle model with specification: ERRORED is now terminal in `is_terminal`, per-phase state validation enforces APPLIED/CONSTRAINED to APPLY-only and COMPLETE to diff --git a/features/m5_acms_smoke.feature b/features/m5_acms_smoke.feature index 315ef28b5..3bb2a2255 100644 --- a/features/m5_acms_smoke.feature +++ b/features/m5_acms_smoke.feature @@ -186,3 +186,96 @@ Feature: M5 ACMS pipeline and large-project context smoke tests Given a m5 smoke context view excluding "**/__pycache__/**" Then the m5 smoke path "src/__pycache__/module.pyc" should be excluded And the m5 smoke path "src/module.py" should not be excluded + + # --- Budget enforcement (enforce_size_budget) --- + + Scenario: M5 smoke enforce_size_budget excludes oversized fragments + Given m5 smoke fragments of sizes 50 200 80 + And a m5 smoke context view with max_file_size 100 + When I m5 smoke enforce the size budget + Then 2 m5 smoke fragments should be accepted + And 1 m5 smoke fragments should be violated + And the m5 smoke violation type should be "max_file_size" + And the m5 smoke total accepted size should be 130 + + Scenario: M5 smoke enforce_size_budget caps total size + Given m5 smoke fragments of sizes 50 60 70 + And a m5 smoke context view with max_total_size 120 + When I m5 smoke enforce the size budget + Then 2 m5 smoke fragments should be accepted + And 1 m5 smoke fragments should be violated + And the m5 smoke violation type should be "max_total_size" + And the m5 smoke total accepted size should be 110 + + Scenario: M5 smoke enforce_size_budget with mixed limits + Given m5 smoke fragments of sizes 30 200 40 50 + And a m5 smoke context view with max_file_size 100 and max_total_size 100 + When I m5 smoke enforce the size budget + Then 2 m5 smoke fragments should be accepted + And 2 m5 smoke fragments should be violated + And the m5 smoke violations should include type "max_file_size" + And the m5 smoke violations should include type "max_total_size" + + Scenario: M5 smoke enforce_size_budget with no limits accepts all + Given m5 smoke fragments of sizes 500 600 700 + And a m5 smoke context view with no size limits + When I m5 smoke enforce the size budget + Then 3 m5 smoke fragments should be accepted + And 0 m5 smoke fragments should be violated + And the m5 smoke total accepted size should be 1800 + + Scenario: M5 smoke enforce_size_budget boundary exact fit + Given m5 smoke fragments of sizes 50 50 + And a m5 smoke context view with max_total_size 100 + When I m5 smoke enforce the size budget + Then 2 m5 smoke fragments should be accepted + And 0 m5 smoke fragments should be violated + And the m5 smoke total accepted size should be 100 + + Scenario: M5 smoke enforce_size_budget reports violation details + Given m5 smoke fragments of sizes 200 + And a m5 smoke context view with max_file_size 100 + When I m5 smoke enforce the size budget + Then the m5 smoke violation reason should contain "exceeded max_file_size" + And the m5 smoke violation content size should be 200 + + Scenario: M5 smoke pipeline applies size budget with context view + Given m5 smoke fragments of sizes 50 200 80 + And a m5 smoke context view with max_file_size 100 + When I m5 smoke assemble via pipeline with the context view + Then the m5 smoke pipeline should return 2 fragments + And the m5 smoke pipeline enforcement result should have 1 violation(s) + + # --- Budget enforcement edge cases --- + + Scenario: M5 smoke enforce_size_budget with empty fragments list + Given m5 smoke fragments of sizes + And a m5 smoke context view with max_file_size 100 + When I m5 smoke enforce the size budget + Then 0 m5 smoke fragments should be accepted + And 0 m5 smoke fragments should be violated + And the m5 smoke total accepted size should be 0 + + Scenario: M5 smoke enforce_size_budget all fragments exceed max_file_size + Given m5 smoke fragments of sizes 200 300 400 + And a m5 smoke context view with max_file_size 100 + When I m5 smoke enforce the size budget + Then 0 m5 smoke fragments should be accepted + And 3 m5 smoke fragments should be violated + And the m5 smoke total accepted size should be 0 + + Scenario: M5 smoke enforce_size_budget single fragment at exact max_file_size boundary + Given m5 smoke fragments of sizes 100 + And a m5 smoke context view with max_file_size 100 + When I m5 smoke enforce the size budget + Then 1 m5 smoke fragments should be accepted + And 0 m5 smoke fragments should be violated + And the m5 smoke total accepted size should be 100 + + Scenario: M5 smoke enforce_size_budget with multi-byte Unicode content + Given m5 smoke fragments with unicode content + And a m5 smoke context view with max_file_size 10 + When I m5 smoke enforce the size budget + Then 1 m5 smoke fragments should be accepted + And 1 m5 smoke fragments should be violated + And the m5 smoke violation type should be "max_file_size" diff --git a/features/steps/m5_acms_smoke_steps.py b/features/steps/m5_acms_smoke_steps.py index 853d2c560..a30f50a42 100644 --- a/features/steps/m5_acms_smoke_steps.py +++ b/features/steps/m5_acms_smoke_steps.py @@ -14,12 +14,21 @@ from behave import given, then, when from behave.runner import Context from pydantic import ValidationError from typer.testing import CliRunner +from ulid import ULID +from cleveragents.application.services.acms_service import ACMSPipeline from cleveragents.cli.commands.context import app as context_app from cleveragents.domain.models.core.context import Context as ContextModel +from cleveragents.domain.models.core.context_fragment import ( + ContextBudget, + ContextFragment, + FragmentProvenance, +) from cleveragents.domain.models.core.context_policy import ( + BudgetEnforcementResult, ContextView, ProjectContextPolicy, + enforce_size_budget, ) _FIXTURES_DIR = Path(__file__).resolve().parents[1] / "fixtures" / "m5" @@ -712,3 +721,158 @@ def step_m5_path_not_excluded(context: Context, path: str) -> None: assert not PurePosixPath(path).match(context.m5_exclude_pattern), ( f"Expected '{path}' NOT to match exclude pattern" ) + + +# --------------------------------------------------------------------------- +# Budget enforcement (enforce_size_budget) +# --------------------------------------------------------------------------- + + +def _make_fragment(size: int, index: int) -> ContextFragment: + """Create a ContextFragment with content of exactly *size* bytes.""" + # ASCII 'x' is 1 byte in UTF-8, so a string of length=size gives size bytes. + return ContextFragment( + content="x" * size, + token_count=max(1, size // 4), + uko_node=f"uko://test/frag-{index}", + provenance=FragmentProvenance( + resource_uri=f"test://res-{index}", + location=f"file_{index}.py", + ), + ) + + +@given("m5 smoke fragments of sizes {sizes}") +def step_m5_fragments_of_sizes(context: Context, sizes: str) -> None: + size_list = [int(s.strip()) for s in sizes.split()] + context.m5_fragments = [_make_fragment(s, i) for i, s in enumerate(size_list)] + + +@given("m5 smoke fragments of sizes") +def step_m5_fragments_empty(context: Context) -> None: + context.m5_fragments = [] + + +@given("m5 smoke fragments with unicode content") +def step_m5_fragments_unicode(context: Context) -> None: + """Create two fragments with multi-byte UTF-8 content. + + Fragment 0: 5 x e-acute = 10 bytes (each is 2 bytes in UTF-8) -- fits in 10. + Fragment 1: 6 x e-acute = 12 bytes -- exceeds max_file_size of 10. + """ + context.m5_fragments = [ + ContextFragment( + content="é" * 5, + token_count=5, + uko_node="uko://test/frag-0", + provenance=FragmentProvenance( + resource_uri="test://res-0", + location="file_0.py", + ), + ), + ContextFragment( + content="é" * 6, + token_count=6, + uko_node="uko://test/frag-1", + provenance=FragmentProvenance( + resource_uri="test://res-1", + location="file_1.py", + ), + ), + ] + + +@given("a m5 smoke context view with max_file_size {mfs:d} and max_total_size {mts:d}") +def step_m5_view_mixed_limits(context: Context, mfs: int, mts: int) -> None: + context.m5_view = ContextView(max_file_size=mfs, max_total_size=mts) + + +@when("I m5 smoke enforce the size budget") +def step_m5_enforce_budget(context: Context) -> None: + context.m5_enforcement = enforce_size_budget(context.m5_fragments, context.m5_view) + + +@then("{count:d} m5 smoke fragments should be accepted") +def step_m5_accepted_count(context: Context, count: int) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert len(result.accepted) == count, ( + f"Expected {count} accepted, got {len(result.accepted)}" + ) + + +@then("{count:d} m5 smoke fragments should be violated") +def step_m5_violation_count(context: Context, count: int) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert len(result.violations) == count, ( + f"Expected {count} violations, got {len(result.violations)}" + ) + + +@then('the m5 smoke violation type should be "{vtype}"') +def step_m5_violation_type(context: Context, vtype: str) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert any(v.violation_type == vtype for v in result.violations), ( + f"No violation with type {vtype!r} found" + ) + + +@then('the m5 smoke violation reason should contain "{text}"') +def step_m5_violation_reason_contains(context: Context, text: str) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert any(text in v.reason for v in result.violations), ( + f"No violation reason contains {text!r}" + ) + + +@then("the m5 smoke violation content size should be {size:d}") +def step_m5_violation_content_size(context: Context, size: int) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert any(v.content_size == size for v in result.violations), ( + f"No violation with content_size={size}" + ) + + +@then("the m5 smoke total accepted size should be {size:d}") +def step_m5_total_size(context: Context, size: int) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + assert result.total_size == size, ( + f"Expected total_size={size}, got {result.total_size}" + ) + + +@then('the m5 smoke violations should include type "{vtype}"') +def step_m5_violations_include_type(context: Context, vtype: str) -> None: + result: BudgetEnforcementResult = context.m5_enforcement + types = {v.violation_type for v in result.violations} + assert vtype in types, f"Expected violation type {vtype!r} in {types}" + + +@when("I m5 smoke assemble via pipeline with the context view") +def step_m5_assemble_pipeline(context: Context) -> None: + pipeline = ACMSPipeline() + plan_id = str(ULID()) + budget = ContextBudget(max_tokens=4096) + payload = pipeline.assemble( + plan_id=plan_id, + fragments=context.m5_fragments, + budget=budget, + context_view=context.m5_view, + ) + context.m5_pipeline_payload = payload + context.m5_pipeline = pipeline + + +@then("the m5 smoke pipeline should return {count:d} fragments") +def step_m5_pipeline_fragment_count(context: Context, count: int) -> None: + assert len(context.m5_pipeline_payload.fragments) == count, ( + f"Expected {count}, got {len(context.m5_pipeline_payload.fragments)}" + ) + + +@then("the m5 smoke pipeline enforcement result should have {count:d} violation(s)") +def step_m5_pipeline_enforcement_violations(context: Context, count: int) -> None: + result = context.m5_pipeline.last_enforcement_result + assert result is not None, "No enforcement result on pipeline" + assert len(result.violations) == count, ( + f"Expected {count} violations, got {len(result.violations)}" + ) diff --git a/robot/acms_pipeline.robot b/robot/acms_pipeline.robot index 1a0efa1cd..df24f7efc 100644 --- a/robot/acms_pipeline.robot +++ b/robot/acms_pipeline.robot @@ -55,3 +55,11 @@ Verify Payload Budget Check Log ${result.stderr} Should Be Equal As Integers ${result.rc} 0 Should Contain ${result.stdout} acms-payload-budget-ok + +Assemble With Size Budget Context View + [Documentation] Verify max_file_size pre-filter and enforcement metadata + ${result}= Run Process ${PYTHON} ${HELPER} assemble-size-budget cwd=${WORKSPACE} + Log ${result.stdout} + Log ${result.stderr} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} acms-assemble-size-budget-ok diff --git a/robot/helper_acms_pipeline.py b/robot/helper_acms_pipeline.py index 47db4e485..9fe9d6735 100644 --- a/robot/helper_acms_pipeline.py +++ b/robot/helper_acms_pipeline.py @@ -10,6 +10,7 @@ Usage: python robot/helper_acms_pipeline.py assemble-recency python robot/helper_acms_pipeline.py assemble-tiered python robot/helper_acms_pipeline.py payload-budget-check + python robot/helper_acms_pipeline.py assemble-size-budget """ from __future__ import annotations @@ -20,8 +21,9 @@ from datetime import UTC, datetime from pathlib import Path _SRC = str(Path(__file__).resolve().parents[1] / "src") -if _SRC not in sys.path: - sys.path.insert(0, _SRC) +if _SRC in sys.path: + sys.path.remove(_SRC) +sys.path.insert(0, _SRC) from cleveragents.application.services.acms_service import ACMSPipeline # noqa: E402 from cleveragents.domain.models.core.context_fragment import ( # noqa: E402 @@ -29,6 +31,7 @@ from cleveragents.domain.models.core.context_fragment import ( # noqa: E402 ContextFragment, FragmentProvenance, ) +from cleveragents.domain.models.core.context_policy import ContextView # noqa: E402 # Default provenance for test fragments. _DEFAULT_PROV = FragmentProvenance(resource_uri="test://robot") @@ -260,6 +263,77 @@ def _cmd_payload_budget_check() -> int: return 0 +def _cmd_assemble_size_budget() -> int: + """Verify ACMSPipeline budget pre-filter with ContextView limits.""" + frags = [ + ContextFragment( + uko_node="project://app/small.py", + content="a" * 50, + relevance_score=0.9, + token_count=60, + provenance=_DEFAULT_PROV, + ), + ContextFragment( + uko_node="project://app/too_big.py", + content="b" * 200, + relevance_score=0.8, + token_count=60, + provenance=_DEFAULT_PROV, + ), + ContextFragment( + uko_node="project://app/fit.py", + content="c" * 80, + relevance_score=0.7, + token_count=60, + provenance=_DEFAULT_PROV, + ), + ] + budget = ContextBudget(max_tokens=1000, reserved_tokens=0) + pipeline = ACMSPipeline() + payload = pipeline.assemble( + plan_id="01JQTESTPN00000000000000AA", + fragments=frags, + budget=budget, + strategy="relevance", + context_view=ContextView(max_file_size=100), + ) + + result = pipeline.last_enforcement_result + if result is None: + print("acms-fail: expected non-empty enforcement result") + return 1 + if len(payload.fragments) != 2: + print(f"acms-fail: expected 2 fragments, got {len(payload.fragments)}") + return 1 + if len(result.accepted) != 2: + print(f"acms-fail: expected 2 accepted ids, got {len(result.accepted)}") + return 1 + if len(result.violations) != 1: + print(f"acms-fail: expected 1 violation, got {len(result.violations)}") + return 1 + if result.violations[0].violation_type != "max_file_size": + print( + "acms-fail: expected violation_type=max_file_size, " + f"got {result.violations[0].violation_type}" + ) + return 1 + + payload_ids = {fragment.fragment_id for fragment in payload.fragments} + if payload_ids != set(result.accepted): + print( + "acms-fail: payload fragment ids do not match accepted ids " + f"payload={sorted(payload_ids)} accepted={sorted(result.accepted)}" + ) + return 1 + + if result.total_size != 130: + print(f"acms-fail: expected total_size=130, got {result.total_size}") + return 1 + + print("acms-assemble-size-budget-ok: accepted=2 violations=1 total_size=130") + return 0 + + _COMMANDS: dict[str, Callable[[], int]] = { "fragment-create": _cmd_fragment_create, "budget-calc": _cmd_budget_calc, @@ -267,6 +341,7 @@ _COMMANDS: dict[str, Callable[[], int]] = { "assemble-recency": _cmd_assemble_recency, "assemble-tiered": _cmd_assemble_tiered, "payload-budget-check": _cmd_payload_budget_check, + "assemble-size-budget": _cmd_assemble_size_budget, } @@ -276,7 +351,8 @@ def main() -> int: print( "Usage: helper_acms_pipeline.py " "" + "|assemble-recency|assemble-tiered|payload-budget-check" + "|assemble-size-budget>" ) return 1 diff --git a/src/cleveragents/application/services/acms_pipeline.py b/src/cleveragents/application/services/acms_pipeline.py index 11a8538fc..8fd2fdc2d 100644 --- a/src/cleveragents/application/services/acms_pipeline.py +++ b/src/cleveragents/application/services/acms_pipeline.py @@ -73,6 +73,7 @@ from cleveragents.domain.models.core.context_fragment import ( build_provenance_map, compute_context_hash, ) +from cleveragents.domain.models.core.context_policy import ContextView logger = structlog.get_logger(__name__) @@ -567,11 +568,14 @@ class ContextAssemblyPipeline(ACMSPipeline): budget: ContextBudget, strategy: str | None = None, request: ContextRequest | None = None, + context_view: ContextView | None = None, ) -> ContextPayload: """Assemble context with per-stage timing instrumentation. Overrides :meth:`ACMSPipeline.assemble` to add structured logging - with per-stage millisecond timings. + with per-stage millisecond timings. When *context_view* is + provided, byte-size limits are enforced as a pre-filter before + strategy orchestration. """ import re @@ -589,6 +593,9 @@ class ContextAssemblyPipeline(ACMSPipeline): ) raise ValueError(msg) + # --- Pre-filter: byte-size budget enforcement --- + fragments = self._apply_budget_enforcement(fragments, context_view) + pipeline_start = time.monotonic() self._pipeline_logger.info( "Pipeline started", diff --git a/src/cleveragents/application/services/acms_service.py b/src/cleveragents/application/services/acms_service.py index 92dd5b8c3..d231fe350 100644 --- a/src/cleveragents/application/services/acms_service.py +++ b/src/cleveragents/application/services/acms_service.py @@ -25,6 +25,7 @@ from __future__ import annotations import re from collections.abc import Sequence from dataclasses import dataclass +from threading import local from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable import structlog @@ -44,6 +45,11 @@ from cleveragents.domain.models.core.context_fragment import ( build_provenance_map, compute_context_hash, ) +from cleveragents.domain.models.core.context_policy import ( + BudgetEnforcementResult, + ContextView, + enforce_size_budget, +) # Lazy import helper — resolved at runtime to avoid circular imports. _GreedyKnapsackPacker: type | None = None @@ -652,6 +658,7 @@ class ACMSPipeline: raise ValueError(msg) self._default_strategy = default_strategy self._logger = logger.bind(service="acms_pipeline") + self._enforcement_result_local = local() # All 10 pipeline components (default to pass-through stubs) self._strategy_selector = strategy_selector or DefaultStrategySelector() @@ -665,6 +672,38 @@ class ACMSPipeline: self._preamble_generator = preamble_generator or DefaultPreambleGenerator() self._skeleton_compressor = skeleton_compressor or DefaultSkeletonCompressor() + # ------------------------------------------------------------------ + # Budget enforcement helper + # ------------------------------------------------------------------ + + def _apply_budget_enforcement( + self, + fragments: Sequence[ContextFragment], + context_view: ContextView | None, + ) -> Sequence[ContextFragment]: + """Pre-filter *fragments* against the byte-size limits in *context_view*. + + Stores the :class:`BudgetEnforcementResult` on the pipeline so + callers can inspect it via :attr:`last_enforcement_result`. + + Returns the (potentially filtered) fragment sequence. + """ + if context_view is not None: + enforcement = enforce_size_budget(fragments, context_view) + self._enforcement_result_local.result = enforcement + if enforcement.violations: + accepted_set = set(enforcement.accepted) + fragments = [f for f in fragments if f.fragment_id in accepted_set] + self._logger.info( + "Size budget enforcement applied", + accepted=len(enforcement.accepted), + violations=len(enforcement.violations), + total_bytes=enforcement.total_size, + ) + else: + self._enforcement_result_local.result = None + return fragments + def assemble( self, plan_id: str, @@ -672,8 +711,16 @@ class ACMSPipeline: budget: ContextBudget, strategy: str | None = None, request: ContextRequest | None = None, + context_view: ContextView | None = None, ) -> ContextPayload: - """Assemble context fragments into a budget-constrained payload.""" + """Assemble context fragments into a budget-constrained payload. + + When *context_view* is provided its ``max_file_size`` and + ``max_total_size`` limits are enforced as a pre-filter before + strategy orchestration. Violations are logged and the + ``BudgetEnforcementResult`` is stored on the pipeline instance + as ``last_enforcement_result`` for caller inspection. + """ if not re.match(ULID_PATTERN, plan_id): msg = f"plan_id must be a valid ULID, got {plan_id!r}" raise ValueError(msg) @@ -687,6 +734,9 @@ class ACMSPipeline: ) raise ValueError(msg) + # --- Pre-filter: byte-size budget enforcement --- + fragments = self._apply_budget_enforcement(fragments, context_view) + self._logger.info( "Assembling context", plan_id=plan_id, @@ -758,6 +808,17 @@ class ACMSPipeline: provenance_map=provenance_map, ) + @property + def last_enforcement_result(self) -> BudgetEnforcementResult | None: + """Return the calling thread's most recent enforcement result. + + The result is stored in thread-local state so concurrent callers + using the same ``ACMSPipeline`` instance do not overwrite each + other's values. + """ + result = getattr(self._enforcement_result_local, "result", None) + return result if isinstance(result, BudgetEnforcementResult) else None + def register_strategy( self, name: str, diff --git a/src/cleveragents/domain/models/core/__init__.py b/src/cleveragents/domain/models/core/__init__.py index ebd23d2bb..05f36f7df 100644 --- a/src/cleveragents/domain/models/core/__init__.py +++ b/src/cleveragents/domain/models/core/__init__.py @@ -68,8 +68,12 @@ from cleveragents.domain.models.core.context_fragment import ( compute_context_hash, ) from cleveragents.domain.models.core.context_policy import ( + VALID_PHASES, + BudgetEnforcementResult, + BudgetViolation, ContextView, ProjectContextPolicy, + enforce_size_budget, ) from cleveragents.domain.models.core.correction import ( CorrectionAttempt, @@ -317,6 +321,7 @@ __all__ = [ "DEFAULT_SAFETY_PROFILE", "ROLE_PERMISSIONS", "VALID_JOB_TRANSITIONS", + "VALID_PHASES", "ActionState", "Actor", "ActorLimits", @@ -328,7 +333,9 @@ __all__ = [ "BindingMode", "BindingResult", "BudgetCheckResult", + "BudgetEnforcementResult", "BudgetLevel", + "BudgetViolation", "Change", "ChangeEntry", "ChangeOperation", @@ -519,6 +526,7 @@ __all__ = [ "classify_error", "compute_context_hash", "deserialize_job_payload", + "enforce_size_budget", "get_builtin_profile", "get_recovery_hints", "merge_invariants", diff --git a/src/cleveragents/domain/models/core/context_policy.py b/src/cleveragents/domain/models/core/context_policy.py index 3afcfabbf..c86de8161 100644 --- a/src/cleveragents/domain/models/core/context_policy.py +++ b/src/cleveragents/domain/models/core/context_policy.py @@ -20,13 +20,27 @@ first explicitly-set ``ContextView`` for that phase (or defaults). | execute | strategize | | apply | execute | +## Budget Enforcement + +``enforce_size_budget`` filters context fragments against the byte-size +limits defined in a ``ContextView``. It returns a +``BudgetEnforcementResult`` containing the accepted fragments and a +structured list of ``BudgetViolation`` entries describing which +fragments were excluded and why. + Based on ``docs/specification.md`` Context section and ADR-004. """ from __future__ import annotations +from collections.abc import Sequence +from typing import TYPE_CHECKING + from pydantic import BaseModel, ConfigDict, Field, field_validator +if TYPE_CHECKING: + from cleveragents.domain.models.core.context_fragment import ContextFragment + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -171,3 +185,151 @@ class ProjectContextPolicy(BaseModel): str_strip_whitespace=True, validate_assignment=True, ) + + +# --------------------------------------------------------------------------- +# Budget enforcement +# --------------------------------------------------------------------------- + + +class BudgetViolation(BaseModel, frozen=True): + """A single budget violation recording why a fragment was excluded. + + Produced by :func:`enforce_size_budget` when a fragment exceeds the + ``max_file_size`` limit or would push cumulative size past + ``max_total_size``. + """ + + fragment_id: str = Field( + ..., + description="ID of the excluded fragment", + ) + reason: str = Field( + ..., + description=( + "Human-readable explanation, e.g. 'exceeded max_file_size (2048 > 1024)'" + ), + ) + content_size: int = Field( + ..., + ge=0, + description="Size in bytes of the fragment content (UTF-8 encoded)", + ) + limit: int | None = Field( + default=None, + description="The limit that was breached (bytes), or None if N/A", + ) + violation_type: str = Field( + ..., + description="'max_file_size' or 'max_total_size'", + ) + + model_config = ConfigDict(str_strip_whitespace=True) + + +class BudgetEnforcementResult(BaseModel, frozen=True): + """Result of enforcing byte-size budget limits on context fragments. + + Contains both the accepted fragments and a structured report of + violations so callers can surface diagnostics to the user. + """ + + accepted: tuple[str, ...] = Field( + default=(), + description="Fragment IDs that passed budget enforcement", + ) + violations: tuple[BudgetViolation, ...] = Field( + default=(), + description="Details on each excluded fragment", + ) + total_size: int = Field( + default=0, + ge=0, + description="Cumulative byte size of all accepted fragments", + ) + + model_config = ConfigDict(str_strip_whitespace=True) + + +def enforce_size_budget( + fragments: Sequence[ContextFragment], + view: ContextView, +) -> BudgetEnforcementResult: + """Filter *fragments* against the byte-size limits in *view*. + + Applies two constraints from the resolved ``ContextView``: + + 1. **max_file_size** — any individual fragment whose UTF-8 encoded + content exceeds this limit is excluded. + 2. **max_total_size** — once the cumulative size of accepted + fragments reaches this limit, remaining fragments are excluded. + + Fragments are evaluated in the order supplied. ``None`` limits + mean "no restriction". + + Args: + fragments: Sequence of ``ContextFragment`` objects to evaluate. + view: The ``ContextView`` whose limits should be enforced. + + Returns: + A :class:`BudgetEnforcementResult` with accepted fragment IDs, + violation details, and cumulative size. + """ + max_file = view.max_file_size + max_total = view.max_total_size + + # Fast path: when no limits are set, accept everything without + # incurring per-fragment encode() overhead. + if max_file is None and max_total is None: + return BudgetEnforcementResult( + accepted=tuple(f.fragment_id for f in fragments), + violations=(), + total_size=sum( + len(f.content.encode("utf-8", errors="replace")) for f in fragments + ), + ) + + accepted_ids: list[str] = [] + violations: list[BudgetViolation] = [] + total_size = 0 + + for fragment in fragments: + content_bytes = len(fragment.content.encode("utf-8", errors="replace")) + + # Check per-fragment size limit + if max_file is not None and content_bytes > max_file: + violations.append( + BudgetViolation( + fragment_id=fragment.fragment_id, + reason=(f"exceeded max_file_size ({content_bytes} > {max_file})"), + content_size=content_bytes, + limit=max_file, + violation_type="max_file_size", + ), + ) + continue + + # Check cumulative size limit + if max_total is not None and total_size + content_bytes > max_total: + violations.append( + BudgetViolation( + fragment_id=fragment.fragment_id, + reason=( + f"would exceed max_total_size " + f"({total_size} + {content_bytes} > {max_total})" + ), + content_size=content_bytes, + limit=max_total, + violation_type="max_total_size", + ), + ) + continue + + accepted_ids.append(fragment.fragment_id) + total_size += content_bytes + + return BudgetEnforcementResult( + accepted=tuple(accepted_ids), + violations=tuple(violations), + total_size=total_size, + )