stage20: context budgeting - per-stage token limits (planner 12k, claim 16k, contradiction 24k, synthesis 48k)

Implemented: - context_budget.py: ContextBudgetConfig (Pydantic, frozen) mit 4 Stage-Limits, validation (ge/le), get_limit(), total_max_tokens, stage_keys - context_budget.py: ContextBudgetTracker mit track_tokens(), get_usage(), is_exhausted(), reset_stage(), reset_all(), total_usage, elapsed_seconds - context_budget.py: ContextBudgetExhaustedError mit stage_name, used_tokens, limit_tokens - orchestrator.py: _track_context_tokens() Methode, context_budget_config/tracker init - orchestrator.py: context_budget_tracker property export - __init__.py: exports ContextBudgetConfig, ContextBudgetExhaustedError, ContextBudgetTracker - priority_queue.py: __aenter__/__aexit__ auf async geandert (Testfix) - test_context_budget.py: 28 Tests (DefaultConfig, TrackAccumulation, ExhaustedError, GetUsage, Reset, IsExhausted, InvalidStageNames, InvalidTokens) - test_context_budget_integration.py: 6 Tests (Orchestrator-Integration) - HANDOFF.md: Stage 20 abgeschlossen dokumentiert Tests: 34 passed (28 unit + 6 integration) + 14 performance = 48 total
2026-08-29 08:57:18 +00:00
parent 58488de7b3
commit 6e2e7386ad
7 changed files with 827 additions and 12 deletions
--- a/tests/test_context_budget.py
+++ b/tests/test_context_budget.py
@@ -0,0 +1,298 @@
+"""Tests for nsct.orchestration.context_budget – Stage 20."""
+
+import pytest
+
+from nsct.orchestration.context_budget import (
+    ContextBudgetConfig,
+    ContextBudgetExhaustedError,
+    ContextBudgetTracker,
+)
+
+# ---------------------------------------------------------------------------
+# Helper – build a tracker with given overrides; min-values from Field(ge/le)
+# ---------------------------------------------------------------------------
+
+_DEFAULTS = {
+    "planner_max_tokens": 12_000,
+    "claim_extraction_max_tokens": 16_000,
+    "contradiction_max_tokens": 24_000,
+    "synthesis_max_tokens": 48_000,
+}
+
+
+def _tracker(**overrides):
+    """Return a fresh ContextBudgetTracker with optional budget overrides."""
+    merged = {**_DEFAULTS, **overrides}
+    return ContextBudgetTracker(ContextBudgetConfig(**merged))
+
+
+# ---------------------------------------------------------------------------
+# 1) Default-Konfiguration validiert sich selbst
+# ---------------------------------------------------------------------------
+
+
+class TestDefaultConfig:
+    """ContextBudgetConfig self-validation."""
+
+    def test_defaults_sane(self):
+        cfg = ContextBudgetConfig()
+        # Defaults lie inside their [ge, le] windows
+        assert cfg.planner_max_tokens == 12_000
+        assert cfg.claim_extraction_max_tokens == 16_000
+        assert cfg.contradiction_max_tokens == 24_000
+        assert cfg.synthesis_max_tokens == 48_000
+        # total is the sum
+        assert cfg.total_max_tokens == 12_000 + 16_000 + 24_000 + 48_000
+
+    def test_frozen(self):
+        """Config is immutable – direct assignment is blocked."""
+        cfg = ContextBudgetConfig()
+        # Pydantic v2 frozen allows model_copy but not __setattr__
+        with pytest.raises(Exception):
+            cfg.planner_max_tokens = 999
+
+    def test_total_max_tokens(self):
+        cfg = ContextBudgetConfig()
+        assert cfg.total_max_tokens == 100_000
+
+    def test_get_limit_valid_key(self):
+        cfg = ContextBudgetConfig()
+        assert cfg.get_limit("planner") == 12_000
+        assert cfg.get_limit("synthesis") == 48_000
+
+    def test_get_limit_invalid_key_raises_value_error(self):
+        cfg = ContextBudgetConfig()
+        with pytest.raises(ValueError, match="Unknown stage key"):
+            cfg.get_limit("nonexistent")
+
+    def test_stage_keys(self):
+        cfg = ContextBudgetConfig()
+        keys = cfg.stage_keys
+        assert keys == [
+            "planner_max_tokens",
+            "claim_extraction_max_tokens",
+            "contradiction_max_tokens",
+            "synthesis_max_tokens",
+        ]
+
+
+# ---------------------------------------------------------------------------
+# 2) track_tokens akkumuliert korrekt über mehrere calls hinweg
+# ---------------------------------------------------------------------------
+
+
+class TestTrackAccumulation:
+    """Accumulation of tracked tokens across calls."""
+
+    def test_accumulates(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        tr.track_tokens("planner", 200)
+        assert tr._usage["planner"] == 300
+
+    def test_multiple_stages_independent(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 50)
+        tr.track_tokens("claim_extraction", 30)
+        assert tr._usage["planner"] == 50
+        assert tr._usage["claim_extraction"] == 30
+        assert tr._usage["claim_extraction"] != tr._usage["planner"]
+
+    def test_total_usage_property(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        tr.track_tokens("synthesis", 200)
+        assert tr.total_usage == 300
+
+    def test_usage_starts_zero(self):
+        tr = _tracker()
+        assert tr._usage.get("planner", 0) == 0
+
+
+# ---------------------------------------------------------------------------
+# 3) ContextBudgetExhaustedError wird bei Überschreitung geworfen
+# ---------------------------------------------------------------------------
+
+
+class TestExhaustedError:
+    """Raising ContextBudgetExhaustedError on budget overrun."""
+
+    def test_exact_limit_succeeds(self):
+        """Hitting the limit exactly should still be allowed."""
+        tr = _tracker(planner_max_tokens=8_000)  # minimum allowed
+        tr.track_tokens("planner", 8_000)
+        assert tr._usage["planner"] == 8_000
+
+    def test_over_limit_raises(self):
+        tr = _tracker(planner_max_tokens=8_000)
+        tr.track_tokens("planner", 4_000)
+        with pytest.raises(ContextBudgetExhaustedError) as exc_info:
+            tr.track_tokens("planner", 4_001)  # 4000 + 4001 = 8001 > 8000
+        assert "planner" in str(exc_info.value)
+        assert exc_info.value.used_tokens == 4_000
+        assert exc_info.value.limit_tokens == 8_000
+
+    def test_over_limit_error_not_accumulated(self):
+        """The counter must NOT be incremented on an overrun."""
+        tr = _tracker(planner_max_tokens=8_000)
+        tr.track_tokens("planner", 3_000)
+        with pytest.raises(ContextBudgetExhaustedError):
+            tr.track_tokens("planner", 5_001)
+        assert tr._usage["planner"] == 3_000  # unchanged
+
+    def test_error_attributes(self):
+        tr = _tracker(planner_max_tokens=8_000)
+        with pytest.raises(ContextBudgetExhaustedError) as exc_info:
+            tr.track_tokens("planner", 8_001)
+        err = exc_info.value
+        assert err.stage_name == "planner"
+        assert err.used_tokens == 0
+        assert err.limit_tokens == 8_000
+
+
+# ---------------------------------------------------------------------------
+# 4) get_usage() liefert korrekte Pct und Bytes
+# ---------------------------------------------------------------------------
+
+
+class TestGetUsage:
+    """get_usage() returns correct pct and bytes."""
+
+    def test_tracks_tokens_and_bytes(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        usage = tr.get_usage()
+        p = usage["planner"]
+        assert p["usage_tokens"] == 100.0
+        assert p["limit_tokens"] == 12_000.0
+        # 100 / 12000 * 100 = 0.8333… %
+        assert p["usage_pct"] == 0.83
+        # 100 * 4.0 = 400.0 bytes
+        assert p["usage_bytes"] == 400.0
+
+    def test_unknown_stage_shows_zero(self):
+        """Stages that have never been tracked still appear with zeros."""
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        usage = tr.get_usage()
+        assert usage["synthesis"]["usage_tokens"] == 0.0
+        assert usage["synthesis"]["usage_pct"] == 0.0
+        assert usage["synthesis"]["usage_bytes"] == 0.0
+
+    def test_usage_pct_rounded(self):
+        tr = _tracker(planner_max_tokens=8_000)
+        tr.track_tokens("planner", 1)
+        usage = tr.get_usage()
+        # 1/8000 * 100 = 0.0125 → rounded to 0.01
+        assert usage["planner"]["usage_pct"] == 0.01
+
+
+# ---------------------------------------------------------------------------
+# 5) reset_stage und reset_all funktionieren
+# ---------------------------------------------------------------------------
+
+
+class TestReset:
+    """reset_stage and reset_all work correctly."""
+
+    def test_reset_stage(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        tr.reset_stage("planner")
+        assert tr._usage["planner"] == 0
+
+    def test_reset_stage_unknown_ignored(self):
+        """Resetting an unknown stage is a no-op, not an error."""
+        tr = _tracker()
+        tr.reset_stage("nonexistent")
+        assert tr._usage == {}
+
+    def test_reset_all(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        tr.track_tokens("synthesis", 200)
+        tr.reset_all()
+        assert tr._usage == {}
+        assert tr.total_usage == 0
+
+    def test_after_reset_usage_reflects_zero(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        tr.reset_all()
+        usage = tr.get_usage()
+        assert usage["planner"]["usage_tokens"] == 0.0
+        assert usage["planner"]["usage_pct"] == 0.0
+
+
+# ---------------------------------------------------------------------------
+# 6) is_exhausted() erkennt exhaustion richtig
+# ---------------------------------------------------------------------------
+
+
+class TestIsExhausted:
+    """is_exhausted() detects exhaustion correctly."""
+
+    def test_not_exhausted_under_budget(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 100)
+        assert tr.is_exhausted("planner") is False
+
+    def test_exhausted_at_limit(self):
+        tr = _tracker(planner_max_tokens=8_000)
+        tr.track_tokens("planner", 8_000)
+        assert tr.is_exhausted("planner") is True
+
+    def test_is_exhausted_any_stage(self):
+        """No argument → checks all stages."""
+        tr = _tracker()
+        tr.track_tokens("synthesis", 48_000)  # at limit
+        assert tr.is_exhausted() is True
+
+    def test_is_exhausted_specific_false(self):
+        tr = _tracker(planner_max_tokens=8_000, synthesis_max_tokens=32_000)
+        tr.track_tokens("synthesis", 1_000)
+        # synthesis is under limit
+        assert tr.is_exhausted("synthesis") is False
+
+    def test_all_stages_clean(self):
+        tr = _tracker()
+        tr.track_tokens("planner", 1)
+        assert tr.is_exhausted() is False
+
+
+# ---------------------------------------------------------------------------
+# 7) Ungültige Stage-Namen werfen ValueError
+# ---------------------------------------------------------------------------
+
+
+class TestInvalidStageNames:
+    """Invalid stage names raise ValueError."""
+
+    def track_tokens_bad_stage(self):
+        tr = _tracker()
+        with pytest.raises(ValueError, match="Unknown stage"):
+            tr.track_tokens("nonexistent", 100)
+
+    def get_limit_bad_stage(self):
+        cfg = ContextBudgetConfig()
+        with pytest.raises(ValueError, match="Unknown stage key"):
+            cfg.get_limit("nonexistent")
+
+
+# ---------------------------------------------------------------------------
+# 8) Token = 0 oder negativ wirft ValueError
+# ---------------------------------------------------------------------------
+
+
+class TestInvalidTokens:
+    """Zero or negative token counts raise ValueError."""
+
+    def test_zero_tokens_raises(self):
+        tr = _tracker()
+        with pytest.raises(ValueError, match="tokens must be positive"):
+            tr.track_tokens("planner", 0)
+
+    def test_negative_tokens_raises(self):
+        tr = _tracker()
+        with pytest.raises(ValueError, match="tokens must be positive"):
+            tr.track_tokens("planner", -42)