[AIGTWY-4756] Cache managed config reads for 5 min between launches

src/ucode/agents/__init__.py

@@ -85,16 +85,22 @@
 }
 
 
-def install_databricks_ai_tools_for_agents(tools: list[str], state: dict) -> None:
+def install_databricks_ai_tools_for_agents(
+    tools: list[str], state: dict, *, force_refresh: bool = False
+) -> None:
     """Install Databricks AI Tools for supported agents.
 
     Gemini and Pi have no ``aitools`` support and are dropped.
+
+    This runs only during ``ug configure``. ``force_refresh`` reads the managed config fresh; a
+    caller that already refreshed this launch (the main configure path) leaves it False so the gate
+    reuses that read instead of adding another control-plane round trip.
     """
     if not state.get("databricks_ai_tools_enabled"):
         return
     # An admin's managed config governs the workspace, so ucode does not
     # self-install AI Tools under one (may become a managed-config option later).
-    if refresh_managed_config(state).manifest is not None:
+    if refresh_managed_config(state, force_refresh=force_refresh).manifest is not None:
         return
     agents = [AITOOLS_AGENT_TOKENS[tool] for tool in tools if tool in AITOOLS_AGENT_TOKENS]
     if not agents:

src/ucode/cli.py

@@ -780,7 +780,9 @@ def configure_workspace_command(
         )
         state = states[0]
         state = configure_single_tool(tool, state)
-        install_databricks_ai_tools_for_agents([tool], state)
+        # No managed refresh precedes this branch, so read fresh here: `ug configure` never decides
+        # from a stale cache.
+        install_databricks_ai_tools_for_agents([tool], state, force_refresh=True)
         spec = TOOL_SPECS[tool]
         console.print(
             Panel(
@@ -807,8 +809,9 @@ def configure_workspace_command(
     save_state(state)
 
     # A published managed config means the admin dictates the setup: apply it to every enabled agent
-    # now rather than prompting the developer to pick.
-    managed, _ = refresh_managed_config(state)
+    # now rather than prompting the developer to pick. Configure always reads fresh so it never
+    # applies a config the admin has since changed.
+    managed, _ = refresh_managed_config(state, force_refresh=True)
     if managed is not None:
         _announce_managed_config(managed)
         for tool_name in managed_enabled_tools(managed):

src/ucode/config_io.py

@@ -170,7 +170,7 @@ def read_json_safe(path: Path) -> dict:
         if not path.exists():
             return {}
         data = json.loads(path.read_text(encoding="utf-8"))
-    except (OSError, json.JSONDecodeError):
+    except (OSError, UnicodeError, json.JSONDecodeError):
         return {}
     return data if isinstance(data, dict) else {}
 

src/ucode/managed_config.py

@@ -7,8 +7,9 @@
 - fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`),
 - normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names,
 - persisting it via :func:`save_managed_state` / :func:`load_managed_state` — the launch path pulls
-  the published copy into this file, and
-- re-reading it on each launch, falling back to the persisted copy when the read fails.
+  the published copy into this file, stamped with a ``retrieved_at`` and its outcome, and
+- re-reading it on each launch (reusing a read younger than :data:`MANAGED_CONFIG_TTL` rather than
+  re-fetching), falling back to the persisted copy when the read fails.
 
 There is deliberately one file: the workspace is the source of truth, so the pulled copy lives in
 ``managed-config.json`` and a launch re-reads it from there.
@@ -26,7 +27,7 @@
 import os
 import re
 from dataclasses import dataclass
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from typing import NamedTuple, cast
 
@@ -64,6 +65,17 @@
 
 MAX_SPEC_VERSION = 1
 
+# Launch-path cache TTL. `ug` / `ug <agent>` reuse a managed-config read younger than this instead of
+# hitting the control plane every launch; `ug configure` forces a fresh read (force_refresh=True).
+MANAGED_CONFIG_TTL = timedelta(minutes=5)
+
+# The last authoritative read's outcome, persisted alongside the config so a cached wrapper can be
+# replayed without a GET. "none" and "feature_disabled" both persist an empty config, so the outcome
+# is what tells them apart.
+_OUTCOME_PUBLISHED = "published"
+_OUTCOME_NONE = "none"
+_OUTCOME_FEATURE_DISABLED = "feature_disabled"
+
 # A per-family default-model key in the `default_models` map, e.g. `default_opus_model`. Matches the
 # server's `default_.+_model` validation so a new Claude family needs no ucode change. The bare
 # `default_model` overall default does not match (no family segment) and is read on its own.
@@ -604,7 +616,12 @@ def _is_unsupported_spec(reason: str) -> bool:
     return "spec_version" in reason.lower()
 
 
-def save_managed_state(workspace: str, config: dict) -> None:
+def _utcnow() -> datetime:
+    """Current UTC time. A seam so the launch-path cache TTL can be exercised deterministically."""
+    return datetime.now(UTC)
+
+
+def save_managed_state(workspace: str, config: dict, *, outcome: str | None = None) -> None:
     """Persist the raw managed config to ``~/.ucode/managed-config.json`` at mode 0600.
 
     ``config`` is stored verbatim as the gateway returned it (byte-identical to the GET), so the file
@@ -615,8 +632,15 @@ def save_managed_state(workspace: str, config: dict) -> None:
     An empty ``config`` records "this workspace has no managed config", which matters because the
     file doubles as the fallback when a later read fails: without it, removing a config server-side
     would leave the old one on disk to be reapplied after a transient outage.
+
+    ``outcome``, when set, stamps the read time and its result (published / none / feature_disabled)
+    so a later launch can reuse this read within :data:`MANAGED_CONFIG_TTL` without a GET. Only an
+    authoritative read passes it; a failed refresh persists nothing and so never advances the stamp.
     """
     payload: dict = {"workspace": workspace, "config": config}
+    if outcome is not None:
+        payload["retrieved_at"] = _utcnow().isoformat()
+        payload["outcome"] = outcome
     if config_io.is_dry_run():
         # Print rather than write, matching how the agent config writers behave under --dry-run.
         console.print(
@@ -681,15 +705,44 @@ def managed_state_workspace() -> str | None:
     return workspace if isinstance(workspace, str) and workspace else None
 
 
-def refresh_managed_config(state: dict) -> ManagedConfigResult:
-    """Fetch the workspace's managed config fresh and persist it as a :class:`ManagedConfigResult`.
+def _cached_result_if_fresh(workspace: str) -> ManagedConfigResult | None:
+    """The persisted read for ``workspace`` replayed as a result, if still within the TTL.
+
+    Returns None (forcing a fresh fetch) when the wrapper is for another workspace, predates this
+    cache format (no ``outcome`` / ``retrieved_at``), or its stamp is missing, unparseable, in the
+    future, or at least :data:`MANAGED_CONFIG_TTL` old.
+    """
+    data = config_io.read_json_safe(MANAGED_CONFIG_PATH)
+    if data.get("workspace") != workspace:
+        return None
+    # Reuses the RFC-3339 parser the update-time watermark uses; None (missing/unparseable) is stale.
+    retrieved_at = _parse_update_time(_str(data.get("retrieved_at")))
+    if retrieved_at is None:
+        return None
+    age = _utcnow() - retrieved_at
+    if age < timedelta(0) or age >= MANAGED_CONFIG_TTL:
+        return None
+    outcome = data.get("outcome")
+    if outcome == _OUTCOME_FEATURE_DISABLED:
+        return ManagedConfigResult(None, True)
+    if outcome == _OUTCOME_NONE:
+        return ManagedConfigResult(None, False)
+    if outcome == _OUTCOME_PUBLISHED and isinstance(data.get("config"), dict):
+        return ManagedConfigResult(normalize_managed_config(data["config"]), False)
+    return None
+
+
+def refresh_managed_config(state: dict, *, force_refresh: bool = False) -> ManagedConfigResult:
+    """Fetch the workspace's managed config and persist it as a :class:`ManagedConfigResult`.
 
     Runs on every launch so a developer picks up an admin's edits without re-running
-    ``ucode configure``. It always hits the control plane; whether the fetched config is *newer* than
-    what was last applied — and so whether the launch re-applies the settings — is the caller's
-    decision, via :func:`managed_config_is_newer` against the persisted applied watermark. The
-    manifest is None when the workspace has no managed config — the normal case for a workspace whose
-    admin hasn't published one.
+    ``ucode configure``. A launch reuses the last read when it is younger than
+    :data:`MANAGED_CONFIG_TTL`, so back-to-back launches don't each hit the control plane;
+    ``force_refresh`` (used by ``ug configure``) skips the cache and always reads fresh. Whether the
+    fetched config is *newer* than what was last applied, and so whether the launch re-applies the
+    settings, is the caller's decision, via :func:`managed_config_is_newer` against the persisted
+    applied watermark. The manifest is None when the workspace has no managed config, the normal
+    case for a workspace whose admin hasn't published one.
 
     A failed fetch never blocks the launch: an unreachable control plane shouldn't stop someone from
     coding. Instead it falls back to the last config persisted for this workspace, so the admin's
@@ -707,25 +760,29 @@ def refresh_managed_config(state: dict) -> ManagedConfigResult:
     workspace = state.get("workspace")
     if not workspace:
         return ManagedConfigResult(None, False)
+    if not force_refresh:
+        cached = _cached_result_if_fresh(workspace)
+        if cached is not None:
+            return cached
     try:
         token = get_databricks_token(workspace, state.get("profile"))
     except RuntimeError as exc:
         return ManagedConfigResult(_persisted_fallback(workspace, str(exc)), False)
     raw, reason = get_managed_config(workspace, token)
     if reason is not None:
         if _is_feature_disabled(reason):
-            save_managed_state(workspace, {})
+            save_managed_state(workspace, {}, outcome=_OUTCOME_FEATURE_DISABLED)
             return ManagedConfigResult(None, True)
         fallback = _persisted_fallback(workspace, reason, refused=_is_permission_denied(reason))
         return ManagedConfigResult(fallback, False)
     if raw is None:
         # Record that this workspace has no config, rather than leaving an earlier one on disk:
         # the file doubles as the fallback above, so a removed policy would otherwise come back
         # into force after the next transient outage.
-        save_managed_state(workspace, {})
+        save_managed_state(workspace, {}, outcome=_OUTCOME_NONE)
         return ManagedConfigResult(None, False)
     # Persist the raw config verbatim; hand callers the normalized manifest they expect.
-    save_managed_state(workspace, raw)
+    save_managed_state(workspace, raw, outcome=_OUTCOME_PUBLISHED)
     return ManagedConfigResult(normalize_managed_config(raw), False)
 
 

tests/test_agents_init.py

@@ -85,7 +85,9 @@ def _capture(self, monkeypatch, *, managed=None):
             lambda agents, profile: captured.update(agents=agents, profile=profile),
         )
         monkeypatch.setattr(
-            agents_mod, "refresh_managed_config", lambda state: ManagedConfigResult(managed, False)
+            agents_mod,
+            "refresh_managed_config",
+            lambda state, **_k: ManagedConfigResult(managed, False),
         )
         return captured
 
@@ -135,6 +137,23 @@ def test_skipped_under_empty_managed_config(self, monkeypatch):
         )
         assert captured == {}  # install_ai_tools never called
 
+    def test_forwards_force_refresh_to_managed_read(self, monkeypatch):
+        # The gate forwards force_refresh so `ug configure --agent` (no prior refresh) reads fresh,
+        # while the main configure path (already refreshed) reuses its read instead of re-fetching.
+        seen: list[bool] = []
+        monkeypatch.setattr(agents_mod, "install_ai_tools", lambda agents, profile: None)
+        monkeypatch.setattr(
+            agents_mod,
+            "refresh_managed_config",
+            lambda state, *, force_refresh=False: (
+                seen.append(force_refresh) or ManagedConfigResult(None, False)
+            ),
+        )
+        state = {"profile": "p", "databricks_ai_tools_enabled": True}
+        install_databricks_ai_tools_for_agents(["claude"], state)
+        install_databricks_ai_tools_for_agents(["claude"], state, force_refresh=True)
+        assert seen == [False, True]
+
 
 class TestConfigureWiresAiToolsInstall:
     """AI Tools install is a `ucode configure`-only step. `configure_selected_tools`
@@ -151,7 +170,9 @@ def _stub_configure(self, monkeypatch):
             lambda agents, profile: captured.update(agents=agents, profile=profile),
         )
         monkeypatch.setattr(
-            agents_mod, "refresh_managed_config", lambda state: ManagedConfigResult(None, False)
+            agents_mod,
+            "refresh_managed_config",
+            lambda state, **_k: ManagedConfigResult(None, False),
         )
         return captured
 

tests/test_cli.py

@@ -2651,7 +2651,7 @@ def _no_managed_config(self, monkeypatch):
         # `ug configure` now fetches the managed config, which shells out to the `databricks` CLI.
         # Default it to absent so these personal-flow tests never hit the CLI (it isn't on CI);
         # the managed-branch test overrides this.
-        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda state: (None, False))
+        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda state, **_k: (None, False))
 
     @pytest.mark.parametrize(("keys", "expected"), [(" \r", ["codex"]), ("\r", [])])
     def test_interactive_picker_installs_only_checked_agents(self, monkeypatch, keys, expected):
@@ -2763,7 +2763,7 @@ def test_managed_config_applies_all_enabled_and_skips_selection(self, monkeypatc
         monkeypatch.setattr(
             cli_mod,
             "refresh_managed_config",
-            lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False),
+            lambda s, **_k: ({"enabled_agents": {"claude": {}, "codex": {}}}, False),
         )
         monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True)
         monkeypatch.setattr(cli_mod, "resolve_state", lambda managed, s, tool: s)
@@ -2794,7 +2794,7 @@ def test_managed_config_registers_mcp_servers_after_configuring_agents(self, mon
             "enabled_agents": {"claude": {}, "codex": {}},
             "mcp_servers": {"names": ["x.y.z"]},
         }
-        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s: (managed, False))
+        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s, **_k: (managed, False))
         monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True)
         monkeypatch.setattr(cli_mod, "resolve_state", lambda m, s, tool: s)
         monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None)
@@ -2826,7 +2826,7 @@ def test_managed_configure_accumulates_available_tools_for_all_agents(self, monk
         monkeypatch.setattr(
             cli_mod,
             "refresh_managed_config",
-            lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False),
+            lambda s, **_k: ({"enabled_agents": {"claude": {}, "codex": {}}}, False),
         )
         monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True)
         # Mirror production: resolve_state hands each iteration a fresh copy of `state`.
@@ -2890,7 +2890,7 @@ def test_unmanaged_workspace_reconciles_managed_mcp_servers(self, monkeypatch):
 
         state = {**MINIMAL_STATE, "available_tools": []}
         monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state)
-        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s: (None, False))
+        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s, **_k: (None, False))
         monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: t == "claude")
         monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True)
         monkeypatch.setattr(
@@ -3590,7 +3590,7 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch):
 class TestConfigureNoLongerValidates:
     @pytest.fixture(autouse=True)
     def _no_managed_config(self, monkeypatch):
-        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda state: (None, False))
+        monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda state, **_k: (None, False))
 
     def test_configure_completes_without_probe(self, monkeypatch):
         import ucode.cli as cli_mod

tests/test_config_io.py

@@ -232,6 +232,13 @@ def test_read_json_safe_non_dict(self, tmp_path):
         p.write_text("[1, 2, 3]", encoding="utf-8")
         assert read_json_safe(p) == {}
 
+    def test_read_json_safe_non_utf8(self, tmp_path):
+        # A file with non-UTF-8 bytes must read as absent rather than raising, so a corrupted
+        # cache file can't crash a launch that reads it.
+        p = tmp_path / "binary.json"
+        p.write_bytes(b"\xff\xfe\x00not utf-8")
+        assert read_json_safe(p) == {}
+
     def test_read_toml_safe_missing_file(self, tmp_path):
         doc = read_toml_safe(tmp_path / "missing.toml")
         assert dict(doc) == {}