Charge auto-research LLM calls to the agenda that owns them

agents/agenda_budget.py

@@ -1,4 +1,13 @@
-"""Context attribution and hard pre-call reservations for agenda resources."""
+"""Context attribution and hard pre-call reservations for agenda resources.
+
+`agenda_scope` attributes the work inside it to one agenda. Legacy `call_llm`
+checks that agenda before contacting a provider and records the call on
+`agenda_token_ledger` plus `token_spent`. Grant-backed role calls keep the
+grant reservation as the `token_spent` clock — a parent settlement already
+moves actual usage from reserved to spent — and instead append a ledger row
+when the sub-reservation settles, so execution spend is visible before the
+grant closes. A budget of NULL, 0, or negative stays accounting-only.
+"""
 
 from __future__ import annotations
 
@@ -88,3 +97,204 @@ def resume(
         agenda_id,
         token_budget=token_budget,
     )
+
+
+PAUSED_BUDGET = "paused_budget"
+ACTIVE = "active"
+
+
+class AgendaBudgetExceededError(RuntimeError):
+    """Raised before an LLM call when the agenda's token budget is spent."""
+
+    def __init__(self, agenda_id: int, token_spent: int, token_budget: int):
+        self.agenda_id = int(agenda_id)
+        self.token_spent = int(token_spent)
+        self.token_budget = int(token_budget)
+        super().__init__(
+            f"agenda {agenda_id} token budget exhausted "
+            f"({token_spent}/{token_budget} tokens); raise token_budget "
+            f"above spent and reserved, then resume the agenda"
+        )
+
+
+def budget_cap(token_budget: object) -> int | None:
+    """Return a positive cap, or None when the budget is accounting-only.
+
+    NULL, 0, and negative budgets record usage and never pause the agenda.
+    """
+    if token_budget is None:
+        return None
+    try:
+        value = int(token_budget)
+    except (TypeError, ValueError):
+        return None
+    if value <= 0:
+        return None
+    return value
+
+
+def _live_grant_tokens(agenda_id: int) -> int:
+    """Settled LLM usage on grants that have not yet folded into token_spent.
+
+    Parent grant settlement adds those tokens to token_spent and leaves the
+    grant, so counting only still-active grants does not charge them twice.
+    """
+    from db import database as db
+
+    if not db.table_exists("resource_grant_usage_reservations"):
+        return 0
+    if not db.table_exists("resource_grants"):
+        return 0
+    row = db.fetchone(
+        """
+        SELECT COALESCE(SUM(u.tokens_used), 0) AS spent
+        FROM resource_grant_usage_reservations u
+        JOIN resource_grants g ON g.id = u.resource_grant_id
+        WHERE u.agenda_id=?
+          AND u.status='settled'
+          AND g.status='active'
+          AND g.expires_at > CURRENT_TIMESTAMP
+        """,
+        (int(agenda_id),),
+    )
+    return int((row or {}).get("spent") or 0)
+
+
+def _agenda_budget_row(agenda_id: int) -> dict | None:
+    from db import database as db
+
+    if not db.table_exists("research_agendas"):
+        return None
+    row = db.fetchone(
+        """
+        SELECT id, token_budget, token_spent, status
+        FROM research_agendas
+        WHERE id=?
+        """,
+        (int(agenda_id),),
+    )
+    return dict(row) if row else None
+
+
+def charged_tokens(agenda_id: int, token_spent: int) -> int:
+    """Tokens already spent plus LLM usage still inside an active grant."""
+    return int(token_spent or 0) + _live_grant_tokens(agenda_id)
+
+
+def _set_status(agenda_id: int, status: str, *, commit: bool) -> None:
+    from db import database as db
+
+    db.execute(
+        """
+        UPDATE research_agendas
+        SET status=?, is_active=?, updated_at=CURRENT_TIMESTAMP
+        WHERE id=?
+        """,
+        (status, 1 if status == ACTIVE else 0, int(agenda_id)),
+    )
+    if commit:
+        db.commit()
+
+
+def _pause_if_exhausted(agenda_id: int, *, commit: bool) -> None:
+    row = _agenda_budget_row(agenda_id)
+    if not row:
+        return
+    cap = budget_cap(row.get("token_budget"))
+    if cap is None:
+        return
+    charged = charged_tokens(agenda_id, int(row.get("token_spent") or 0))
+    if charged >= cap and str(row.get("status") or "") != PAUSED_BUDGET:
+        _set_status(agenda_id, PAUSED_BUDGET, commit=commit)
+
+
+def check_budget(agenda_id: int) -> None:
+    """Raise before a provider request when this agenda may not spend more.
+
+    A NULL, 0, or negative budget disables the cap. A paused agenda whose
+    budget was raised back above spent plus in-flight grant usage is
+    reactivated, so increasing the budget is enough to continue.
+    """
+    row = _agenda_budget_row(agenda_id)
+    if row is None:
+        return
+    cap = budget_cap(row.get("token_budget"))
+    if cap is None:
+        return
+    charged = charged_tokens(agenda_id, int(row.get("token_spent") or 0))
+    if charged >= cap:
+        if str(row.get("status") or "") != PAUSED_BUDGET:
+            _set_status(agenda_id, PAUSED_BUDGET, commit=True)
+        raise AgendaBudgetExceededError(agenda_id, charged, cap)
+    if str(row.get("status") or "") == PAUSED_BUDGET:
+        _set_status(agenda_id, ACTIVE, commit=True)
+
+
+def _insert_ledger_row(
+    agenda_id: int,
+    operation: str,
+    tokens: int,
+    cost_usd: float | None,
+) -> bool:
+    from db import database as db
+
+    if not db.table_exists("agenda_token_ledger"):
+        return False
+    db.execute(
+        """
+        INSERT INTO agenda_token_ledger (agenda_id, operation, tokens, cost_usd)
+        VALUES (?, ?, ?, ?)
+        """,
+        (int(agenda_id), str(operation or "llm_call"), int(tokens or 0), cost_usd),
+    )
+    return True
+
+
+def record_usage(
+    agenda_id: int,
+    operation: str,
+    tokens: int,
+    cost_usd: float | None = None,
+) -> None:
+    """Record one legacy LLM call and bump token_spent.
+
+    Grant-backed calls must not use this. Their parent reservation already
+    adds actual tokens to token_spent when the grant settles.
+    """
+    from db import database as db
+
+    tokens = int(tokens or 0)
+    agenda_id = int(agenda_id)
+    try:
+        if not _insert_ledger_row(agenda_id, operation, tokens, cost_usd):
+            return
+        if db.table_exists("research_agendas"):
+            db.execute(
+                """
+                UPDATE research_agendas
+                SET token_spent = COALESCE(token_spent, 0) + ?,
+                    updated_at=CURRENT_TIMESTAMP
+                WHERE id=?
+                """,
+                (tokens, agenda_id),
+            )
+        _pause_if_exhausted(agenda_id, commit=False)
+        db.commit()
+    except Exception:
+        db.rollback()
+        raise
+
+
+def attribute_settled_grant_usage(
+    agenda_id: int,
+    operation: str,
+    tokens: int,
+    cost_usd: float | None = None,
+) -> None:
+    """Append a ledger row for one settled grant sub-reservation.
+
+    Does not touch token_spent. The caller owns the transaction.
+    """
+    if not _insert_ledger_row(agenda_id, operation, int(tokens or 0), cost_usd):
+        return
+    _pause_if_exhausted(int(agenda_id), commit=False)

agents/llm_client.py

@@ -1131,11 +1131,60 @@ def is_llm_transient_provider_error(exc: Exception) -> bool:
     return any(marker in msg for marker in markers)
 
 
+def _agenda_budget_scope() -> tuple[int, str, str] | None:
+    """Active agenda attribution, if this call is running inside one."""
+    from agents.agenda_budget import current_scope
+
+    return current_scope()
+
+
+def _enforce_scoped_budget(agenda_id: int | None = None) -> tuple[int, str, str] | None:
+    """Check the agenda cap before a provider request.
+
+    A scope for a different agenda is ignored so one worker cannot charge
+    or block another agenda's call.
+    """
+    scope = _agenda_budget_scope()
+    if scope is None:
+        return None
+    if agenda_id is not None and int(scope[0]) != int(agenda_id):
+        return None
+    from agents.agenda_budget import check_budget
+
+    check_budget(int(scope[0]))
+    return scope
+
+
+def _record_scoped_usage(
+    scope: tuple[int, str, str] | None,
+    tokens: int,
+    cost_usd: float | None,
+) -> None:
+    """Record a legacy, non-grant call. Failures keep the paid response."""
+    if scope is None or not tokens:
+        return
+    try:
+        from agents.agenda_budget import record_usage
+
+        record_usage(int(scope[0]), str(scope[1]), int(tokens), cost_usd=cost_usd)
+    except Exception as acct_err:  # noqa: BLE001
+        print(
+            f"[LLM] WARNING: agenda token accounting failed "
+            f"(agenda {scope[0]}): {acct_err}",
+            flush=True,
+        )
+
+
 def call_llm(system_prompt: str, user_prompt: str, temperature: float = 0.0,
              max_tokens: int = None) -> tuple[str, int]:
-    """Call LLM with automatic provider selection, rate limiting, and failover."""
+    """Call LLM with automatic provider selection, rate limiting, and failover.
+
+    Inside ``agenda_scope``, the agenda cap is checked before any provider
+    request and a successful response is recorded on the agenda ledger.
+    """
     max_tokens = max_tokens or LLM_MAX_OUTPUT_TOKENS
     _init_providers()
+    budget_scope = _enforce_scoped_budget()
 
     last_error = None
     tried = set()
@@ -1161,7 +1210,7 @@ def call_llm(system_prompt: str, user_prompt: str, temperature: float = 0.0,
 
             start = time.time()
             try:
-                text, tokens, cached_toks, input_toks, _cost_usd = _call_provider(
+                text, tokens, cached_toks, input_toks, cost_usd = _call_provider(
                     provider, system_prompt, user_prompt, max_tokens)
                 latency = time.time() - start
 
@@ -1182,6 +1231,7 @@ def call_llm(system_prompt: str, user_prompt: str, temperature: float = 0.0,
                     stats["cached_tokens"] += cached_toks
                     stats["input_tokens"] += input_toks
                 _release_provider(provider["name"])
+                _record_scoped_usage(budget_scope, tokens, cost_usd)
                 return text, tokens
 
             except Exception as e:
@@ -1309,14 +1359,15 @@ def call_llm_with_provider(
         provider = _providers[0]
 
     name = provider["name"]
+    budget_scope = _enforce_scoped_budget()
     limiter = _rate_limiters.get(name)
     if limiter:
         limiter.wait()
     with _provider_lock:
         _provider_stats[name]["in_flight"] = _provider_stats[name].get("in_flight", 0) + 1
     start = time.time()
     try:
-        text, tokens, cached_toks, input_toks, _cost_usd = _call_provider(
+        text, tokens, cached_toks, input_toks, cost_usd = _call_provider(
             provider, system_prompt, user_prompt, max_tokens
         )
         latency = time.time() - start
@@ -1329,6 +1380,7 @@ def call_llm_with_provider(
             stats["input_tokens"] += input_toks
         if not text or len(text.strip()) < 10:
             raise RuntimeError(f"{name} returned empty response")
+        _record_scoped_usage(budget_scope, tokens, cost_usd)
         return text, tokens, {
             "name": provider.get("name"),
             "model": provider.get("model"),
@@ -1396,6 +1448,10 @@ def call_llm_for_role(
     )
     if not grant_row:
         raise PermissionError("active scoped ResourceGrant is required for LLM role")
+    # Execution workers enter agenda_scope before this call. The grant still
+    # owns token_spent; this only refuses a provider request once spent plus
+    # in-flight grant usage has already reached the agenda cap.
+    _enforce_scoped_budget(int(agenda_id))
     # Trim the request to what this grant still has, rather than sending the
     # caller's own default and letting the sub-reservation refuse it. A small
     # short grant is the whole point of the bounded pilot path: without this a

meta_harness/grant_usage.py

@@ -235,6 +235,17 @@ def settle(
                 """,
                 (tokens_used, cost_usd, reservation_id, self.resource_grant_id),
             )
+            # Execution workers spend through this sub-reservation. The parent
+            # grant still owns token_spent; the compatibility ledger is the
+            # per-call record, written in this same transaction.
+            from agents.agenda_budget import attribute_settled_grant_usage
+
+            attribute_settled_grant_usage(
+                int(row["agenda_id"]),
+                str(row.get("operation") or "llm_call"),
+                int(tokens_used),
+                cost_usd,
+            )
             db.commit()
         except Exception:
             db.rollback()

orchestrator/auto_research.py

@@ -14,6 +14,7 @@
 import os
 import threading
 import time
+from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from pathlib import Path
@@ -3835,7 +3836,37 @@ def _launch_candidates_to_capacity() -> dict:
     }
 
 
+@contextmanager
+def _auto_research_budget_scope(insight: dict):
+    """Attribute this candidate's LLM calls to its agenda.
+
+    Discovery enters an agenda scope before it spends. Execution workers
+    dispatched after a selection is handed to auto-research used to call the
+    provider outside any scope, so those tokens never reached the agenda
+    ledger. The scope is attribution only; it does not reserve budget.
+    """
+    agenda_id = int(insight.get("agenda_id") or 0)
+    insight_id = int(insight.get("id") or 0)
+    if agenda_id <= 0 or insight_id <= 0:
+        yield
+        return
+    from agents.agenda_budget import agenda_scope
+
+    with agenda_scope(
+        agenda_id,
+        "auto_research_execution",
+        f"auto-research:{agenda_id}:{insight_id}",
+    ):
+        yield
+
+
 def _process_candidate(insight: dict) -> None:
+    """Run one auto-research candidate inside its agenda token scope."""
+    with _auto_research_budget_scope(insight):
+        _process_candidate_in_scope(insight)
+
+
+def _process_candidate_in_scope(insight: dict) -> None:
     insight_id = insight["id"]
     tier = insight.get("tier")
     initial_auto_stage = str(insight.get("auto_stage") or "").strip()