blisspixel
diff --git a/‎deepr/api/app.py‎
Lines changed: 18 additions & 3 deletions b/‎deepr/api/app.py‎
Lines changed: 18 additions & 3 deletions
diff --git a/‎deepr/experts/beliefs.py‎
Lines changed: 27 additions & 7 deletions b/‎deepr/experts/beliefs.py‎
Lines changed: 27 additions & 7 deletions
diff --git a/‎deepr/experts/chat.py‎
Lines changed: 114 additions & 23 deletions b/‎deepr/experts/chat.py‎
Lines changed: 114 additions & 23 deletions
diff --git a/‎deepr/experts/confidence_calibration.py‎
Lines changed: 3 additions & 2 deletions b/‎deepr/experts/confidence_calibration.py‎
Lines changed: 3 additions & 2 deletions
@@ -115,7 +115,12 @@ def _check_auth():
 ```
 
 ## Authentication
-Currently, the API does not require authentication. Rate limiting is based on client IP address.
+Bearer token via ``Authorization: Bearer <token>`` is required when the server is
+started with ``DEEPR_API_TOKEN`` set, which is the recommended configuration for
+any non-loopback bind. Local-development binds to 127.0.0.1 may run without a
+token; the server refuses to start on a non-loopback host without either
+``DEEPR_API_TOKEN`` or ``DEEPR_ALLOW_PUBLIC_BIND=1``. Rate limiting is applied
+per client IP on top of authentication.
         """,
         "version": "1.0.0",
         "contact": {"name": "Deepr Support", "url": "https://github.com/deepr-ai/deepr"},
@@ -621,8 +626,18 @@ def submit_job():
                     },
                 }
             ), 429
-    except Exception as _e:  # pragma: no cover - defensive; cost paths are well-tested
-        logger.warning("Cost guard skipped due to internal error: %s", _e)
+    except ImportError as _e:
+        # CostController/Estimator missing is a deployment problem, not a
+        # situation we should silently fall through. Fail-closed so the
+        # cost-gate cannot be bypassed by a broken dependency.
+        logger.error("Cost guard unavailable (ImportError): %s", _e)
+        return jsonify({"error": "Cost controller unavailable"}), 503
+    except (ValueError, TypeError) as _e:
+        # Estimator-input problems are the caller's; everything else
+        # should propagate and surface as a 500 instead of bypassing the
+        # gate the way the previous broad except did.
+        logger.warning("Cost guard input rejected: %s", _e)
+        return jsonify({"error": f"Invalid cost-guard input: {_e}"}), 400
 
     # Create job
     job_id = str(uuid.uuid4())
 
@@ -21,13 +21,16 @@
 """
 
 import json
+import logging
 import math
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from deepr.core.contracts import Claim
 
@@ -629,16 +632,32 @@ def _save(self):
             ],
         }
 
-        with open(self.storage_path, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+        from deepr.utils.atomic_io import atomic_write_json
+
+        atomic_write_json(self.storage_path, data)
 
     def _load(self):
-        """Load beliefs from disk."""
+        """Load beliefs from disk.
+
+        Catches corrupt/oversized files and starts fresh rather than
+        crashing the expert load entirely. A 50 MB ceiling guards
+        against poisoned belief files from corpus imports.
+        """
         if not self.storage_path.exists():
             return
 
-        with open(self.storage_path, encoding="utf-8") as f:
-            data = json.load(f)
+        try:
+            if self.storage_path.stat().st_size > 50 * 1024 * 1024:
+                logger.error(
+                    "Belief store at %s exceeds 50 MB; refusing to load. Inspect and reduce manually.",
+                    self.storage_path,
+                )
+                return
+            with open(self.storage_path, encoding="utf-8") as f:
+                data = json.load(f)
+        except (json.JSONDecodeError, OSError) as exc:
+            logger.error("Failed to load beliefs from %s: %s. Starting fresh.", self.storage_path, exc)
+            return
 
         self.beliefs = {bid: Belief.from_dict(bdata) for bid, bdata in data.get("beliefs", {}).items()}
 
@@ -891,8 +910,9 @@ def _save(self):
             "contributors": {bid: list(experts) for bid, experts in self.contributors.items()},
         }
 
-        with open(self.storage_path, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+        from deepr.utils.atomic_io import atomic_write_json
+
+        atomic_write_json(self.storage_path, data)
 
     def _load(self):
         """Load shared beliefs from disk."""
 
@@ -40,10 +40,11 @@
 def _chat_token_cost(usage: Any, model_name: str) -> float:
     """Compute chat-completion cost from token usage using the model registry.
 
-    Replaces the previous hard-coded GPT-5 rates: when the chat session
-    routes to gpt-5.2 the registry rates ($1.75/$14 per 1M) are applied
-    instead of the GPT-5 rates ($1.25/$10 per 1M), so cost_accumulated and
-    budget_remaining no longer under-count gpt-5.2 spend by ~28.6%.
+    Uses ``get_token_pricing`` so any model in the registry (gpt-5.2 at
+    $1.75/$14, etc.) is priced correctly. When ``usage`` exposes
+    ``prompt_tokens_details.cached_tokens`` (OpenAI caching), the cached
+    portion is billed at 50% — without this discount users hit their
+    session budget earlier than necessary on cache-hit-heavy workloads.
     """
     if not usage:
         return 0.0
@@ -58,7 +59,16 @@ def _chat_token_cost(usage: Any, model_name: str) -> float:
         output_price = _DEFAULT_CHAT_OUTPUT_PRICE_PER_1M
     prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
     completion_tokens = getattr(usage, "completion_tokens", 0) or 0
-    return (prompt_tokens / 1_000_000) * input_price + (completion_tokens / 1_000_000) * output_price
+
+    cached_tokens = 0
+    details = getattr(usage, "prompt_tokens_details", None)
+    if details is not None:
+        cached_tokens = getattr(details, "cached_tokens", 0) or 0
+    uncached_input = max(prompt_tokens - cached_tokens, 0)
+
+    input_cost = (uncached_input / 1_000_000) * input_price + (cached_tokens / 1_000_000) * input_price * 0.5
+    output_cost = (completion_tokens / 1_000_000) * output_price
+    return input_cost + output_cost
 
 
 class ExpertChatSession:
@@ -125,10 +135,19 @@ def __init__(
         from deepr.experts.cost_safety import get_cost_safety_manager
 
         self.cost_safety = get_cost_safety_manager()
+        # Sanitize identifiers feeding into session_id — agent_identity
+        # is caller-supplied (downstream MCP / A2A may forward arbitrary
+        # strings) and the session_id ends up in cost-safety keys and
+        # potentially thought-stream log filenames.
+        import re as _re
+
+        _safe_re = _re.compile(r"[^\w\-]+")
+        _expert_part = _safe_re.sub("_", str(expert.name))[:64]
         if self.agent_identity:
-            self.session_id = f"chat_{expert.name}_{self.agent_identity.agent_id}"
+            _agent_part = _safe_re.sub("_", str(self.agent_identity.agent_id))[:64]
+            self.session_id = f"chat_{_expert_part}_{_agent_part}"
         else:
-            self.session_id = f"chat_{expert.name}_{uuid.uuid4().hex[:8]}"
+            self.session_id = f"chat_{_expert_part}_{uuid.uuid4().hex[:8]}"
         self.cost_session = self.cost_safety.create_session(
             session_id=self.session_id, session_type="chat", budget_limit=self.budget
         )
@@ -610,6 +629,21 @@ async def _quick_lookup(self, query: str) -> dict:
         Returns:
             Dict with answer and sources
         """
+        # Pre-flight budget check. The previous implementation called the
+        # model first and accounted afterwards; that can blow past a tight
+        # session budget on a single call. Estimate via worst-case upper
+        # bound (4K input + 2K output at gpt-5.2 rates ≈ $0.035) and let
+        # the cost-safety layer veto if there's no room.
+        estimated_cost = 0.05
+        allowed, reason, _ = self.cost_safety.check_operation(
+            session_id=self.session_id,
+            operation_type="quick_lookup",
+            estimated_cost=estimated_cost,
+            require_confirmation=False,
+        )
+        if not allowed:
+            return {"error": f"Quick lookup blocked: {reason}", "mode": "quick_lookup_gpt52", "status": "blocked"}
+
         try:
             # Use GPT-5.2 with low reasoning effort for knowledge lookups
             response = await self.client.chat.completions.create(
@@ -626,14 +660,19 @@ async def _quick_lookup(self, query: str) -> dict:
 
             answer = response.choices[0].message.content or ""
 
-            # Track cost (GPT-5.2: $1.75 input, $14 output per 1M tokens)
-            if response.usage:
-                input_cost = (response.usage.prompt_tokens / 1_000_000) * 1.75
-                output_cost = (response.usage.completion_tokens / 1_000_000) * 14.00
-                cost = input_cost + output_cost
-                self.cost_accumulated += cost
-            else:
-                cost = 0.01  # Estimate ~5-10 cents for typical query
+            # Token-priced via the registry so future rate changes for
+            # gpt-5.2 don't silently drift from the budget bookkeeping.
+            cost = _chat_token_cost(response.usage, "gpt-5.2") if response.usage else 0.01
+            self.cost_accumulated += cost
+            self.cost_safety.record_cost(
+                session_id=self.session_id,
+                operation_type="quick_lookup",
+                actual_cost=cost,
+                provider="openai",
+                model="gpt-5.2",
+                tokens_input=getattr(response.usage, "prompt_tokens", 0) if response.usage else 0,
+                tokens_output=getattr(response.usage, "completion_tokens", 0) if response.usage else 0,
+            )
 
             return {"answer": answer, "mode": "quick_lookup_gpt52", "cost": cost}
         except Exception as e:
@@ -694,8 +733,14 @@ async def _standard_research(self, query: str) -> dict:
             )
             chat.append(user(query))
 
-            # Get response with automatic agentic search
-            response = chat.sample()
+            # xAI SDK's chat.sample() is a synchronous network call. Running
+            # it directly here would freeze the event loop for 5-15s per
+            # call, stalling every concurrent chat session, WebSocket emit,
+            # and MCP request handler that shares this loop. Push it to a
+            # worker thread so async behaviour holds.
+            import asyncio as _asyncio_local
+
+            response = await _asyncio_local.to_thread(chat.sample)
 
             # Extract answer and citations
             answer = response.content
@@ -751,11 +796,7 @@ async def _standard_research(self, query: str) -> dict:
 
                 answer = f"{response.choices[0].message.content or ''}\n\n[Note: Grok web search unavailable, using GPT-5.2 knowledge instead]"
 
-                cost = 0.01
-                if response.usage:
-                    input_cost = (response.usage.prompt_tokens / 1_000_000) * 1.75
-                    output_cost = (response.usage.completion_tokens / 1_000_000) * 14.00
-                    cost = input_cost + output_cost
+                cost = _chat_token_cost(response.usage, "gpt-5.2") if response.usage else 0.01
                 self.cost_accumulated += cost
 
                 # Record fallback cost
@@ -784,7 +825,16 @@ async def _deep_research(self, query: str) -> dict:
         Returns:
             Dict with job_id and estimated_cost
         """
-        estimated_cost = 0.20  # Average estimate
+        # Use the registry estimate as the budget reservation. The previous
+        # hard-coded $0.20 was ~10x lower than the registry cost ($2.00)
+        # for o4-mini-deep-research, so the session/daily budgets were
+        # silently exhausted ten times faster than tracked.
+        from deepr.providers.registry import get_cost_estimate as _get_cost_estimate
+
+        try:
+            estimated_cost = float(_get_cost_estimate("o4-mini-deep-research"))
+        except Exception:
+            estimated_cost = 2.00  # registry default for o4-mini-deep-research
 
         # Check cost safety before proceeding
         allowed, reason, _needs_confirm = self.cost_safety.check_operation(
@@ -1293,10 +1343,35 @@ def report_status(status: str):
 
             max_rounds = 5  # Prevent infinite loops
             round_count = 0
+            # Worst-case per-round estimate used for the mid-loop budget
+            # guard. Cheaper models will under-use this, but a tool round
+            # that runs a multi-thousand-token reasoning call shouldn't
+            # silently blow past the session budget between rounds.
+            _per_round_estimate = max(
+                _chat_token_cost(getattr(first_response, "usage", None), selected_model.model) * 1.5, 0.05
+            )
 
             while current_message.tool_calls and round_count < max_rounds:
                 round_count += 1
 
+                # Re-check the session budget between rounds. The previous
+                # implementation only checked once at the start of the
+                # request, so a 5-round loop on a frontier model could
+                # blow well past `self.budget` before any guard tripped.
+                _can_continue, _round_reason = self.cost_session.can_proceed(_per_round_estimate)
+                if not _can_continue:
+                    logger.warning(
+                        "Tool loop aborted at round %d: %s (accumulated $%.4f, budget $%.4f)",
+                        round_count,
+                        _round_reason,
+                        self.cost_accumulated,
+                        self.budget,
+                    )
+                    current_message.content = (
+                        current_message.content or ""
+                    ) + f"\n\n[Tool loop stopped after {round_count - 1} rounds: {_round_reason}]"
+                    break
+
                 # Process each tool call
                 tool_messages = []
 
@@ -1881,9 +1956,25 @@ def report_status(status: str):
             conversation_messages = [{"role": "system", "content": self.get_system_message()}, *self.messages]
             max_rounds = 5
             round_count = 0
+            _per_round_estimate = max(
+                _chat_token_cost(getattr(first_response, "usage", None), selected_model.model) * 1.5, 0.05
+            )
 
             while current_message.tool_calls and round_count < max_rounds:
                 round_count += 1
+                # Mirror the non-streaming branch: re-check budget between
+                # rounds so the tool loop can't run away on its own.
+                _can_continue, _round_reason = self.cost_session.can_proceed(_per_round_estimate)
+                if not _can_continue:
+                    logger.warning(
+                        "Tool loop (streaming) aborted at round %d: %s",
+                        round_count,
+                        _round_reason,
+                    )
+                    current_message.content = (
+                        current_message.content or ""
+                    ) + f"\n\n[Tool loop stopped after {round_count - 1} rounds: {_round_reason}]"
+                    break
                 tool_messages = []
 
                 for tool_call in current_message.tool_calls:
 
@@ -366,8 +366,9 @@ def save(self, path: Path):
         }
 
         path.parent.mkdir(parents=True, exist_ok=True)
-        with open(path, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+        from deepr.utils.atomic_io import atomic_write_json
+
+        atomic_write_json(path, data)
 
     @classmethod
     def load(cls, path: Path) -> "ConfidenceCalibrator":