Fix stale retrieval note and doctor env leak from review

TEN: reset _last_grounding/_last_sdk_ms at the start of _query_moss so a
failed second search in a turn no longer replays the previous hit in the
retrieval note, and send the note the current grounding return value.

custom-llm: snapshot MOCK and CUSTOM_LLM_API_KEY before the first
create_app() and restore both in an outer finally so run_doctor() (called
in-process by the tests) stops leaking env into later tests.

Add regressions for both.

Author: abojja9

apps/agora-custom-llm-moss/server/src/llm.py

@@ -346,26 +346,40 @@ async def health():
 def run_doctor() -> None:
     from fastapi.testclient import TestClient
 
-    os.environ["MOCK"] = "1"
-    payload = {
-        "model": "mock",
-        "stream": True,
-        "messages": [{"role": "user", "content": "How long do refunds take?"}],
-    }
-    for mode in ("ambient", "tool"):
-        with TestClient(create_app(mode)) as client:
-            ok = client.post("/chat/completions", json=payload)
-        if ok.status_code != 200 or "data: [DONE]" not in ok.text:
-            raise SystemExit(f"doctor {mode} failed: {ok.status_code} {ok.text}")
-        print(f"doctor {mode}: ok")
-    os.environ["MOCK"] = "0"
-    with TestClient(create_app("ambient")) as client:
-        denied = client.post("/chat/completions", json=payload)
-    if denied.status_code != 401:
-        raise SystemExit(f"doctor bearer: expected 401, got {denied.status_code}")
-    print("doctor bearer: rejected missing Authorization")
-    saved_key = os.environ.pop("CUSTOM_LLM_API_KEY", None)
+    # run_doctor runs in-process from the test suite, so snapshot every env var
+    # it flips and restore the originals on the way out (None means "was absent").
+    saved_mock = os.environ.get("MOCK")
+    saved_key = os.environ.get("CUSTOM_LLM_API_KEY")
+
+    def _restore(name: str, value: str | None) -> None:
+        if value is None:
+            os.environ.pop(name, None)
+        else:
+            os.environ[name] = value
+
     try:
+        os.environ["MOCK"] = "1"
+        payload = {
+            "model": "mock",
+            "stream": True,
+            "messages": [{"role": "user", "content": "How long do refunds take?"}],
+        }
+        for mode in ("ambient", "tool"):
+            with TestClient(create_app(mode)) as client:
+                ok = client.post("/chat/completions", json=payload)
+            if ok.status_code != 200 or "data: [DONE]" not in ok.text:
+                raise SystemExit(f"doctor {mode} failed: {ok.status_code} {ok.text}")
+            print(f"doctor {mode}: ok")
+        os.environ["MOCK"] = "0"
+        with TestClient(create_app("ambient")) as client:
+            denied = client.post("/chat/completions", json=payload)
+        if denied.status_code != 401:
+            raise SystemExit(f"doctor bearer: expected 401, got {denied.status_code}")
+        print("doctor bearer: rejected missing Authorization")
+        # Temporarily remove the key so the next request runs against an unset
+        # key; create_app reloads .env, so pop again after building the app. The
+        # outer finally owns the real restore.
+        os.environ.pop("CUSTOM_LLM_API_KEY", None)
         test_app = create_app("ambient")
         os.environ.pop("CUSTOM_LLM_API_KEY", None)
         with TestClient(test_app) as client:
@@ -378,13 +392,11 @@ def run_doctor() -> None:
             raise SystemExit(
                 f"doctor bearer: unset key should reject any token, got {any_token.status_code}"
             )
+        print("doctor bearer: rejected any token while CUSTOM_LLM_API_KEY is unset")
+        print("doctor: ok")
     finally:
-        if saved_key is None:
-            os.environ.pop("CUSTOM_LLM_API_KEY", None)
-        else:
-            os.environ["CUSTOM_LLM_API_KEY"] = saved_key
-    print("doctor bearer: rejected any token while CUSTOM_LLM_API_KEY is unset")
-    print("doctor: ok")
+        _restore("MOCK", saved_mock)
+        _restore("CUSTOM_LLM_API_KEY", saved_key)
 
 
 if __name__ == "__main__":

apps/agora-custom-llm-moss/tests/test_llm.py

@@ -194,3 +194,23 @@ def _load(_server_dir=None) -> None:
     out = capsys.readouterr().out
     assert "rejected any token while CUSTOM_LLM_API_KEY is unset" in out
     assert "doctor: ok" in out
+
+
+def test_doctor_restores_preset_env(
+    monkeypatch: pytest.MonkeyPatch, moss_ok
+) -> None:
+    monkeypatch.setenv("MOCK", "preset")
+    monkeypatch.setenv("CUSTOM_LLM_API_KEY", "preset-key")
+    llm.run_doctor()
+    assert os.environ["MOCK"] == "preset"
+    assert os.environ["CUSTOM_LLM_API_KEY"] == "preset-key"
+
+
+def test_doctor_restores_absent_env(
+    monkeypatch: pytest.MonkeyPatch, moss_ok
+) -> None:
+    monkeypatch.delenv("MOCK", raising=False)
+    monkeypatch.delenv("CUSTOM_LLM_API_KEY", raising=False)
+    llm.run_doctor()
+    assert "MOCK" not in os.environ
+    assert "CUSTOM_LLM_API_KEY" not in os.environ

apps/ten-moss/tenapp/ten_packages/extension/main_python/extension.py

@@ -247,7 +247,7 @@ async def _on_tool_call(self, cmd: Cmd) -> None:
             else:
                 self._moss_tool_calls += 1
                 grounding = await self._query_moss(query)
-                await self._send_retrieval_note(self._last_grounding, self._last_sdk_ms)
+                await self._send_retrieval_note(grounding, self._last_sdk_ms)
         else:
             self.ten_env.log_error(
                 f"[MainControlExtension] unknown tool_call name={name!r}"
@@ -261,6 +261,10 @@ async def _on_tool_call(self, cmd: Cmd) -> None:
         await self.ten_env.return_result(result)
 
     async def _query_moss(self, user_text: str) -> str:
+        # Reset per-query state up front so a failed search never replays the
+        # previous hit's grounding/latency in the retrieval note.
+        self._last_grounding = ""
+        self._last_sdk_ms = None
         if self.moss is None:
             return ""
         try:

apps/ten-moss/tests/test_retrieval_note.py

@@ -0,0 +1,168 @@
+"""Regression for the retrieval note after a failed second search in a turn.
+
+extension.py imports the TEN runtime, which is not installed offline, so we load
+just that one file against light stubs and drive its Moss methods directly.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+EXTENSION = (
+    Path(__file__).resolve().parents[1]
+    / "tenapp"
+    / "ten_packages"
+    / "extension"
+    / "main_python"
+    / "extension.py"
+)
+
+
+def _load_extension_module():
+    pkg = "main_python"
+
+    def stub(name: str) -> types.ModuleType:
+        mod = types.ModuleType(name)
+        sys.modules[name] = mod
+        return mod
+
+    runtime = stub("ten_runtime")
+    for symbol in ("AsyncTenEnv", "Cmd", "Data"):
+        setattr(runtime, symbol, type(symbol, (), {}))
+    runtime.AsyncExtension = type(
+        "AsyncExtension", (), {"__init__": lambda self, name: None}
+    )
+    runtime.StatusCode = type("StatusCode", (), {"OK": "ok"})
+
+    class _CmdResult:
+        def __init__(self):
+            self.content = None
+
+        @classmethod
+        def create(cls, _status, _cmd):
+            return cls()
+
+        def set_property_from_json(self, _key, value):
+            self.content = value
+
+    runtime.CmdResult = _CmdResult
+
+    ten_moss = stub("ten_moss")
+    ten_moss.MossSessionManager = type("MossSessionManager", (), {})
+
+    const = stub("ten_ai_base.const")
+    const.CMD_PROPERTY_RESULT = "result"
+    ai_types = stub("ten_ai_base.types")
+    ai_types.LLMToolMetadata = type("LLMToolMetadata", (), {})
+    ai_types.LLMToolMetadataParameter = type("LLMToolMetadataParameter", (), {})
+    ai_base = stub("ten_ai_base")
+    ai_base.const = const
+    ai_base.types = ai_types
+
+    parent = stub(pkg)
+    parent.__path__ = []
+    agent_pkg = stub(f"{pkg}.agent")
+    agent_pkg.__path__ = []
+    decorators = stub(f"{pkg}.agent.decorators")
+    decorators.agent_event_handler = lambda *a, **k: (lambda fn: fn)
+    agent_mod = stub(f"{pkg}.agent.agent")
+    agent_mod.Agent = type("Agent", (), {})
+    events = stub(f"{pkg}.agent.events")
+    for evt in (
+        "ASRResultEvent",
+        "LLMResponseEvent",
+        "ToolRegisterEvent",
+        "UserJoinedEvent",
+        "UserLeftEvent",
+    ):
+        setattr(events, evt, type(evt, (), {}))
+    helper = stub(f"{pkg}.helper")
+    helper._send_cmd = helper._send_data = lambda *a, **k: None
+    helper.parse_sentences = lambda frag, delta: ([], "")
+    config = stub(f"{pkg}.config")
+    config.MainControlConfig = type("MainControlConfig", (), {})
+
+    spec = importlib.util.spec_from_file_location(f"{pkg}.extension", EXTENSION)
+    module = importlib.util.module_from_spec(spec)
+    module.__package__ = pkg
+    spec.loader.exec_module(module)
+    return module
+
+
+extension = _load_extension_module()
+
+
+class _RecordingEnv:
+    def __init__(self):
+        self.results = []
+
+    def log_info(self, *_a):
+        pass
+
+    def log_error(self, *_a):
+        pass
+
+    async def return_result(self, result):
+        self.results.append(result)
+
+
+class _FlakyMoss:
+    """Succeeds once, then raises - a second search failing mid-turn."""
+
+    last_time_taken_ms = 12
+
+    def __init__(self):
+        self.calls = 0
+
+    async def query_context(self, _text: str) -> str:
+        self.calls += 1
+        if self.calls == 1:
+            return "Refunds land in 3-5 business days."
+        raise RuntimeError("moss backend unavailable")
+
+
+class _ToolCallCmd:
+    """A tool_call Cmd carrying a search_knowledge_base payload."""
+
+    def __init__(self, query: str):
+        self._raw = json.dumps(
+            {"name": "search_knowledge_base", "arguments": {"query": query}}
+        )
+
+    def get_property_to_json(self, _key):
+        return self._raw, None
+
+
+@pytest.mark.asyncio
+async def test_failed_second_search_does_not_replay_first_hit() -> None:
+    ext = extension.MainControlExtension("main_control")
+    ext.ten_env = _RecordingEnv()
+    ext.moss = _FlakyMoss()
+
+    notes: list[str] = []
+
+    async def capture(role, text, final, stream_id, data_type="text"):
+        notes.append(text)
+
+    ext._send_transcript = capture
+
+    # Drive the real tool-call handler so the note reflects production flow.
+    await ext._on_tool_call(_ToolCallCmd("How long do refunds take?"))
+    assert "3-5 business days" in notes[0]
+    assert "3-5 business days" in ext.ten_env.results[0].content
+
+    await ext._on_tool_call(_ToolCallCmd("What about exchanges?"))
+
+    # Second search failed: the note must not replay the first hit, and the
+    # tool result is empty - the two must agree.
+    assert "3-5 business days" not in notes[1]
+    assert "no match" in notes[1]
+    assert '"content": ""' in ext.ten_env.results[1].content
+    assert ext._last_grounding == ""
+    assert ext._last_sdk_ms is None