feat(acms): budget enforcement with max_file_size and max_total_size constraints
CI / build (pull_request) Successful in 22s
CI / lint (pull_request) Successful in 3m20s
CI / quality (pull_request) Successful in 3m55s
CI / security (pull_request) Successful in 4m2s
CI / typecheck (pull_request) Successful in 4m13s
CI / integration_tests (pull_request) Successful in 6m50s
CI / unit_tests (pull_request) Successful in 7m19s
CI / docker (pull_request) Successful in 1m8s
CI / e2e_tests (pull_request) Successful in 9m53s
CI / coverage (pull_request) Successful in 11m28s
CI / status-check (pull_request) Successful in 1s
CI / benchmark-publish (pull_request) Has been skipped
CI / benchmark-regression (pull_request) Successful in 1h9m52s

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
This commit is contained in:
2026-03-23 14:05:55 +00:00
committed by Aditya Chhabra
parent 336a79b681
commit d8d08facde
9 changed files with 592 additions and 5 deletions
+8
View File
@@ -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
+93
View File
@@ -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"
+164
View File
@@ -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)}"
)
+8
View File
@@ -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
+79 -3
View File
@@ -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 "
"<fragment-create|budget-calc|assemble-relevance"
"|assemble-recency|assemble-tiered|payload-budget-check>"
"|assemble-recency|assemble-tiered|payload-budget-check"
"|assemble-size-budget>"
)
return 1
@@ -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",
@@ -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,
@@ -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",
@@ -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,
)