memorymaster 3.17.1__tar.gz → 3.19.0__tar.gz

{memorymaster-3.17.1/memorymaster.egg-info → memorymaster-3.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: memorymaster
-Version: 3.17.1
+Version: 3.19.0
 Summary: Production-grade memory reliability system for AI coding agents. Lifecycle-managed claims with citations, conflict detection, steward governance, and MCP integration.
 Author: wolverin0
 License: MIT

{memorymaster-3.17.1 → memorymaster-3.19.0}/memorymaster/config.py

@@ -67,6 +67,17 @@ MEMORYMASTER_RRF_TIEBREAKER_THRESHOLD
     Maximum adjacent score gap considered a near tie.
     Default: ``0.01``
 
+MEMORYMASTER_RETRIEVAL_PROFILE_<TYPE>
+    Per-question-type retrieval weight override. ``<TYPE>`` is the
+    upper-snake-case form of the question type (LongMemEval-S labels:
+    SINGLE_SESSION_USER, SINGLE_SESSION_ASSISTANT, SINGLE_SESSION_PREFERENCE,
+    TEMPORAL_REASONING, KNOWLEDGE_UPDATE, MULTI_SESSION). Value format
+    matches ``MEMORYMASTER_RETRIEVAL_WEIGHTS``: comma-separated floats
+    ``lex,conf,fresh,vec``. When set, ``Config.retrieval_profile(qtype)``
+    returns the profile and ``retrieval._compute_claim_score`` consumes it
+    instead of the default weights. Profiles are opt-in; missing types fall
+    back to the global ``cfg.retrieval_weights``.
+
 MEMORYMASTER_CONFIG_FILE
     Path to a JSON config file. Keys match attribute names on ``Config``.
 """
@@ -218,6 +229,15 @@ class Config:
         default_factory=lambda: dict(INITIAL_CONFIDENCE_BY_TYPE)
     )
 
+    # --- Per-question-type retrieval weight profiles (S3, opt-in) ---
+    # Maps the canonical question_type slug (lowercase, hyphens preserved,
+    # e.g. "single-session-preference") to a (W_LEX, W_CONF, W_FRESH, W_VEC)
+    # tuple. Populated by env vars MEMORYMASTER_RETRIEVAL_PROFILE_<TYPE>.
+    # Empty by default — retrieval falls back to ``retrieval_weights``.
+    retrieval_profiles: Dict[str, tuple[float, float, float, float]] = field(
+        default_factory=dict
+    )
+
     # --- Derived convenience dicts ---
 
     @property
@@ -270,6 +290,20 @@ class Config:
             self.lexical_weight_prefix,
         )
 
+    def retrieval_profile(
+        self, question_type: str | None
+    ) -> tuple[float, float, float, float] | None:
+        """Per-question-type weight override, or None to fall back to default.
+
+        Accepts the canonical bench label (e.g. ``single-session-preference``)
+        or the classifier output (e.g. ``preference``) — both lookup forms are
+        normalized to lowercase before lookup. Returning None signals the
+        caller to use ``retrieval_weights`` as usual.
+        """
+        if not question_type:
+            return None
+        return self.retrieval_profiles.get(question_type.strip().lower())
+
 
 # ---------------------------------------------------------------------------
 # Module-level singleton
@@ -363,6 +397,7 @@ def load_config(config_path: str | Path | None = None) -> Config:
     _apply_env_bool(overrides, "MEMORYMASTER_LLM_RERANK", "llm_rerank")
     _apply_env_bool(overrides, "MEMORYMASTER_RRF_TIEBREAKER", "rrf_tiebreaker_enabled")
     _apply_env_float(overrides, "MEMORYMASTER_RRF_TIEBREAKER_THRESHOLD", "rrf_tiebreaker_threshold")
+    _apply_env_retrieval_profiles(overrides)
 
     # Filter to only valid Config fields
     valid_fields = {f.name for f in Config.__dataclass_fields__.values()}
@@ -404,3 +439,29 @@ def _apply_env_bool(overrides: dict[str, object], env_var: str, key: str) -> Non
     if not raw:
         return
     overrides[key] = raw in {"1", "true", "yes", "on"}
+
+
+_PROFILE_PREFIX = "MEMORYMASTER_RETRIEVAL_PROFILE_"
+
+
+def _apply_env_retrieval_profiles(overrides: dict[str, object]) -> None:
+    """Collect per-type profile overrides from env into a single dict.
+
+    Scans every env var starting with ``MEMORYMASTER_RETRIEVAL_PROFILE_``.
+    The suffix becomes the question_type slug (lowercased, underscores
+    converted to hyphens to match the LongMemEval-S canonical labels like
+    ``single-session-preference``).
+    """
+    profiles: dict[str, tuple[float, float, float, float]] = {}
+    for key, raw in os.environ.items():
+        if not key.startswith(_PROFILE_PREFIX):
+            continue
+        raw = raw.strip()
+        if not raw:
+            continue
+        suffix = key[len(_PROFILE_PREFIX):]
+        qtype = suffix.lower().replace("_", "-")
+        values = _parse_floats(raw, 4)
+        profiles[qtype] = (values[0], values[1], values[2], values[3])
+    if profiles:
+        overrides["retrieval_profiles"] = profiles

{memorymaster-3.17.1 → memorymaster-3.19.0}/memorymaster/dashboard.py

@@ -19,6 +19,7 @@ import urllib.error
 import urllib.request
 from urllib.parse import parse_qs, urlparse
 
+from memorymaster import dashboard_auth
 from memorymaster.config import get_config
 from memorymaster.review import build_review_queue, queue_to_dicts
 from memorymaster.security import is_sensitive_claim
@@ -626,6 +627,9 @@ class DashboardHTTPServer(ThreadingHTTPServer):
         self.db_target = str(db_target) if db_target is not None else "memorymaster.db"
         self.workspace_root = Path(workspace_root) if workspace_root is not None else Path.cwd()
         self._operator_proc: subprocess.Popen[str] | None = None
+        # Pin the configured host:port string so the CSRF check can compare it
+        # against the request's Origin/Referer header in authenticate POST routes.
+        self.configured_host_port = f"{server_address[0]}:{server_address[1]}"
         super().__init__(server_address, DashboardRequestHandler)
 
     def operator_status(self) -> dict[str, Any]:
@@ -701,10 +705,45 @@ class DashboardRequestHandler(BaseHTTPRequestHandler):
     def _server(self) -> DashboardHTTPServer:
         return self.server  # type: ignore[return-value]
 
+    def _enforce_auth(self, *, method: str, route: str) -> bool:
+        """Run the v3.19.0-H2 auth + role + CSRF gates. Returns True on pass.
+
+        On failure, the JSON error response has already been written; the
+        caller must return immediately without dispatching the route.
+        Health endpoints (/health, /healthz, /readyz) are intentionally
+        exempt so external monitors can probe without credentials.
+        """
+        if route in {"/health", "/healthz", "/readyz"}:
+            return True
+
+        decision = dashboard_auth.authenticate(self.headers)
+        decision = dashboard_auth.authorize(decision, method=method, route=route)
+        if not decision.ok:
+            self._write_json(
+                {"ok": False, "error": decision.reason},
+                status=decision.status,
+            )
+            return False
+
+        if method.upper() == "POST":
+            csrf = dashboard_auth.check_csrf(
+                self.headers,
+                configured_host_port=getattr(self._server, "configured_host_port", None),
+            )
+            if not csrf.ok:
+                self._write_json(
+                    {"ok": False, "error": csrf.reason},
+                    status=csrf.status,
+                )
+                return False
+        return True
+
     def do_GET(self) -> None:  # noqa: N802
         parsed = urlparse(self.path)
         route = parsed.path
         try:
+            if not self._enforce_auth(method="GET", route=route):
+                return
             if _route_get_request(self, route, parsed.query):
                 return
             self._write_json({"ok": False, "error": "Not found"}, status=HTTPStatus.NOT_FOUND)
@@ -717,6 +756,8 @@ class DashboardRequestHandler(BaseHTTPRequestHandler):
         parsed = urlparse(self.path)
         route = parsed.path
         try:
+            if not self._enforce_auth(method="POST", route=route):
+                return
             payload = self._read_json_body()
             if route == "/api/triage/action":
                 self._handle_triage_action(payload)
@@ -1443,6 +1484,12 @@ def create_dashboard_server(
     port: int = 8765,
     operator_log_jsonl: str | Path = "artifacts/operator/operator_events.jsonl",
 ) -> DashboardHTTPServer:
+    # v3.19.0-H2: refuse non-loopback bind without an auth secret unless the
+    # operator explicitly opts in. Raises dashboard_auth.BindUnsafeError on
+    # an unsafe configuration; callers should let it propagate so the
+    # mistake is visible at startup instead of silently exposed.
+    dashboard_auth.check_bind_safety(host)
+
     if service is None:
         if db_target is None:
             db_target = "memorymaster.db"

memorymaster-3.19.0/memorymaster/dashboard_auth.py

@@ -0,0 +1,196 @@
+"""HTTP-layer authentication, CSRF, and bind-safety for the dashboard (v3.19.0-H2).
+
+The dashboard's prior posture was "trusted local-only" — zero HTTP auth,
+default loopback bind. Anyone who could reach the port had full read +
+operator-control. This module adds opt-in token auth with viewer/operator
+role separation, CSRF for browser POSTs, and refusal to bind non-loopback
+hosts without an explicit secret.
+
+Backwards compatibility: when no auth secrets are configured (both
+``MEMORYMASTER_DASHBOARD_TOKEN_VIEWER`` and ``MEMORYMASTER_DASHBOARD_TOKEN_OPERATOR``
+are empty), the dashboard runs in legacy mode — no auth enforced. Bind
+safety still applies: legacy mode refuses non-loopback bind unless the
+operator explicitly opts in with ``MEMORYMASTER_DASHBOARD_UNSAFE_BIND=1``.
+
+Env vars:
+    MEMORYMASTER_DASHBOARD_TOKEN_VIEWER     — bearer token granting read-only access
+    MEMORYMASTER_DASHBOARD_TOKEN_OPERATOR   — bearer token granting full mutating access
+    MEMORYMASTER_DASHBOARD_UNSAFE_BIND      — set to 1 to allow non-loopback bind
+                                              without an auth secret (logs WARNING)
+"""
+from __future__ import annotations
+
+import hmac
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+
+logger = logging.getLogger(__name__)
+
+
+class DashboardRole(str, Enum):
+    VIEWER = "viewer"
+    OPERATOR = "operator"
+
+
+# Routes that require operator role regardless of HTTP method. POST routes
+# always require operator role separately. This list captures GET routes that
+# are mutating-adjacent (start/stop operator control, stream that may carry
+# sensitive operator events).
+OPERATOR_ONLY_GET_ROUTES: frozenset[str] = frozenset({
+    "/api/operator/control",  # POST in practice but pin it here for completeness
+    "/api/operator/stream",   # SSE — operator-only since it streams operator events
+})
+
+# POST routes — all require operator.
+POST_OPERATOR_ROUTES: frozenset[str] = frozenset({
+    "/api/triage/action",
+    "/api/operator/control",
+    "/api/action-proposals/status",
+})
+
+
+@dataclass(frozen=True)
+class AuthDecision:
+    """Outcome of an authentication or authorization check."""
+
+    ok: bool
+    role: DashboardRole | None = None
+    reason: str = ""
+    status: int = 200
+
+
+class BindUnsafeError(RuntimeError):
+    """Raised when the dashboard refuses to bind a non-loopback host without auth."""
+
+
+def _env_token(name: str) -> str:
+    return os.environ.get(name, "").strip()
+
+
+def legacy_mode() -> bool:
+    """True when no auth secrets are configured (back-compat path).
+
+    Legacy mode preserves the pre-v3.19 behaviour for loopback bind — no
+    auth check, no CSRF. Bind safety still applies regardless.
+    """
+    return not (
+        _env_token("MEMORYMASTER_DASHBOARD_TOKEN_VIEWER")
+        or _env_token("MEMORYMASTER_DASHBOARD_TOKEN_OPERATOR")
+    )
+
+
+def _extract_bearer(headers) -> str:
+    raw = headers.get("Authorization", "") if headers else ""
+    if not raw or not raw.lower().startswith("bearer "):
+        return ""
+    return raw[7:].strip()
+
+
+def authenticate(headers) -> AuthDecision:
+    """Map an incoming request's Authorization header to a role.
+
+    Behaviour:
+    - Legacy mode (no tokens set): return ok with role=operator (no enforcement).
+    - Missing token: 401 (``missing_token``).
+    - Token matches operator: ok with role=operator.
+    - Token matches viewer: ok with role=viewer.
+    - Token unrecognized: 401 (``invalid_token``).
+
+    Constant-time comparison (``hmac.compare_digest``) prevents timing attacks.
+    """
+    if legacy_mode():
+        return AuthDecision(True, role=DashboardRole.OPERATOR)
+
+    token = _extract_bearer(headers)
+    if not token:
+        return AuthDecision(False, reason="missing_token", status=401)
+
+    op_token = _env_token("MEMORYMASTER_DASHBOARD_TOKEN_OPERATOR")
+    vw_token = _env_token("MEMORYMASTER_DASHBOARD_TOKEN_VIEWER")
+
+    if op_token and hmac.compare_digest(token, op_token):
+        return AuthDecision(True, role=DashboardRole.OPERATOR)
+    if vw_token and hmac.compare_digest(token, vw_token):
+        return AuthDecision(True, role=DashboardRole.VIEWER)
+    return AuthDecision(False, reason="invalid_token", status=401)
+
+
+def authorize(decision: AuthDecision, *, method: str, route: str) -> AuthDecision:
+    """Refine an authentication decision against the role-vs-route policy.
+
+    Returns 403 (``role_required_operator``) for routes where viewer is
+    insufficient. Passes ``decision`` through unchanged on success.
+    """
+    if not decision.ok:
+        return decision
+
+    method_u = method.upper()
+    operator_required = (
+        method_u == "POST"
+        or route in OPERATOR_ONLY_GET_ROUTES
+        or route in POST_OPERATOR_ROUTES
+    )
+    if operator_required and decision.role != DashboardRole.OPERATOR:
+        return AuthDecision(
+            False,
+            role=decision.role,
+            reason="role_required_operator",
+            status=403,
+        )
+    return decision
+
+
+def check_csrf(headers, *, configured_host_port: str | None) -> AuthDecision:
+    """Validate Origin/Referer for browser-originated POSTs.
+
+    Non-browser clients (curl, scripts, MCP-style integrations) typically
+    omit ``Origin`` — those requests pass through unchallenged. Browsers
+    always set ``Origin``; when present, it must contain the configured
+    host:port string. Returns 403 (``csrf_origin_mismatch``) on mismatch.
+
+    Legacy mode skips CSRF entirely.
+    """
+    if legacy_mode():
+        return AuthDecision(True)
+
+    origin = ""
+    if headers:
+        origin = headers.get("Origin", "") or headers.get("Referer", "") or ""
+    if not origin:
+        return AuthDecision(True)  # non-browser caller
+    if configured_host_port and configured_host_port not in origin:
+        return AuthDecision(False, reason="csrf_origin_mismatch", status=403)
+    return AuthDecision(True)
+
+
+def check_bind_safety(host: str) -> None:
+    """Refuse non-loopback bind without an auth secret or explicit opt-in.
+
+    Raises ``BindUnsafeError`` if all of: (a) host is non-loopback,
+    (b) no auth tokens configured, (c) ``MEMORYMASTER_DASHBOARD_UNSAFE_BIND``
+    not set. Otherwise returns silently. Logs a WARNING for the unsafe opt-in
+    case so operators see they're running exposed.
+    """
+    loopback_hosts = {"127.0.0.1", "::1", "localhost", ""}
+    if host in loopback_hosts:
+        return
+    if not legacy_mode():
+        return  # token-based auth is enforced; non-loopback bind is acceptable
+
+    unsafe_raw = os.environ.get("MEMORYMASTER_DASHBOARD_UNSAFE_BIND", "").strip().lower()
+    if unsafe_raw in {"1", "true", "yes", "on"}:
+        logger.warning(
+            "dashboard binding to non-loopback host '%s' with no auth secret "
+            "and MEMORYMASTER_DASHBOARD_UNSAFE_BIND=1 — running exposed",
+            host,
+        )
+        return
+
+    raise BindUnsafeError(
+        f"Refusing to bind dashboard to non-loopback host '{host}' without auth secret. "
+        f"Set MEMORYMASTER_DASHBOARD_TOKEN_OPERATOR (and optionally "
+        f"MEMORYMASTER_DASHBOARD_TOKEN_VIEWER), or explicitly opt-in with "
+        f"MEMORYMASTER_DASHBOARD_UNSAFE_BIND=1."
+    )

{memorymaster-3.17.1 → memorymaster-3.19.0}/memorymaster/jobs/daydream_ingest.py

@@ -33,7 +33,51 @@ def ingest_insights(
     scope: str = "user",
     dry_run: bool = False,
 ) -> dict:
-    """Ingest accepted daydream insights as candidate hypothesis claims."""
+    """Ingest accepted daydream insights as candidate hypothesis claims.
+
+    Honours per-cycle LLM budget caps from ``llm_budget``. Today this path
+    does not itself call ``call_llm`` (insights arrive pre-scored), but the
+    wrapper is in place so downstream changes (re-scoring, paraphrase
+    detection, etc.) inherit the same abort semantics as ``service.run_cycle``
+    and ``wiki_engine.absorb``. If a parent scope is already open, defers
+    to it instead of opening a nested one.
+    """
+    from memorymaster import llm_budget
+
+    if llm_budget.get_current() is not None:
+        return _ingest_insights_impl(
+            service, insights_dir, min_score=min_score, scope=scope, dry_run=dry_run
+        )
+
+    with llm_budget.cycle_scope() as budget:
+        try:
+            result = _ingest_insights_impl(
+                service, insights_dir, min_score=min_score, scope=scope, dry_run=dry_run
+            )
+        except llm_budget.LLMBudgetExceeded as exc:
+            result = {
+                "ingested": 0,
+                "skipped": 0,
+                "errors": [
+                    f"daydream ingest aborted by llm budget: reason={exc.reason}"
+                ],
+                "aborted": True,
+                "aborted_reason": exc.reason,
+                "aborted_provider": exc.provider,
+            }
+        result["budget"] = budget.snapshot()
+        return result
+
+
+def _ingest_insights_impl(
+    service: MemoryService,
+    insights_dir: Path,
+    *,
+    min_score: float = 7.0,
+    scope: str = "user",
+    dry_run: bool = False,
+) -> dict:
+    """Original ingest_insights implementation, called inside a budget scope."""
     root = Path(insights_dir)
     result: dict[str, Any] = {"ingested": 0, "skipped": 0, "errors": []}
     if not root.exists() or not root.is_dir():

memorymaster-3.19.0/memorymaster/llm_budget.py

@@ -0,0 +1,190 @@
+"""Per-cycle LLM budget caps with reason-coded hard stops.
+
+A cycle scope tracks LLM call count, estimated tokens, and per-provider
+failures. When any cap is exceeded, the next ``llm_provider.call_llm``
+raises ``LLMBudgetExceeded`` with a reason code so callers can record
+the abort visibly instead of silently overspending.
+
+Caps are read from env vars at scope-entry time (any change requires a
+new scope to take effect):
+
+- ``MEMORYMASTER_MAX_LLM_CALLS_PER_CYCLE``       — hard cap on total calls
+- ``MEMORYMASTER_MAX_TOKENS_PER_CYCLE``          — hard cap on summed estimated tokens
+- ``MEMORYMASTER_MAX_PROVIDER_FAILURES_PER_CYCLE`` — per-provider failure ceiling
+                                                    (also acts as circuit breaker)
+
+A value of ``0`` (default) means "unlimited" for that axis — preserves
+backwards compatibility when env vars are unset.
+
+Usage::
+
+    from memorymaster import llm_budget
+
+    with llm_budget.cycle_scope() as budget:
+        ...
+        try:
+            response = llm_provider.call_llm(prompt, text)
+        except llm_budget.LLMBudgetExceeded as exc:
+            # exc.reason in {"calls_exhausted", "tokens_exhausted",
+            #                "provider_failures_exhausted"}
+            # exc.provider is set only for provider-failures reason.
+            ...
+        ...
+    snapshot = budget.snapshot()  # totals after the scope
+"""
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Iterator
+
+
+# ---------------------------------------------------------------------------
+# Exception + dataclass
+# ---------------------------------------------------------------------------
+
+
+class LLMBudgetExceeded(Exception):
+    """Raised by ``call_llm`` when a per-cycle budget cap is hit.
+
+    Attributes:
+        reason: one of ``"calls_exhausted"``, ``"tokens_exhausted"``,
+            ``"provider_failures_exhausted"``.
+        provider: provider name (set only for the failures reason).
+    """
+
+    def __init__(self, reason: str, provider: str | None = None) -> None:
+        self.reason = reason
+        self.provider = provider
+        suffix = f" provider={provider}" if provider else ""
+        super().__init__(f"llm budget exceeded: reason={reason}{suffix}")
+
+
+@dataclass
+class CycleBudget:
+    """Per-cycle counters and limits. Lives for the duration of one scope.
+
+    Limits of 0 mean unlimited. ``aborted_reason`` is set the first time
+    a cap is hit (subsequent overruns don't overwrite the original reason).
+    """
+
+    max_calls: int = 0
+    max_tokens: int = 0
+    max_provider_failures: int = 0
+    calls: int = 0
+    tokens: int = 0
+    provider_failures: dict[str, int] = field(default_factory=dict)
+    aborted_reason: str | None = None
+    aborted_provider: str | None = None
+
+    def snapshot(self) -> dict[str, object]:
+        return {
+            "max_calls": self.max_calls,
+            "max_tokens": self.max_tokens,
+            "max_provider_failures": self.max_provider_failures,
+            "calls": self.calls,
+            "tokens": self.tokens,
+            "provider_failures": dict(self.provider_failures),
+            "aborted_reason": self.aborted_reason,
+            "aborted_provider": self.aborted_provider,
+        }
+
+
+# ---------------------------------------------------------------------------
+# Context variable
+# ---------------------------------------------------------------------------
+
+
+_current: ContextVar[CycleBudget | None] = ContextVar(
+    "memorymaster_llm_cycle_budget", default=None
+)
+
+
+def _read_int_env(name: str) -> int:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return 0
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        return 0
+
+
+def _new_from_env() -> CycleBudget:
+    return CycleBudget(
+        max_calls=_read_int_env("MEMORYMASTER_MAX_LLM_CALLS_PER_CYCLE"),
+        max_tokens=_read_int_env("MEMORYMASTER_MAX_TOKENS_PER_CYCLE"),
+        max_provider_failures=_read_int_env("MEMORYMASTER_MAX_PROVIDER_FAILURES_PER_CYCLE"),
+    )
+
+
+@contextmanager
+def cycle_scope() -> Iterator[CycleBudget]:
+    """Open a new per-cycle budget scope. Yields the tracker; auto-cleans on exit."""
+    budget = _new_from_env()
+    token = _current.set(budget)
+    try:
+        yield budget
+    finally:
+        _current.reset(token)
+
+
+def get_current() -> CycleBudget | None:
+    """Return the active cycle budget, or None if no scope is open."""
+    return _current.get()
+
+
+# ---------------------------------------------------------------------------
+# Enforcement helpers — invoked from llm_provider.call_llm
+# ---------------------------------------------------------------------------
+
+
+def _abort(budget: CycleBudget, reason: str, provider: str | None = None) -> None:
+    """Record the first abort reason and raise."""
+    if budget.aborted_reason is None:
+        budget.aborted_reason = reason
+        budget.aborted_provider = provider
+    raise LLMBudgetExceeded(reason, provider)
+
+
+def estimate_tokens(*parts: str) -> int:
+    """Rough char/4 estimator. Sufficient for cap accounting; not for billing."""
+    total = sum(len(p) for p in parts if p)
+    return (total + 3) // 4
+
+
+def check_before_call(provider: str) -> None:
+    """Raise if the calls cap is already at its ceiling, or this provider's
+    failure breaker is open. Called before contacting the provider."""
+    budget = get_current()
+    if budget is None:
+        return
+    if budget.max_calls and budget.calls >= budget.max_calls:
+        _abort(budget, "calls_exhausted")
+    if budget.max_provider_failures:
+        if budget.provider_failures.get(provider, 0) >= budget.max_provider_failures:
+            _abort(budget, "provider_failures_exhausted", provider)
+
+
+def record_call(provider: str, *, tokens: int = 0) -> None:
+    """Record a successful (or attempted) call. Raises if tokens cap is hit."""
+    budget = get_current()
+    if budget is None:
+        return
+    budget.calls += 1
+    budget.tokens += max(0, tokens)
+    if budget.max_tokens and budget.tokens >= budget.max_tokens:
+        _abort(budget, "tokens_exhausted")
+
+
+def record_failure(provider: str) -> None:
+    """Increment per-provider failure counter. Raises if breaker hits limit."""
+    budget = get_current()
+    if budget is None:
+        return
+    new_count = budget.provider_failures.get(provider, 0) + 1
+    budget.provider_failures[provider] = new_count
+    if budget.max_provider_failures and new_count >= budget.max_provider_failures:
+        _abort(budget, "provider_failures_exhausted", provider)

{memorymaster-3.17.1 → memorymaster-3.19.0}/memorymaster/llm_provider.py

@@ -17,6 +17,8 @@ import urllib.request
 import urllib.error
 from typing import Any
 
+from memorymaster import llm_budget
+
 
 def _env(key: str, default: str = "") -> str:
     return os.environ.get(key, default)
@@ -489,13 +491,28 @@ def call_llm(prompt: str, text: str) -> str:
     if not primary_fn:
         return ""
 
+    # Per-cycle budget gate. If a cycle_scope() is active and the calls cap
+    # is already hit (or this provider's failure breaker is open), raise
+    # LLMBudgetExceeded so the caller can record the abort visibly instead
+    # of silently overspending. When no scope is active, this is a no-op
+    # — preserves backwards-compat for callers outside run_cycle/wiki/daydream.
+    llm_budget.check_before_call(primary_name)
+
     primary_response = primary_fn(prompt, text)
+    llm_budget.record_call(
+        primary_name,
+        tokens=llm_budget.estimate_tokens(prompt, text, primary_response),
+    )
 
     # Happy path: primary returned something that doesn't look like a quota error.
     if primary_response and not _looks_like_quota_error(primary_response):
         _FALLBACK_STATS["primary_ok"] += 1
         return primary_response
 
+    # Treat empty or quota-shaped response as a provider failure for breaker purposes.
+    # record_failure may raise if the per-provider failure cap is hit.
+    llm_budget.record_failure(primary_name)
+
     fallback_name = _env("MEMORYMASTER_LLM_FALLBACK_PROVIDER", "").lower()
     if not fallback_name:
         return primary_response
@@ -512,6 +529,9 @@ def call_llm(prompt: str, text: str) -> str:
     log.info("llm_fallback_fired primary=%s reason=%s", primary_name, reason)
     _FALLBACK_STATS["fired"] += 1
 
+    # Budget gate for the fallback provider too.
+    llm_budget.check_before_call(fallback_name)
+
     # Swap MEMORYMASTER_LLM_MODEL to fallback model for the duration of the call.
     fallback_model = _env("MEMORYMASTER_LLM_FALLBACK_MODEL", "")
     saved_model = os.environ.get("MEMORYMASTER_LLM_MODEL")
@@ -523,6 +543,10 @@ def call_llm(prompt: str, text: str) -> str:
             # own default, not the primary's model (which may be Gemini-specific).
             del os.environ["MEMORYMASTER_LLM_MODEL"]
         fallback_response = fallback_fn(prompt, text)
+        llm_budget.record_call(
+            fallback_name,
+            tokens=llm_budget.estimate_tokens(prompt, text, fallback_response),
+        )
     finally:
         if saved_model is None:
             os.environ.pop("MEMORYMASTER_LLM_MODEL", None)
@@ -533,6 +557,8 @@ def call_llm(prompt: str, text: str) -> str:
         return fallback_response
 
     # Both failed — match legacy contract, return primary's (possibly empty) response.
+    # Record fallback failure for breaker purposes; may raise if cap is hit.
+    llm_budget.record_failure(fallback_name)
     return primary_response