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

tests/stages/test_performance.py

@@ -85,9 +85,15 @@ class TestPriorityLimiter:
 
     def test_release_only_for_non_high(self) -> None:
         limiter = PriorityLimiter(max_concurrency=1)
-        limiter.acquire(Priority.HIGH)
-        limiter.release(Priority.HIGH)
-        assert limiter.active_count() == 0
+
+        async def test():
+            await limiter.acquire(Priority.HIGH)
+            # HIGH release is a no-op
+            limiter.release(Priority.HIGH)
+            # NORMAL still not acquired, so active_count should be 0
+            assert limiter.active_count() == 0
+
+        asyncio_run(test())
 
     def test_max_concurrency_reflection(self) -> None:
         limiter = PriorityLimiter(max_concurrency=5)
@@ -219,7 +225,7 @@ class TestLLMConcurrency:
         async def test():
             normal_done = False
 
-            async def slow():
+            async def slow(**kwargs):
                 nonlocal normal_done
                 await asyncio.sleep(0.1)
                 normal_done = True
@@ -311,5 +317,10 @@ class TestAsyncContextManager:
 
     def test_context_manager(self) -> None:
         limiter = PriorityLimiter(max_concurrency=1)
-        asyncio_run(limiter.__aenter__())
-        asyncio_run(limiter.__aexit__(None, None, None))
+
+        async def test():
+            ctx = await limiter.__aenter__()
+            assert ctx is limiter
+            await limiter.__aexit__(None, None, None)
+
+        asyncio_run(test())

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)

tests/test_context_budget_integration.py

@@ -0,0 +1,142 @@
+"""Integration tests for Stage 20: Context Budgeting with ResearchOrchestrator."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from nsct.orchestration.context_budget import (
+    ContextBudgetConfig,
+    ContextBudgetExhaustedError,
+    ContextBudgetTracker,
+)
+from nsct.orchestration.orchestrator import ResearchOrchestrator
+from nsct.config import AppSettings, AudioConfig, DatabaseConfig, LLMConfig, VisionConfig
+
+
+@pytest.fixture
+def mock_config():
+    """Minimal AppSettings for orchestrator tests."""
+    llm = LLMConfig(
+        base_url="http://test:8000/v1",
+        model="test-model",
+        max_concurrency=3,
+    )
+    vision = VisionConfig(
+        base_url="http://test:8001/v1",
+        model="test-vision",
+    )
+    audio = AudioConfig(
+        base_url="http://test:8002/v1",
+        model="test-audio",
+    )
+    postgres = DatabaseConfig(url="sqlite+aiosqlite:///")
+    return AppSettings(llm=llm, vision=vision, audio=audio, postgres=postgres)
+
+
+def test_orchestrator_has_context_budget_tracker(mock_config):
+    """Orchestrator muss context_budget_tracker als Property haben."""
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+    )
+    tracker = orchestrator.context_budget_tracker
+    assert isinstance(tracker, ContextBudgetTracker)
+    assert tracker.total_usage == 0
+
+
+def test_orchestrator_custom_context_budget(mock_config):
+    """Orchestrator kann ein custom ContextBudgetConfig akzeptieren."""
+    custom_config = ContextBudgetConfig(
+        planner_max_tokens=16000,
+        claim_extraction_max_tokens=24000,
+        contradiction_max_tokens=32000,
+        synthesis_max_tokens=64000,
+    )
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+        context_budget_config=custom_config,
+    )
+    assert orchestrator._context_budget_config.planner_max_tokens == 16000
+    assert orchestrator.context_budget_tracker._config.planner_max_tokens == 16000
+
+
+def test_orchestrator_track_tokens(mock_config):
+    """Orchestrator muss track_tokens für Stages unterstützen."""
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+    )
+    # Track planner tokens
+    orchestrator._track_context_tokens("planner", 5000)
+    assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 5000
+
+    # Track again — accumulation works
+    orchestrator._track_context_tokens("planner", 3000)
+    assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 8000
+
+    # get_usage returns correct data
+    usage = orchestrator.context_budget_tracker.get_usage()
+    assert usage["planner"]["usage_tokens"] == 8000.0
+    assert usage["planner"]["limit_tokens"] == 12000.0
+
+
+def test_orchestrator_context_budget_exhaustion(mock_config):
+    """Orchestrator wirft ContextBudgetExhaustedError bei Budget-Überschreitung."""
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+    )
+    # Track 12000 tokens to reach the limit
+    orchestrator._track_context_tokens("planner", 12000)
+    # 12000 tokens tracked → at limit → is_exhausted returns True
+    assert orchestrator.context_budget_tracker.is_exhausted("planner") is True
+
+    # One more token should raise
+    with pytest.raises(ContextBudgetExhaustedError) as exc_info:
+        orchestrator._track_context_tokens("planner", 1)
+
+    assert exc_info.value.stage_name == "planner"
+    assert exc_info.value.used_tokens == 12000
+    assert exc_info.value.limit_tokens == 12000
+
+
+def test_orchestrator_reset_clears_context_budget(mock_config):
+    """Orchestrator reset sollte auch context_budget_tracker zurücksetzen."""
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+    )
+    orchestrator._track_context_tokens("planner", 5000)
+    orchestrator._track_context_tokens("synthesis", 10000)
+    assert orchestrator.context_budget_tracker.total_usage == 15000
+
+    # Simulate reset of context budget tracker
+    orchestrator.context_budget_tracker.reset_all()
+    assert orchestrator.context_budget_tracker.total_usage == 0
+
+
+def test_orchestrator_budget_usage_in_report(mock_config):
+    """Context budget usage sollte im Report enthalten sein."""
+    orchestrator = ResearchOrchestrator(
+        config=mock_config,
+        research_id="test-research-id",
+        query="test query",
+    )
+    orchestrator._track_context_tokens("planner", 5000)
+    orchestrator._track_context_tokens("claim_extraction", 8000)
+
+    usage = orchestrator.context_budget_tracker.get_usage()
+
+    assert "planner" in usage
+    assert "claim_extraction" in usage
+    assert usage["planner"]["usage_tokens"] == 5000.0
+    assert usage["claim_extraction"]["usage_tokens"] == 8000.0
+    # synthesis and contradiction should be present but zero
+    assert usage["synthesis"]["usage_tokens"] == 0.0
+    assert usage["contradiction"]["usage_tokens"] == 0.0