From 74fa448bfdec48173e9844fcb3dbc216e42facb1 Mon Sep 17 00:00:00 2001
From: Andy Xu <andy.xu@databricks.com>
Date: Wed, 16 Sep 2026 21:57:03 +0000
Subject: [PATCH] Define scoped model discovery policy

---
 src/ucode/constants.py  | 25 +++++++++++++++++++++++++
 tests/test_constants.py | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 tests/test_constants.py

diff --git a/src/ucode/constants.py b/src/ucode/constants.py
index c129da93..ba8fe18e 100644
--- a/src/ucode/constants.py
+++ b/src/ucode/constants.py
@@ -1,11 +1,36 @@
 """Shared UCode constants."""
 
+import os
+
 LOCALHOST = "localhost"
 LOOPBACK_HOST = "127.0.0.1"
 
 MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service"
 MODEL_SERVICE_PARENT_SCHEMA_HEADER = "Databricks-Model-Service-Parent-Schema"
 
+# Public policy switch for agent-owned catalogs scoped by a Model Provider
+# Service or Unity Catalog location. Ordinary ``system.ai`` discovery is
+# intentionally independent of this flag.
+MODEL_DISCOVERY_ENV_VAR = "UG_ENABLE_MODEL_DISCOVERY"
+
+
+def scoped_model_discovery_enabled(
+    *,
+    override: bool | None = None,
+    force: bool = False,
+) -> bool:
+    """Whether provider/location-scoped agent discovery should be enabled.
+
+    ``override`` is transient launch state, not persisted configuration. Managed
+    workspace policy uses ``force`` because it outranks a developer's environment.
+    """
+    if force:
+        return True
+    if override is not None:
+        return override
+    return os.environ.get(MODEL_DISCOVERY_ENV_VAR) != "0"
+
+
 # MCP server registration scopes. Claude Code supports local/project/user; the
 # other CLIs only take the user-scope name. Kept here (a leaf module) so both
 # `ucode.mcp` and `ucode.agents.claude` can import them without an import cycle.
diff --git a/tests/test_constants.py b/tests/test_constants.py
new file mode 100644
index 00000000..29868a62
--- /dev/null
+++ b/tests/test_constants.py
@@ -0,0 +1,38 @@
+import pytest
+
+from ucode.constants import MODEL_DISCOVERY_ENV_VAR, scoped_model_discovery_enabled
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (None, True),
+        ("", True),
+        ("1", True),
+        ("false", True),
+        ("0", False),
+    ],
+)
+def test_scoped_model_discovery_environment_policy(monkeypatch, value, expected):
+    if value is None:
+        monkeypatch.delenv(MODEL_DISCOVERY_ENV_VAR, raising=False)
+    else:
+        monkeypatch.setenv(MODEL_DISCOVERY_ENV_VAR, value)
+
+    assert scoped_model_discovery_enabled() is expected
+
+
+@pytest.mark.parametrize(
+    ("environment", "override", "expected"),
+    [("0", True, True), ("1", False, False)],
+)
+def test_scoped_model_discovery_override_wins(monkeypatch, environment, override, expected):
+    monkeypatch.setenv(MODEL_DISCOVERY_ENV_VAR, environment)
+
+    assert scoped_model_discovery_enabled(override=override) is expected
+
+
+def test_scoped_model_discovery_force_wins(monkeypatch):
+    monkeypatch.setenv(MODEL_DISCOVERY_ENV_VAR, "0")
+
+    assert scoped_model_discovery_enabled(override=False, force=True) is True