Merge ab4da6f5cf into 20e7fc0164

src/llm_core.py

@@ -8,9 +8,10 @@ import hashlib
 import threading
 import re
 import os
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from fastapi import HTTPException
-from typing import Optional, Dict, List, Tuple
+from typing import Callable, Optional, Dict, List, Tuple
 from src.model_context import get_context_length, DEFAULT_CONTEXT, is_local_endpoint
 from urllib.parse import urlparse
 
@@ -19,6 +20,37 @@ logger = logging.getLogger(__name__)
 _LOCAL_MODEL_LOCK = asyncio.Lock()
 _LOCAL_MODEL_WAITING_FOREGROUND = 0
 _LOCAL_MODEL_CURRENT: Dict[str, object] = {}
+_MODEL_USAGE_RECORDER: ContextVar[Optional[Callable[[str], None]]] = ContextVar(
+    "model_usage_recorder",
+    default=None,
+)
+
+
+@contextmanager
+def capture_model_usage(recorder: Callable[[str], None]):
+    """Report successful fallback candidates within the current context.
+
+    Context-local storage keeps concurrent chat and scheduled-task calls from
+    attributing one another's models. Callers can use this without changing
+    the return type of the shared LLM helpers.
+    """
+    token = _MODEL_USAGE_RECORDER.set(recorder)
+    try:
+        yield
+    finally:
+        _MODEL_USAGE_RECORDER.reset(token)
+
+
+def _report_model_usage(model: str) -> None:
+    recorder = _MODEL_USAGE_RECORDER.get()
+    if recorder is None:
+        return
+    try:
+        recorder(model)
+    except Exception:
+        # Usage accounting must never turn a successful model response into a
+        # failed request.
+        logger.debug("Model usage recorder failed", exc_info=True)
 
 
 def _local_model_gate_enabled() -> bool:
@@ -1932,7 +1964,9 @@ def llm_call_with_fallback(candidates, messages, **kwargs) -> str:
     last_err = None
     for i, (url, model, headers) in enumerate(cands):
         try:
-            return llm_call(url, model, messages, headers=headers, **kwargs)
+            result = llm_call(url, model, messages, headers=headers, **kwargs)
+            _report_model_usage(model)
+            return result
         except Exception as e:
             last_err = e
             tag = "primary" if i == 0 else "candidate"
@@ -1949,7 +1983,9 @@ async def llm_call_async_with_fallback(candidates, messages, **kwargs) -> str:
     last_err = None
     for i, (url, model, headers) in enumerate(cands):
         try:
-            return await llm_call_async(url, model, messages, headers=headers, **kwargs)
+            result = await llm_call_async(url, model, messages, headers=headers, **kwargs)
+            _report_model_usage(model)
+            return result
         except Exception as e:
             last_err = e
             tag = "primary" if i == 0 else "candidate"

src/task_scheduler.py

@@ -1230,6 +1230,7 @@ class TaskScheduler:
     async def _execute_action(self, task, run_id: str | None = None) -> tuple:
         """Execute a built-in action (no LLM needed)."""
         from src.builtin_actions import BUILTIN_ACTIONS
+        from src.llm_core import capture_model_usage
 
         action_fn = BUILTIN_ACTIONS.get(task.action)
         if not action_fn:
@@ -1250,6 +1251,14 @@ class TaskScheduler:
             # through as `command` so action_cookbook_serve can json.loads it.
             elif task.action == "cookbook_serve" and task.prompt:
                 kwargs["command"] = task.prompt
+            def _record_model(model: str):
+                self._last_run_model = model
+
+            # Some built-in actions use the shared LLM fallback chain even
+            # though their task type is "action". Capture the candidate that
+            # actually succeeds so TaskRun.model is populated just like it is
+            # for ordinary LLM and research tasks.
+            with capture_model_usage(_record_model):
                 result, success = await action_fn(**kwargs)
             return result, success
         except TaskNoop:

tests/test_task_action_model_usage.py

@@ -0,0 +1,121 @@
+from types import SimpleNamespace
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_async_fallback_reports_only_successful_model(monkeypatch):
+    from src import llm_core
+
+    attempted = []
+
+    async def fake_call(url, model, messages, **kwargs):
+        attempted.append(model)
+        if model == "primary":
+            raise RuntimeError("primary unavailable")
+        return "fallback response"
+
+    monkeypatch.setattr(llm_core, "llm_call_async", fake_call)
+    used = []
+
+    with llm_core.capture_model_usage(used.append):
+        result = await llm_core.llm_call_async_with_fallback(
+            [
+                ("http://primary/v1", "primary", {}),
+                ("http://fallback/v1", "fallback", {}),
+            ],
+            messages=[{"role": "user", "content": "hello"}],
+        )
+
+    assert result == "fallback response"
+    assert attempted == ["primary", "fallback"]
+    assert used == ["fallback"]
+
+
+def test_sync_primary_success_reports_model(monkeypatch):
+    from src import llm_core
+
+    monkeypatch.setattr(llm_core, "llm_call", lambda *args, **kwargs: "ok")
+    used = []
+
+    with llm_core.capture_model_usage(used.append):
+        result = llm_core.llm_call_with_fallback(
+            [("http://primary/v1", "primary", {})],
+            messages=[{"role": "user", "content": "hello"}],
+        )
+
+    assert result == "ok"
+    assert used == ["primary"]
+
+
+@pytest.mark.asyncio
+async def test_builtin_action_records_actual_fallback_model(monkeypatch):
+    from src import builtin_actions, llm_core
+    from src.task_scheduler import TaskScheduler
+
+    async def fake_call(url, model, messages, **kwargs):
+        if model == "primary":
+            raise RuntimeError("primary unavailable")
+        return "summary"
+
+    async def model_backed_action(**kwargs):
+        result = await llm_core.llm_call_async_with_fallback(
+            [
+                ("http://primary/v1", "primary", {}),
+                ("http://fallback/v1", "fallback", {}),
+            ],
+            messages=[{"role": "user", "content": "summarize"}],
+        )
+        return result, True
+
+    monkeypatch.setattr(llm_core, "llm_call_async", fake_call)
+    monkeypatch.setitem(
+        builtin_actions.BUILTIN_ACTIONS,
+        "test_model_usage",
+        model_backed_action,
+    )
+
+    scheduler = TaskScheduler(session_manager=None)
+    scheduler._last_run_model = None
+    task = SimpleNamespace(
+        action="test_model_usage",
+        owner="alice",
+        name="Model-backed action",
+        prompt=None,
+    )
+
+    result, success = await scheduler._execute_action(task)
+
+    assert success is True
+    assert result == "summary"
+    assert scheduler._last_run_model == "fallback"
+
+
+@pytest.mark.asyncio
+async def test_non_model_action_leaves_usage_empty(monkeypatch):
+    from src import builtin_actions
+    from src.task_scheduler import TaskScheduler
+
+    async def housekeeping_action(**kwargs):
+        return "clean", True
+
+    monkeypatch.setitem(
+        builtin_actions.BUILTIN_ACTIONS,
+        "test_housekeeping",
+        housekeeping_action,
+    )
+
+    scheduler = TaskScheduler(session_manager=None)
+    scheduler._last_run_model = None
+    task = SimpleNamespace(
+        action="test_housekeeping",
+        owner="alice",
+        name="Housekeeping",
+        prompt=None,
+    )
+
+    result, success = await scheduler._execute_action(task)
+
+    assert success is True
+    assert result == "clean"
+    assert scheduler._last_run_model is None