diff --git a/anton/cloud_turn/session.py b/anton/cloud_turn/session.py index 6e22e588f..52b51270b 100644 --- a/anton/cloud_turn/session.py +++ b/anton/cloud_turn/session.py @@ -799,6 +799,7 @@ def build_cloud_chat_session( from anton.utils.datasources import restore_namespaced_env data_vault = TurnKeyDataVault(request.oauth) + data_vault.prefetch() restore_namespaced_env(data_vault) # ENG-1816: discover any MCP-method connections' tools now, before diff --git a/anton/core/datasources/data_vault.py b/anton/core/datasources/data_vault.py index 83bbfd1e1..5c1fb74a6 100644 --- a/anton/core/datasources/data_vault.py +++ b/anton/core/datasources/data_vault.py @@ -5,6 +5,7 @@ import os import re import urllib.error +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from pathlib import Path from typing import Any, Protocol, runtime_checkable @@ -391,6 +392,10 @@ def clear_ds_env(self) -> None: _TURNKEY_EXTRA_FIELDS = ("developer_token", "login_customer_id") _TURNKEY_SECURE_KEYS = ("access_token", "developer_token") +#: Upper bound on simultaneous token fetches while a turn's session is built, +#: so a turn with many connectors does not burst auth. +_PREFETCH_MAX_WORKERS = 4 + class TurnKeyDataVault: """Read-only DataVault backed by a live call to auth's turn-key token @@ -493,6 +498,33 @@ def clear_ds_env(self) -> None: _sweep_ds_env_vars() self._cache.clear() + def prefetch(self) -> None: + """Fetch every listed connection not yet cached, a few at a time. + + Called once while a turn's session is built, so the later reads + (registry rebuild, MCP discovery, prompt) hit the cache instead of paying + one auth round trip per connection in sequence. Workers only fetch; the + calling thread writes the cache, and no context or process state is + touched. Returns after every fetch finished, successful or not. + """ + pending = [ + key + for key in dict.fromkeys((c["engine"], c["name"]) for c in self._connections) + if key not in self._cache + ] + if not pending: + return + workers = min(_PREFETCH_MAX_WORKERS, len(pending)) + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="turnkey-prefetch") as pool: + futures = {key: pool.submit(self._fetch_uncached, *key) for key in pending} + for (engine, name), future in futures.items(): + try: + self._cache[(engine, name)] = future.result() + except Exception: + # Only an unexpected error lands here; auth and parse failures already + # came back as None and are cached. Left uncached so a later read retries. + logger.warning("turn-key prefetch failed for %s/%s", engine, name, exc_info=True) + def _fetch(self, engine: str, name: str) -> dict[str, str] | None: cache_key = (engine, name) if cache_key in self._cache: diff --git a/anton/core/memory/cortex.py b/anton/core/memory/cortex.py index 1fd3ea7d1..d34866b1d 100644 --- a/anton/core/memory/cortex.py +++ b/anton/core/memory/cortex.py @@ -22,7 +22,9 @@ from __future__ import annotations import asyncio +import json import logging +import re import time from dataclasses import replace from pathlib import Path @@ -33,6 +35,7 @@ from anton.core.llm.tracing import tagged_trace from anton.core.llm.structured import ( generate_with_truncation_retry, + looks_truncated, no_preamble_instruction, ) from anton.core.memory.base import ( @@ -69,6 +72,8 @@ # observation to count). _RULE_RETRIEVAL_TAG = "rule-retrieval" _RULE_RETRIEVAL_MAX_TOKENS = 4096 +# The filter blocks the first reply token; past this, every conditional rule is kept. +_RULE_RETRIEVAL_TIMEOUT_SECONDS = 3.0 def learned_memory_notice( @@ -135,6 +140,36 @@ def _emit_rule_retrieval(**shape) -> None: logger.debug("rule_retrieval analytics event failed", exc_info=True) +# Models often fence JSON even when told not to; one surrounding fence is unwrapped. +_JSON_FENCE_RE = re.compile(r"\A```(?:json)?[ \t]*\n(.*)\n[ \t]*```\Z", re.DOTALL | re.IGNORECASE) + + +def _parse_rule_selection(text: str, count: int) -> list[int] | None: + """Parse the rule filter's `{"keep": [1, 3]}` answer into 0-based indices. + + `count` is how many rules were numbered (from 1). One surrounding code fence + is removed first. Returns the sorted, de-duplicated indices, an empty list + when the model chose none, or None when the answer is not exactly that shape + or names a rule that was not shown. + """ + text = text.strip() + fenced = _JSON_FENCE_RE.match(text) + if fenced: + text = fenced.group(1) + try: + data = json.loads(text) + except ValueError: + return None + if not isinstance(data, dict) or set(data) != {"keep"} or not isinstance(data["keep"], list): + return None + selected: set[int] = set() + for item in data["keep"]: + if type(item) is not int or not 1 <= item <= count: + return None + selected.add(item - 1) + return sorted(selected) + + # ───────────────────────────────────────────────────────────────────────────── # Pydantic schemas — used by LLMClient.generate_object # ───────────────────────────────────────────────────────────────────────────── @@ -355,10 +390,10 @@ def _move_project_identity(self) -> None: _LESSONS_BUDGET_TOKENS = 1000 _RULES_RETRIEVAL_PROMPT = """\ -Given the user's current message, select only the conditional (When/If) rules that are \ -relevant. Return the selected rules exactly as they appear, one per line (keep the "- " prefix). -If all rules are relevant, return them all. If none are relevant, return "NONE". -Do NOT add, modify, or summarize rules — return them verbatim. +Given the user's current message, select only the numbered conditional (When/If) rules \ +that are relevant. Answer with JSON only, exactly {"keep": []}, listing the \ +numbers of the relevant rules, for example {"keep": [1, 3]}. If none are relevant, \ +answer {"keep": []}. Do not repeat the rule text. """ # Rendering first doesn't tell the model which tier wins, so the first @@ -545,7 +580,8 @@ async def _retrieve_relevant_rules( Always/Never rules are behavioral constraints — always loaded in full. Only conditional (When/If) rules are filtered by relevance. - If rules are under budget or no LLM is available, returns as-is. + If rules are under budget or no LLM is available, returns as-is. A slow, + malformed, or truncated selection keeps every conditional rule. """ if not user_message or self._llm is None: return engrams @@ -562,15 +598,11 @@ async def _retrieve_relevant_rules( when_text = "\n".join(f"- {e.text}" for e in when_engrams) if len(when_text) < 1000: return engrams + numbered_text = "\n".join(f"{i}. {e.text}" for i, e in enumerate(when_engrams, start=1)) - # ENG-1390 instrumentation. Purely observational — every branch below - # produces exactly the value it produced before. What is new is that the - # branches are now DISTINGUISHABLE: `filtered_when = when_engrams` is - # reached three different ways (no exact match, empty response, exception) - # and all three are byte-identical in their effect on the prompt, so a - # count of invocations could never tell a working filter from a broken one. - # Measured shares before this shipped: 55.7% filtered, 30.2% dropped_all, - # 14.2% kept_all_no_match, plus an exception rate that was unobservable. + # Every keep-all fallback below (unusable answer, empty response, timeout, + # exception) is byte-identical in its effect on the prompt, so only + # `outcome` tells a working filter from a broken one. outcome = "error" kept = len(when_engrams) stop_reason = "" @@ -580,36 +612,44 @@ async def _retrieve_relevant_rules( try: # Isolate this call from every other `llm.code` caller in the trace. with tagged_trace(_RULE_RETRIEVAL_TAG): - response = await self._llm.code( - system=self._RULES_RETRIEVAL_PROMPT, - messages=[ - { - "role": "user", - "content": f"User message: {user_message}\n\nRules:\n{when_text}", - } - ], - max_tokens=_RULE_RETRIEVAL_MAX_TOKENS, + response = await asyncio.wait_for( + self._llm.code( + system=self._RULES_RETRIEVAL_PROMPT, + messages=[ + { + "role": "user", + "content": f"User message: {user_message}\n\nRules:\n{numbered_text}", + } + ], + max_tokens=_RULE_RETRIEVAL_MAX_TOKENS, + ), + timeout=_RULE_RETRIEVAL_TIMEOUT_SECONDS, ) usage = getattr(response, "usage", None) stop_reason = getattr(response, "stop_reason", None) or "" result = response.content.strip() - if result.upper() == "NONE": - # A fourth outcome the ticket did not name: EVERY conditional rule - # is discarded. The opposite of the permissive fallback, and ~30% - # of calls. + selection = None + if result and not looks_truncated(response, _RULE_RETRIEVAL_MAX_TOKENS): + selection = _parse_rule_selection(result, len(when_engrams)) + if not result: + outcome, filtered_when = "kept_all_empty", when_engrams + elif selection is None: + # Malformed, out of range, or cut at the ceiling: a partial list + # would silently drop rules, so keep everything instead. + outcome, filtered_when = "kept_all_no_match", when_engrams + elif not selection: + # The opposite of the permissive fallback: EVERY conditional rule + # is discarded. outcome, kept, early = "dropped_all", 0, mandatory - elif result: - returned = {line.lstrip("- ").strip() for line in result.splitlines() if line.strip()} - filtered_when = [e for e in when_engrams if e.text in returned] - if not filtered_when: - # Exact string equality against the model's echo. Any rewrap, - # added punctuation, or a response truncated at the ceiling - # yields zero matches — and then we silently keep everything. - outcome, filtered_when = "kept_all_no_match", when_engrams - else: - outcome, kept = "filtered", len(filtered_when) else: - outcome, filtered_when = "kept_all_empty", when_engrams + # Original engrams in file order; model text never becomes a rule. + filtered_when = [when_engrams[i] for i in selection] + outcome, kept = "filtered", len(filtered_when) + except TimeoutError: + filtered_when = when_engrams + # A provider's own timeout is an error, not a deadline hit. + reached = time.monotonic() - started >= _RULE_RETRIEVAL_TIMEOUT_SECONDS + outcome = "timeout" if reached else "error" except asyncio.CancelledError: # A user pressing STOP, or an abandoned SSE stream, cancels the turn # mid-call. `CancelledError` is a BaseException, so `except Exception` @@ -630,11 +670,8 @@ async def _retrieve_relevant_rules( when_rules=len(when_engrams), kept_rules=kept, rules_chars=len(when_text), - # `stop_reason`, not output_tokens == max_tokens: a truncated - # verbatim echo is silently accepted as "the relevant rules", so - # this is the difference between a real filter and one that just - # ran out of room. The provider reason is authoritative where the - # exact-cap heuristic gives false positives. + # A cap hit can still report "stop", so `stop_reason` and + # `output_tokens` together separate it from a malformed selection. stop_reason=stop_reason, input_tokens=getattr(usage, "input_tokens", 0) or 0, output_tokens=getattr(usage, "output_tokens", 0) or 0, diff --git a/tests/test_cloud_turn_session.py b/tests/test_cloud_turn_session.py index aa4c471f0..0ef0fa423 100644 --- a/tests/test_cloud_turn_session.py +++ b/tests/test_cloud_turn_session.py @@ -15,6 +15,7 @@ import subprocess import sys import textwrap +import threading from pathlib import Path # Bound at import so the constants below are built from the real class, which # `_pin_today` later replaces on the module. @@ -34,6 +35,7 @@ resolve_trusted_workspace_path, ) from anton.core.backends.local import local_scratchpad_runtime_factory +from anton.utils.datasources import scrub_credentials class _FakeSession: @@ -1280,6 +1282,34 @@ def fake_minds_request(url, api_key, *, method="GET", payload=None, verify=True, assert "DS_GMAIL_PRIMARY__ACCESS_TOKEN" not in os.environ +def test_oauth_connections_are_fetched_together_once_and_registered_for_scrubbing(tmp_path, monkeypatch): + """Each connection's fetch waits for the other three, so the build only + gets tokens when the fetches overlap; every value must then be scrubbable.""" + barrier = threading.Barrier(4) + calls: list[str] = [] + lock = threading.Lock() + + def fake_minds_request(url, api_key, *, method="GET", payload=None, verify=True, timeout=30): + name = json.loads(payload)["name"] + with lock: + calls.append(name) + barrier.wait(timeout=5) + return json.dumps({"access_token": f"live-token-{name}"}).encode() + + monkeypatch.setattr("anton.minds_client.minds_request", fake_minds_request) + + names = ["alpha", "beta", "gamma", "delta"] + oauth = {"turn_key": "tk_abc", "connections": [{"engine": "google_drive", "name": n} for n in names]} + _, cfg = _build(tmp_path, monkeypatch, oauth=oauth) + + assert sorted(calls) == sorted(names) + for n in names: + token = f"live-token-{n}" + assert cfg.data_vault.load("google_drive", n)["access_token"] == token + assert token not in scrub_credentials(f"leaked {token}") + assert len(calls) == len(names) + + # ── ENG-1638: the pod injects a configured-LLM block ───────────────────────── def test_pod_injects_a_configured_llm_block_for_the_runtime_identity(tmp_path, monkeypatch): diff --git a/tests/test_cortex.py b/tests/test_cortex.py index c48198688..be9de3adb 100644 --- a/tests/test_cortex.py +++ b/tests/test_cortex.py @@ -11,10 +11,12 @@ import pytest +import anton.core.memory.cortex as cortex_mod from anton.core.backends.base import Cell +from anton.core.llm.provider import LLMResponse, Usage from anton.core.memory.cerebellum import Cerebellum, _DiffPassResult, _LessonDraft from anton.core.memory.base import MemoryRead -from anton.core.memory.cortex import Cortex, _CompactionResult +from anton.core.memory.cortex import Cortex, _CompactionResult, _parse_rule_selection from anton.core.memory.hippocampus import Engram, Hippocampus from anton.core.session import ChatSession from anton.core.tools.tool_handlers import handle_memorize @@ -246,7 +248,7 @@ async def test_counts_the_rules_the_filter_keeps(self, cortex, dirs, captured_ev rules = _when_engrams(12) lines = "\n".join(f"- {e.text}" for e in rules) (g / "rules.md").write_text(f"## When\n{lines}\n") - cortex._llm = _FakeLLM(content="\n".join(f"- {rules[i].text}" for i in (0, 1))) + cortex._llm = _FakeLLM(content='{"keep": [1, 2]}') await cortex.build_memory_context("please do the august campaign") @@ -617,22 +619,20 @@ async def _run(self, cortex, llm, rules): async def test_filtered_outcome_is_reported(self, cortex, captured_events): rules = _when_engrams(12) - # Echo back exactly two rules, verbatim, as the prompt demands. - echo = "\n".join(f"- {rules[i].text}" for i in (0, 1)) - llm = _FakeLLM(content=echo) + llm = _FakeLLM(content='{"keep": [1, 2]}') out = await self._run(cortex, llm, rules) assert llm.calls == 1 - assert len(out) == 2, "behaviour must be unchanged: only the echoed rules survive" + assert out == rules[:2], "only the selected rules survive, as the original engrams" ev = _one_event(captured_events) assert ev["outcome"] == "filtered" assert ev["when_rules"] == 12 and ev["kept_rules"] == 2 async def test_dropped_all_is_its_own_outcome(self, cortex, captured_events): - """`NONE` discards every conditional rule — measured at 30% and absent - from the ticket's original three-outcome taxonomy.""" + """An empty selection discards every conditional rule, an outcome + separate from the original three.""" rules = _when_engrams(12) - out = await self._run(cortex, _FakeLLM(content="NONE"), rules) + out = await self._run(cortex, _FakeLLM(content='{"keep": []}'), rules) assert out == [], "no mandatory rules present, so nothing should remain" ev = _one_event(captured_events) assert ev["outcome"] == "dropped_all" @@ -668,19 +668,126 @@ async def test_exception_is_reported_at_all(self, cortex, captured_events): assert len(out) == 12, "the permissive fallback must still apply" assert _one_event(captured_events)["outcome"] == "error" - async def test_truncation_is_reported_via_stop_reason(self, cortex, captured_events): - """A verbatim echo cut at the ceiling is accepted as 'the relevant rules', - so truncation masquerades as a successful filter — 8 of 9 truncated calls - in production were classified that way. `stop_reason` is what separates - them; output_tokens == max_tokens is a heuristic with false positives.""" + async def test_a_selection_at_the_output_cap_keeps_every_rule_even_on_stop(self, cortex, captured_events): + """A gateway may report "stop" at the cap, so the token count alone rejects it.""" + + class _AtTheCap: + async def code(self, **kwargs): + return LLMResponse( + content='{"keep": [1, 2]}', + usage=Usage(output_tokens=cortex_mod._RULE_RETRIEVAL_MAX_TOKENS), + stop_reason="stop", + ) + + rules = _when_engrams(12) + out = await self._run(cortex, _AtTheCap(), rules) + assert out == rules + assert _one_event(captured_events)["outcome"] == "kept_all_no_match" + + async def test_truncated_selection_keeps_every_rule(self, cortex, captured_events): + """A selection cut at the ceiling may be missing rules even when what + arrived parses, so it is rejected whole rather than trusted.""" rules = _when_engrams(12) - echo = "\n".join(f"- {rules[i].text}" for i in (0, 1)) - llm = _FakeLLM(content=echo, stop_reason="max_tokens") - await self._run(cortex, llm, rules) + llm = _FakeLLM(content='{"keep": [1, 2]}', stop_reason="max_tokens") + out = await self._run(cortex, llm, rules) + assert out == rules ev = _one_event(captured_events) - assert ev["outcome"] == "filtered" + assert ev["outcome"] == "kept_all_no_match" assert ev["stop_reason"] == "max_tokens", "a truncated filter must be separable" + async def test_the_model_is_shown_numbered_rules_and_asked_for_numbers(self, cortex): + rules = _when_engrams(12) + seen: dict = {} + + class _Capturing: + async def code(self, **kwargs): + seen.update(kwargs) + return LLMResponse(content='{"keep": [3]}', usage=Usage()) + + cortex._llm = _Capturing() + out = await cortex._retrieve_relevant_rules(rules, "do the august campaign") + content = seen["messages"][0]["content"] + assert f"1. {rules[0].text}" in content + assert f"12. {rules[11].text}" in content + assert '{"keep"' in seen["system"] + assert out == [rules[2]] + + async def test_selection_keeps_file_order_and_collapses_duplicates(self, cortex, captured_events): + rules = _when_engrams(12) + out = await self._run(cortex, _FakeLLM(content='{"keep": [5, 2, 5]}'), rules) + assert out == [rules[1], rules[4]] + assert _one_event(captured_events)["kept_rules"] == 2 + + async def test_a_fenced_selection_still_filters(self, cortex, captured_events): + rules = _when_engrams(12) + out = await self._run(cortex, _FakeLLM(content='```json\n{"keep": [2]}\n```'), rules) + assert out == [rules[1]] + assert _one_event(captured_events)["outcome"] == "filtered" + + async def test_one_invalid_index_keeps_every_rule(self, cortex, captured_events): + rules = _when_engrams(12) + out = await self._run(cortex, _FakeLLM(content='{"keep": [1, 13]}'), rules) + assert out == rules + assert _one_event(captured_events)["outcome"] == "kept_all_no_match" + + async def test_mandatory_rules_survive_every_selection(self, cortex): + mandatory = [ + Engram(text="Always cite sources", kind="always", scope="global", confidence="high"), + Engram(text="Never delete data", kind="never", scope="global", confidence="high"), + ] + when = _when_engrams(12) + for content in ('{"keep": [1]}', '{"keep": []}', "garbage", ""): + cortex._llm = _FakeLLM(content=content) + out = await cortex._retrieve_relevant_rules(mandatory + when, "do the august campaign") + assert out[:2] == mandatory, content + + async def test_a_slow_filter_times_out_and_keeps_every_rule(self, cortex, captured_events, monkeypatch): + monkeypatch.setattr(cortex_mod, "_RULE_RETRIEVAL_TIMEOUT_SECONDS", 0.05) + + class _Hanging: + async def code(self, **kwargs): + await asyncio.sleep(30) + + rules = _when_engrams(12) + cortex._llm = _Hanging() + started = asyncio.get_running_loop().time() + out = await cortex._retrieve_relevant_rules(rules, "do the august campaign") + assert asyncio.get_running_loop().time() - started < 5, "must not wait for the 30 s hang" + assert out == rules + assert _one_event(captured_events)["outcome"] == "timeout" + + async def test_a_provider_timeout_before_the_deadline_is_an_error(self, cortex, captured_events): + class _ProviderTimeout: + async def code(self, **kwargs): + raise TimeoutError("read timed out") + + rules = _when_engrams(12) + cortex._llm = _ProviderTimeout() + out = await cortex._retrieve_relevant_rules(rules, "do the august campaign") + assert out == rules + assert _one_event(captured_events)["outcome"] == "error" + + async def test_both_scopes_time_out_together(self, dirs, captured_events, monkeypatch): + """Prompt assembly waits for one deadline, not one per memory scope.""" + monkeypatch.setattr(cortex_mod, "_RULE_RETRIEVAL_TIMEOUT_SECONDS", 0.2) + for d in dirs: + lines = "\n".join(f"- {e.text}" for e in _when_engrams(12)) + (d / "rules.md").write_text(f"## When\n{lines}\n") + + class _Hanging: + async def code(self, **kwargs): + await asyncio.sleep(30) + + g, p = dirs + cortex = Cortex( + global_hc=Hippocampus(g), project_hc=Hippocampus(p), mode="copilot", llm_client=_Hanging() + ) + started = asyncio.get_running_loop().time() + result = await cortex.build_memory_context("do the august campaign") + assert asyncio.get_running_loop().time() - started < 5, "must not wait for the 30 s hang" + assert [ev["outcome"] for ev in captured_events] == ["timeout", "timeout"] + assert result.count("When condition 011 applies") == 2 + async def test_the_call_is_tagged_so_it_is_isolatable_in_a_trace(self, cortex): """Done-when #1: distinguishable from every other `llm.code` call.""" from anton.core.llm.tracing import ( @@ -1136,6 +1243,51 @@ async def test_entries_under_an_unknown_heading_are_not_lost(self, dirs): assert len(hc.get_rules()) == 11 +class TestParseRuleSelection: + @pytest.mark.parametrize( + "text, expected", + [ + ('{"keep": [1, 3]}', [0, 2]), + ('{"keep": []}', []), + ('{"keep": [3, 1, 3]}', [0, 2]), + (' {"keep": [2]} ', [1]), + ('```json\n{"keep": [1, 3]}\n```', [0, 2]), + ('```\n{"keep": []}\n```', []), + (' ```JSON\n{"keep": [2]}\n``` ', [1]), + ], + ) + def test_accepts_exact_index_selections(self, text, expected): + assert _parse_rule_selection(text, 3) == expected + + @pytest.mark.parametrize( + "text", + [ + "", + "NONE", + "not json", + '{"keep": [1, 2]', + "[1, 2]", + '{"keep": "1"}', + '{"keep": [1], "why": "x"}', + '{"rules": [1]}', + '{"keep": [0]}', + '{"keep": [4]}', + '{"keep": [1, 4]}', + '{"keep": [true]}', + '{"keep": [1.0]}', + '{"keep": ["1"]}', + '{"keep": [null]}', + '```json\n{"keep": [4]}\n```', + '```json\n```json\n{"keep": [1]}\n```\n```', + '```json\n{"keep": [1]}', + 'Here you go: ```json\n{"keep": [1]}\n```', + '```python\n{"keep": [1]}\n```', + ], + ) + def test_rejects_anything_else(self, text): + assert _parse_rule_selection(text, 3) is None + + # ───────────────────────────────────────────────────────────────────────────── # The operator's read-only layer # ───────────────────────────────────────────────────────────────────────────── @@ -1447,7 +1599,7 @@ async def test_long_operator_when_rules_are_filtered( managed = read_only_dir( "managed", {"rules.md": f"## Always\n- Keep it short\n## When\n{lines}\n"} ) - llm = _FakeLLM(content="\n".join(f"- {rules[i].text}" for i in (0, 1))) + llm = _FakeLLM(content='{"keep": [1, 2]}') cortex = _managed_cortex(dirs, managed, llm_client=llm) prompt = await cortex.build_memory_context("do the august campaign") diff --git a/tests/test_turn_key_data_vault.py b/tests/test_turn_key_data_vault.py index 6939a6bd3..2755df9f4 100644 --- a/tests/test_turn_key_data_vault.py +++ b/tests/test_turn_key_data_vault.py @@ -3,6 +3,8 @@ """ from __future__ import annotations +import json +import threading import urllib.error import pytest @@ -279,3 +281,142 @@ def test_developer_token_is_marked_secret_and_login_customer_id_is_not(monkeypat monkeypatch.setattr("anton.minds_client.minds_request", lambda *a, **kw: _GOOGLE_ADS_RESPONSE) record = _google_ads_vault().read_record("google_ads", "primary") assert record["secure_keys"] == ["access_token", "developer_token"] + + +# ── prefetch ───────────────────────────────────────────────────────────────── + + +def _vault_with(count: int) -> TurnKeyDataVault: + return _vault(connections=[{"engine": "google_drive", "name": f"c{i}"} for i in range(count)]) + + +class _CountingAuth: + """Fake token endpoint that records each fetch and the peak overlap.""" + + def __init__(self, barrier: threading.Barrier | None = None, fail_names: tuple[str, ...] = ()) -> None: + self.calls: list[str] = [] + self.peak = 0 + self._active = 0 + self._lock = threading.Lock() + self._barrier = barrier + self._fail_names = fail_names + + def __call__(self, url, api_key, *, method="GET", payload=None, verify=True, timeout=30): + name = json.loads(payload)["name"] + with self._lock: + self.calls.append(name) + self._active += 1 + self.peak = max(self.peak, self._active) + try: + if self._barrier is not None: + self._barrier.wait(timeout=5) + if name in self._fail_names: + raise urllib.error.HTTPError(url, 403, "Forbidden", {}, None) + return json.dumps({"access_token": f"tok-{name}"}).encode() + finally: + with self._lock: + self._active -= 1 + + +def test_prefetch_with_no_connections_makes_no_call(monkeypatch): + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + _vault(connections=[]).prefetch() + assert auth.calls == [] + + +def test_prefetch_of_a_single_connection_caches_it_with_one_fetch(monkeypatch): + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault_with(1) + vault.prefetch() + assert auth.calls == ["c0"] + assert vault.load("google_drive", "c0")["access_token"] == "tok-c0" + assert auth.calls == ["c0"] + + +def test_prefetch_runs_fetches_concurrently_and_caches_every_result(monkeypatch): + """Four fetches that each wait for the other three only complete when they + overlap; a sequential fetch would break the barrier and cache no token.""" + auth = _CountingAuth(barrier=threading.Barrier(4)) + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault_with(4) + vault.prefetch() + assert sorted(auth.calls) == ["c0", "c1", "c2", "c3"] + for i in range(4): + assert vault.load("google_drive", f"c{i}")["access_token"] == f"tok-c{i}" + assert len(auth.calls) == 4 + + +def test_prefetch_never_runs_more_than_four_fetches_at_once(monkeypatch): + auth = _CountingAuth(barrier=threading.Barrier(4)) + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault_with(8) + vault.prefetch() + assert sorted(auth.calls) == sorted(f"c{i}" for i in range(8)) + assert auth.peak == 4 + for i in range(8): + assert vault.load("google_drive", f"c{i}")["access_token"] == f"tok-c{i}" + assert len(auth.calls) == 8 + + +def test_prefetch_fetches_a_duplicated_connection_once(monkeypatch): + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault(connections=[ + {"engine": "google_drive", "name": "a"}, + {"engine": "google_drive", "name": "a"}, + {"engine": "google_drive", "name": "b"}, + ]) + vault.prefetch() + assert sorted(auth.calls) == ["a", "b"] + + +def test_prefetch_skips_connections_already_cached(monkeypatch): + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault_with(3) + vault.load("google_drive", "c1") + vault.prefetch() + assert sorted(auth.calls) == ["c0", "c1", "c2"] + + +def test_prefetch_caches_a_needs_reconnect_failure_like_a_plain_load(monkeypatch): + auth = _CountingAuth(fail_names=("c1",)) + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault_with(3) + vault.prefetch() + assert vault.load("google_drive", "c1") is None + assert vault.load("google_drive", "c0")["access_token"] == "tok-c0" + assert len(auth.calls) == 3 + + +def test_prefetch_without_a_turn_key_makes_no_call(monkeypatch): + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + vault = _vault(turn_key="", connections=[{"engine": "google_drive", "name": f"c{i}"} for i in range(3)]) + vault.prefetch() + assert auth.calls == [] + assert vault.load("google_drive", "c0") is None + + +def test_a_prefetch_worker_that_raises_leaves_the_connection_to_a_later_read(monkeypatch): + """An unexpected error in one worker must not fail the others or cache a + miss; the first real read fetches that connection again.""" + vault = _vault_with(3) + real = vault._fetch_uncached + failed = [] + + def flaky(engine, name): + if name == "c1" and not failed: + failed.append(name) + raise RuntimeError("unexpected") + return real(engine, name) + + auth = _CountingAuth() + monkeypatch.setattr("anton.minds_client.minds_request", auth) + monkeypatch.setattr(vault, "_fetch_uncached", flaky) + vault.prefetch() + assert sorted(auth.calls) == ["c0", "c2"] + assert vault.load("google_drive", "c1")["access_token"] == "tok-c1" + assert sorted(auth.calls) == ["c0", "c1", "c2"]