From a5979eac505ee752406a2788fdeb86df932f00fd Mon Sep 17 00:00:00 2001
From: Liang Jiaqing <l.j.q.light@gmail.com>
Date: Fri, 16 Oct 2026 09:30:02 +0800
Subject: [PATCH] fix: track applied mykeys snapshot per agent

---
 agentmain.py                          |  14 +-
 frontends/desktop_bridge.py           |  34 -----
 frontends/tests/test_mykeys_reload.py | 203 ++++++++++++++++++++++++++
 frontends/tuiapp_v2.py                |   7 +-
 llmcore.py                            |  95 +++++++-----
 5 files changed, 275 insertions(+), 78 deletions(-)
 create mode 100644 frontends/tests/test_mykeys_reload.py

diff --git a/agentmain.py b/agentmain.py
index d8cab8cf..dfa01e6c 100644
--- a/agentmain.py
+++ b/agentmain.py
@@ -6,7 +6,7 @@ if sys.stderr is None: sys.stderr = open(os.devnull, "w")
 elif hasattr(sys.stderr, 'reconfigure'): sys.stderr.reconfigure(errors='replace')
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from llmcore import reload_mykeys, ToolClient, MixinSession, NativeToolClient, NativeClaudeSession, NativeOAISession, resolve_client
+from llmcore import mykeys_store, ToolClient, MixinSession, NativeToolClient, NativeClaudeSession, NativeOAISession, resolve_client
 from agent_loop import agent_runner_loop
 try:
     from plugins.hooks import discover_and_load; discover_and_load()
@@ -64,8 +64,8 @@ class GenericAgent:
         self.intervene = self.extrakeyinfo = None
 
     def load_llm_sessions(self):
-        mykeys, changed = reload_mykeys()
-        if not changed and hasattr(self, 'llmclients'): return
+        mykeys, updated_at = mykeys_store.snapshot()
+        if hasattr(self, 'llmclients') and getattr(self, '_mykeys_updated_at', None) == updated_at: return
         try: oldhistory, oldname = self.llmclient.backend.history, self.llmclient.backend.name
         except: oldhistory = oldname = None
         llm_sessions = []
@@ -73,7 +73,7 @@ class GenericAgent:
             if not any(x in k for x in ['api', 'config', 'cookie']): continue
             try:
                 if 'mixin' in k: llm_sessions += [{'mixin_cfg': cfg}]
-                elif c := resolve_client(k): llm_sessions += [c]
+                elif c := resolve_client(k, cfg): llm_sessions += [c]
             except: pass
         for i, s in enumerate(llm_sessions):
             if isinstance(s, dict) and 'mixin_cfg' in s:
@@ -83,11 +83,15 @@ class GenericAgent:
                     else: llm_sessions[i] = ToolClient(mixin)
                 except Exception as e: print(f'\n\n\n[ERROR] Failed to init MixinSession with cfg {s["mixin_cfg"]}: {e}!!!\n\n')
         self.llmclients = llm_sessions
-        if not self.llmclients: return
+        if not self.llmclients:
+            self.llmclient = None
+            self._mykeys_updated_at = updated_at
+            return
         names = [c.backend.name if not isinstance(c, dict) else f'BADMIXIN_{i}' for i, c in enumerate(self.llmclients)]
         if oldname in names: self.llm_no = names.index(oldname)
         self.llmclient = self.llmclients[self.llm_no%len(self.llmclients)]
         if oldhistory: self.llmclient.backend.history = oldhistory
+        self._mykeys_updated_at = updated_at
     
     def next_llm(self, n=-1):
         self.load_llm_sessions()
diff --git a/frontends/desktop_bridge.py b/frontends/desktop_bridge.py
index f44ad287..726a4ac6 100644
--- a/frontends/desktop_bridge.py
+++ b/frontends/desktop_bridge.py
@@ -595,13 +595,6 @@ class AgentManager:
         lines = [f"    '{k}': {json.dumps(v, ensure_ascii=False)}," if isinstance(v, str) else f"    '{k}': {v}," for k, v in d.items()]
         return "{\n" + "\n".join(lines) + "\n}"
 
-    def _invalidate_mykey_cache(self) -> None:
-        self.ensure_ga_import_path()
-        sys.modules.pop("mykey", None)
-        with contextlib.suppress(Exception):
-            import llmcore
-            llmcore._mykey_mtime = None
-
     def _profile_keys(self) -> List[str]:
         self.ensure_ga_import_path()
         from llmcore import reload_mykeys
@@ -679,35 +672,8 @@ class AgentManager:
 
     def _save_mykey_text(self, text: str) -> list:
         self._mykey_file().write_text(text, encoding="utf-8")
-        self._invalidate_mykey_cache()
-        self._reload_live_agents()
         return self.list_model_profiles()
 
-    def _reload_live_agents(self) -> None:
-        """mykey.py 改动后，强制所有活着的会话 agent 重建 LLM session，让新 key/模型
-        立即生效（无需重启）。重建保留对话 history（agentmain 内部用 oldhistory 接回）。
-
-        纯 bridge 侧实现，不改 agentmain：每次调 agent.load_llm_sessions() 前，把
-        llmcore 的全局 mtime 标志清空（与 _invalidate_mykey_cache 同一手法），使其内部
-        reload_mykeys() 报告 changed=True、从而真正重建——否则刷新模型列表等路径会先
-        消费掉变更标志，常驻 agent 的 load_llm_sessions 会因 changed=False 跳过重建。"""
-        self.ensure_ga_import_path()
-        try:
-            import llmcore
-        except Exception:
-            return
-        with self.lock:
-            agents = [s.agent for s in self.sessions.values() if getattr(s, "agent", None) is not None]
-        for agent in agents:
-            fn = getattr(agent, "load_llm_sessions", None)
-            if not callable(fn):
-                continue
-            try:
-                llmcore._mykey_mtime = None   # 让本次 reload_mykeys() 视为“已变更”，触发真正重建
-                fn()
-            except Exception as e:
-                bridge_print(f"[bridge] reload live agent failed: {e}")
-
     def add_model_profile(self, data: dict) -> dict:
         cfg = self._build_cfg(data)
         text = self._mykey_file().read_text(encoding="utf-8")
diff --git a/frontends/tests/test_mykeys_reload.py b/frontends/tests/test_mykeys_reload.py
new file mode 100644
index 00000000..f18121ea
--- /dev/null
+++ b/frontends/tests/test_mykeys_reload.py
@@ -0,0 +1,203 @@
+"""Isolate configuration loading from real keys, network and agent startup."""
+import ast
+import importlib.util
+import json
+import os
+import sys
+import threading
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def extract(path, names, namespace):
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    nodes = [n for n in tree.body if getattr(n, "name", None) in names]
+    exec(compile(ast.Module(body=nodes, type_ignores=[]), str(path), "exec"), namespace)
+
+
+@pytest.fixture
+def core(tmp_path, monkeypatch):
+    # Never resolve or import the user's actual mykey module.
+    source = tmp_path / "mykey.py"
+    def find_spec(name):
+        assert name == "mykey"
+        return importlib.util.spec_from_file_location(name, source) if source.exists() else None
+    monkeypatch.setattr(importlib.util, "find_spec", find_spec)
+    monkeypatch.delitem(sys.modules, "mykey", raising=False)
+    ns = dict(os=os, sys=sys, json=json, time=time, threading=threading,
+              importlib=importlib, __file__=str(tmp_path / "llmcore.py"))
+    extract(ROOT / "llmcore.py", {"MyKeys", "reload_mykeys", "_load_mykeys", "__getattr__", "resolve_session", "resolve_client"}, ns)
+    ns["mykeys_store"] = ns["MyKeys"]()
+    yield ns, source
+    sys.modules.pop("mykey", None)
+
+
+def write(source, value, stamp=1700000000000000100):
+    source.write_text("oai_config = " + repr({"name": "test", "model": value}) + "\n", encoding="utf-8")
+    os.utime(source, ns=(stamp, stamp))
+
+
+def agent_type(ns):
+    tree = ast.parse((ROOT / "agentmain.py").read_text(encoding="utf-8"))
+    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "GenericAgent")
+    cls.body = [n for n in cls.body if getattr(n, "name", None) == "load_llm_sessions"]
+    exec(compile(ast.Module(body=[cls], type_ignores=[]), "agentmain.py", "exec"), ns)
+    return ns["GenericAgent"]
+
+
+def test_two_agents_update_even_if_bridge_read_first(core):
+    ns, source = core
+    write(source, "A")
+    built = []
+    def resolve(name, cfg):
+        built.append(cfg["model"])
+        return SimpleNamespace(backend=SimpleNamespace(name=cfg["name"], model=cfg["model"], history=[]))
+    ns["resolve_client"] = resolve
+    Agent = agent_type(ns)
+    a, b = Agent(), Agent()
+    for ga in (a, b):
+        ga.llm_no = 0
+        ga.load_llm_sessions()
+    history = [{"role": "user", "content": "preserve"}]
+    a.llmclient.backend.history = history
+    old = a.llmclient
+    a.load_llm_sessions()
+    assert a.llmclient is old
+    write(source, "B", 1700000000000000200)
+    assert ns["reload_mykeys"]()[1] is True  # bridge consumes old changed signal
+    assert ns["reload_mykeys"]()[1] is False
+    for ga in (a, b):
+        ga.load_llm_sessions()
+        assert ga.llmclient.backend.model == "B"
+    assert a.llmclient.backend.history is history
+    assert built == ["A", "A", "B", "B"]
+
+
+def test_same_second_same_size_and_backwards_mtime(core):
+    ns, source = core
+    write(source, "A")
+    store = ns["mykeys_store"]
+    first = store.snapshot()[1]
+    write(source, "B", 1700000000000000200)
+    data, second = store.snapshot()
+    assert data["oai_config"]["model"] == "B" and second > first
+    write(source, "C", 1600000000000000000)
+    data, third = store.snapshot()
+    assert data["oai_config"]["model"] == "C" and third > second
+
+
+def test_json_fallback_and_legacy_interfaces(core):
+    ns, source = core
+    source.with_suffix(".json").write_text('{"oai_config": {"model": "A"}}', encoding="utf-8")
+    data, changed = ns["reload_mykeys"]()
+    assert changed is True
+    assert ns["__getattr__"]("mykeys") is data
+    assert ns["reload_mykeys"]() == (data, False)
+    assert ns["_load_mykeys"]() == data
+    with pytest.raises(AttributeError): ns["__getattr__"]("unknown")
+
+
+def test_failure_keeps_last_good_snapshot_and_recovers(core):
+    ns, source = core
+    write(source, "A")
+    store = ns["mykeys_store"]
+    before = store.snapshot()
+    source.write_text("oai_config = {", encoding="utf-8")
+    assert store.snapshot() == before
+    source.unlink()
+    assert store.snapshot() == before
+    write(source, "B", 1700000000000000300)
+    assert store.snapshot()[0]["oai_config"]["model"] == "B"
+
+
+def test_force_refresh_publishes_once_for_all_readers(core):
+    ns, source = core
+    write(source, "A")
+    store = ns["mykeys_store"]
+    before = store.snapshot()[1]
+    store.reload(force=True)
+    after = store.snapshot()[1]
+    assert after > before
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        results = list(pool.map(lambda _: store.snapshot()[1], range(32)))
+    assert results == [after] * 32
+
+
+def test_resolvers_accept_snapshot_without_reading_global_or_mutating_it(core):
+    ns, source = core
+    class Session:
+        def __init__(self, cfg): self.cfg = cfg
+    ns.update(LLMSession=Session, ClaudeSession=Session, NativeClaudeSession=type("NC", (), {}),
+              NativeOAISession=type("NO", (), {}), ToolClient=lambda s: s)
+    cfg = {"model": "A"}
+    ns["reload_mykeys"] = lambda: pytest.fail("unexpected global config read")
+    session = ns["resolve_client"]("oai_config", cfg)
+    assert session.cfg == {"model": "A", "_mykey_name": "oai_config"}
+    assert cfg == {"model": "A"}
+    ns["reload_mykeys"] = lambda: ({"oai_config": cfg}, False)
+    assert ns["resolve_session"]("oai_config").cfg == session.cfg
+
+
+def test_bridge_save_only_writes_config_and_reads_profiles(tmp_path):
+    tree = ast.parse((ROOT / "frontends/desktop_bridge.py").read_text(encoding="utf-8"))
+    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "AgentManager")
+    method = next(n for n in cls.body if getattr(n, "name", None) == "_save_mykey_text")
+    ns = {}
+    exec(compile(ast.Module(body=[method], type_ignores=[]), "desktop_bridge.py", "exec"), ns)
+    file = tmp_path / "synthetic_config.py"
+    manager = SimpleNamespace(_mykey_file=lambda: file, list_model_profiles=lambda: ["profile"])
+    assert ns["_save_mykey_text"](manager, "synthetic = {}") == ["profile"]
+    assert file.read_text(encoding="utf-8") == "synthetic = {}"
+
+
+def test_ga_preserves_selected_model_when_config_order_changes(core):
+    ns, source = core
+    ns["resolve_client"] = lambda name, cfg: SimpleNamespace(backend=SimpleNamespace(name=cfg["name"], history=[]))
+    Agent = agent_type(ns)
+    ga = Agent()
+    ga.llm_no = 1
+    source.write_text("oai_config_a = {'name': 'A'}\noai_config_b = {'name': 'B'}\n", encoding="utf-8")
+    ga.load_llm_sessions()
+    assert ga.llmclient.backend.name == "B"
+    source.write_text("oai_config_b = {'name': 'B'}\noai_config_a = {'name': 'A'}\n", encoding="utf-8")
+    os.utime(source, ns=(1700000000000000000, 1700000000000000000))
+    ga.load_llm_sessions()
+    assert ga.llmclient.backend.name == "B" and ga.llm_no == 0
+
+
+def test_ga_drops_stale_client_when_models_removed(core):
+    ns, source = core
+    ns["resolve_client"] = lambda name, cfg: SimpleNamespace(backend=SimpleNamespace(name=cfg["name"], history=[]))
+    Agent = agent_type(ns)
+    ga = Agent()
+    ga.llm_no = 0
+    write(source, "A")
+    ga.load_llm_sessions()
+    source.write_text("", encoding="utf-8")
+    os.utime(source, ns=(1700000000000000500, 1700000000000000500))
+    ga.load_llm_sessions()
+    assert ga.llmclients == [] and ga.llmclient is None
+
+
+def test_loader_does_not_reuse_same_second_bytecode(core):
+    import py_compile
+    ns, source = core
+    write(source, "A")
+    py_compile.compile(str(source), doraise=True)
+    write(source, "B", 1700000000000000200)
+    assert ns["mykeys_store"].snapshot()[0]["oai_config"]["model"] == "B"
+
+
+def test_legacy_direct_loader_raises_and_cleans_failed_module(core):
+    ns, source = core
+    source.write_text("partial = 1\nraise RuntimeError('synthetic failure')\n", encoding="utf-8")
+    with pytest.raises(RuntimeError, match="synthetic failure"):
+        ns["_load_mykeys"]()
+    assert "mykey" not in sys.modules
+    assert ns["reload_mykeys"]() == ({}, False)
diff --git a/frontends/tuiapp_v2.py b/frontends/tuiapp_v2.py
index 5234e70a..8a6a5626 100644
--- a/frontends/tuiapp_v2.py
+++ b/frontends/tuiapp_v2.py
@@ -5620,22 +5620,17 @@ class GenericAgentTUI(App[None]):
         self._refresh_all()
 
     def _cmd_reload_keys(self, args, raw):
-        # Force rebuild of every session's llmclients from a fresh mykey.py.
-        # reload_mykeys() uses a module-level mtime cache, so the first agent
-        # to call it consumes the "changed" signal and subsequent agents see
-        # changed=False (and skip rebuild). Invalidate the cache before each
-        # agent so every session picks up the new config.
         try:
             import llmcore
         except Exception as e:
             self._system(f"❌ 无法 import llmcore: {e}"); return
+        llmcore.mykeys_store.reload(force=True)
         n_ok = n_fail = 0
         for sess in self.sessions.values():
             agent = getattr(sess, "agent", None)
             if agent is None:
                 continue
             try:
-                llmcore._mykey_mtime = None
                 agent.load_llm_sessions()
                 n_ok += 1
             except Exception:
diff --git a/llmcore.py b/llmcore.py
index 3b0e6b06..611e8355 100644
--- a/llmcore.py
+++ b/llmcore.py
@@ -10,38 +10,67 @@ _RESP_CACHE_KEY = str(uuid.uuid4()); _RESP_CODEX_KEY = str(uuid.uuid4())
 _ROOT = os.path.dirname(os.path.abspath(__file__))
 if _ROOT not in sys.path: sys.path.append(_ROOT)
 
-def _load_mykeys():
-    global _mykey_path
-    try:
+class MyKeys:
+    def __init__(self):
+        self.path = self.mtime = self.updated_at = None
+        self.data = {}
+        self.lock = threading.Lock()
+
+    def _load(self):
         sys.modules.pop('mykey', None)
-        import mykey; _mykey_path = mykey.__file__
-        return {k: v for k, v in vars(mykey).items() if not k.startswith('_')}
-    except ImportError as e:
-        if getattr(e, 'name', None) != 'mykey':
-            raise Exception(f'[ERROR] mykey.py found but failed to import: {e}') from e
-    except SyntaxError as e:
-        raise Exception(f'[ERROR] mykey.py has syntax error: {e}') from e
-    p = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'mykey.json')
-    if not os.path.exists(p): raise Exception('[ERROR] mykey.py not found in sys.path and mykey.json not found. Run "python configure_mykey.py" or copy mykey_template.py to mykey.py and fill in your keys.')
-    with open(_mykey_path := p, encoding='utf-8') as f: mk = json.load(f)
-    if isinstance(mk, dict) and 'remote_url' in mk: return requests.get(mk['remote_url'], timeout=10).json()
-    return mk
+        spec = importlib.util.find_spec('mykey')
+        if spec is not None:
+            path = spec.origin
+            mt = os.stat(path).st_mtime_ns
+            module = importlib.util.module_from_spec(spec)
+            sys.modules['mykey'] = module
+            # Compile source directly: pyc timestamps only have second precision.
+            try:
+                exec(compile(spec.loader.get_source('mykey'), path, 'exec'), module.__dict__)
+            except BaseException:
+                sys.modules.pop('mykey', None)
+                raise
+            data = {k: v for k, v in vars(module).items() if not k.startswith('_')}
+        else:
+            path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'mykey.json')
+            mt = os.stat(path).st_mtime_ns
+            with open(path, encoding='utf-8') as f: data = json.load(f)
+            if isinstance(data, dict) and 'remote_url' in data:
+                data = requests.get(data['remote_url'], timeout=10).json()
+        if not isinstance(data, dict): raise ValueError('mykey configuration must be a dict')
+        if os.stat(path).st_mtime_ns != mt: raise RuntimeError('mykey changed while loading')
+        return data, path, mt
+
+    def reload(self, force=False):
+        with self.lock:
+            try:
+                mt = os.stat(self.path).st_mtime_ns if self.path else None
+                if not force and self.updated_at is not None and mt == self.mtime:
+                    return self.data, self.updated_at, False
+                data, path, mt = self._load()
+                self.data, self.path, self.mtime = data, path, mt
+                self.updated_at = max(time.monotonic_ns(), (self.updated_at or 0) + 1)
+                print(f'[Info] Load mykeys from {self.path}')
+                return self.data, self.updated_at, True
+            except Exception:
+                return self.data, self.updated_at, False
+
+    def snapshot(self):
+        data, updated_at, _ = self.reload()
+        return data, updated_at
+
+
+mykeys_store = MyKeys()
 
-_mykey_lock = threading.Lock()
-_mykey_path = _mykey_mtime = None
 def reload_mykeys():
-    global _mykey_mtime
-    try:
-        mt = os.stat(_mykey_path).st_mtime_ns if _mykey_path else -1
-        if mt == _mykey_mtime: return globals().get('mykeys', {}), False
-        with _mykey_lock: mk = _load_mykeys()
-        _mykey_mtime = os.stat(_mykey_path).st_mtime_ns
-        print(f'[Info] Load mykeys from {_mykey_path}')
-        globals().update(mykeys=mk)
-        return mk, True
-    except: return globals().get('mykeys', {}), False
+    data, _, changed = mykeys_store.reload()
+    return data, changed
 
-def __getattr__(name):  # once guard in PEP 562
+def _load_mykeys():
+    with mykeys_store.lock:
+        return mykeys_store._load()[0]
+
+def __getattr__(name):
     if name == 'mykeys': return reload_mykeys()[0]
     raise AttributeError(f"module 'llmcore' has no attribute {name}")
 
@@ -1223,16 +1252,16 @@ class NativeToolClient:
         if resp and hasattr(resp, 'tool_calls') and resp.tool_calls: self._pending_tool_ids = [tc.id for tc in resp.tool_calls]
         return resp
 
-def resolve_session(cfg_name):
-    cfg = reload_mykeys()[0].get(cfg_name)
+def resolve_session(cfg_name, cfg=None):
+    if cfg is None: cfg = reload_mykeys()[0].get(cfg_name)
     if not cfg: raise ValueError(f"Config '{cfg_name}' not in mykey")
-    cfg['_mykey_name'] = cfg_name
+    cfg = {**cfg, '_mykey_name': cfg_name}
     if 'native' in cfg_name: return (NativeClaudeSession if 'claude' in cfg_name else NativeOAISession)(cfg=cfg)
     if 'claude' in cfg_name: return ClaudeSession(cfg=cfg)
     return LLMSession(cfg=cfg) if 'oai' in cfg_name else None
 
-def resolve_client(cfg_name):
-    s = resolve_session(cfg_name)
+def resolve_client(cfg_name, cfg=None):
+    s = resolve_session(cfg_name, cfg)
     return (NativeToolClient(s) if isinstance(s, (NativeClaudeSession, NativeOAISession)) else ToolClient(s)) if s else None
 
 def fast_ask(prompt, cfg_name):