fix(python): ship credential driver protocol bindings (#4760)

Closes #4759

Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>

Author: maxdubrinsky

docs/sdk/python.mdx

@@ -148,6 +148,23 @@ workspace for repeated exec, stop, start, and delete operations. List methods
 return lazy `Pager` instances; use `list_all()` only when you want to fetch the
 complete collection.
 
+## Implement a Credential Driver
+
+The SDK includes generated message classes, type stubs, and gRPC bindings for
+the credential-driver and extension protocols:
+
+```python
+from openshell._proto import credential_driver_pb2, credential_driver_pb2_grpc
+from openshell._proto import datamodel_pb2, extension_pb2
+```
+
+Subclass `credential_driver_pb2_grpc.CredentialDriverServicer` and register your
+implementation with `add_CredentialDriverServicer_to_server()` from the same
+module. Clients use `CredentialDriverStub`. The generated messages share
+`datamodel_pb2.CredentialHandle` and `extension_pb2.PeerMetadata` with the SDK.
+Use bindings from the same OpenShell release as your gateway; these modules
+expose the internal credential-driver contract.
+
 ## Next Steps
 
 - Review [Manage Sandboxes](/how-it-works/sandboxes/overview) for labels, templates, services, and Python SDK examples.

python/credential_driver_proto_test.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Credential-driver bindings must share SDK datamodel types and dispatch RPCs."""
+
+import grpc.aio
+import pytest
+
+from openshell._proto import credential_driver_pb2 as credentials
+from openshell._proto import credential_driver_pb2_grpc as driver_grpc
+from openshell._proto import datamodel_pb2 as datamodel
+from openshell._proto import extension_pb2 as extension
+from openshell._proto import openshell_pb2 as gateway
+
+
+def test_driver_messages_share_sdk_descriptors() -> None:
+    request = credentials.ResolveCredentialRequest(
+        handle=datamodel.CredentialHandle(driver="test", handle="opaque"),
+    )
+    assert isinstance(request.handle, datamodel.CredentialHandle)
+    assert request.handle.DESCRIPTOR is datamodel.CredentialHandle.DESCRIPTOR
+    assert (
+        credentials.GetCredentialDriverCapabilitiesRequest(
+            gateway=extension.PeerMetadata()
+        ).gateway.DESCRIPTOR
+        is extension.PeerMetadata.DESCRIPTOR
+    )
+    assert (
+        gateway.CreateProviderRequest().provider.DESCRIPTOR
+        is datamodel.Provider.DESCRIPTOR
+    )
+
+
+@pytest.mark.asyncio
+async def test_generated_servicer_and_stub_exchange() -> None:
+    class Driver(driver_grpc.CredentialDriverServicer):
+        async def StoreCredential(self, request, context):  # noqa: ARG002
+            assert request.value == "registration-ticket"
+            return credentials.StoreCredentialResponse(
+                handle=datamodel.CredentialHandle(driver="test", handle="opaque")
+            )
+
+    server = grpc.aio.server()
+    driver_grpc.add_CredentialDriverServicer_to_server(Driver(), server)
+    address = "127.0.0.1"
+    port = server.add_insecure_port(f"{address}:0")
+    await server.start()
+    try:
+        async with grpc.aio.insecure_channel(f"{address}:{port}") as channel:
+            stub = driver_grpc.CredentialDriverStub(channel)
+            result = await stub.StoreCredential(
+                credentials.StoreCredentialRequest(value="registration-ticket"),
+                timeout=5,
+            )
+            assert isinstance(result.handle, datamodel.CredentialHandle)
+            assert result.handle.handle == "opaque"
+    finally:
+        await server.stop(0)

python/generate_python_proto_test.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise generation in a clean tree so stale bindings cannot hide failures."""
+
+import ast
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize("no_pyi", [False, True])
+def test_generate_package_relative_bindings(tmp_path: Path, no_pyi: bool) -> None:
+    root = Path(__file__).resolve().parents[1]
+    shutil.copytree(root / "proto", tmp_path / "proto")
+    output = tmp_path / "python/openshell/_proto"
+    output.mkdir(parents=True)
+    command = [
+        sys.executable,
+        str(root / "tasks/scripts/generate_python_proto.py"),
+        *(["--no-pyi"] if no_pyi else []),
+    ]
+    subprocess.run(command, cwd=tmp_path, check=True, capture_output=True, text=True)
+
+    stems = {
+        "credential_driver",
+        "datamodel",
+        "extension",
+        "openshell",
+        "options",
+        "sandbox",
+    }
+    suffixes = ["_pb2.py", "_pb2_grpc.py"]
+    if not no_pyi:
+        suffixes.append("_pb2.pyi")
+    assert {path.name for path in output.iterdir()} == {
+        f"{stem}{suffix}" for stem in stems for suffix in suffixes
+    }
+    for path in output.iterdir():
+        for node in ast.walk(ast.parse(path.read_text())):
+            if isinstance(node, ast.Import):
+                assert not any(alias.name.endswith("_pb2") for alias in node.names), (
+                    f"Absolute generated import in {path.name}"
+                )
+
+    # Skipping type generation must also leave existing stubs untouched.
+    sentinel = output / "credential_driver_pb2.pyi"
+    sentinel.write_text("# existing type stub\n")
+    subprocess.run(
+        command if no_pyi else [*command, "--no-pyi"],
+        cwd=tmp_path,
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    assert sentinel.read_text() == "# existing type stub\n"

python/wheel_verifier_test.py

@@ -32,7 +32,14 @@ def _wheel_files() -> set[str]:
         "openshell/py.typed",
         "openshell/_proto/__init__.py",
     }
-    for stem in ("datamodel", "openshell", "options", "sandbox"):
+    for stem in (
+        "credential_driver",
+        "datamodel",
+        "extension",
+        "openshell",
+        "options",
+        "sandbox",
+    ):
         files.add(f"openshell/_proto/{stem}_pb2.py")
         files.add(f"openshell/_proto/{stem}_pb2.pyi")
         files.add(f"openshell/_proto/{stem}_pb2_grpc.py")
@@ -79,11 +86,15 @@ def test_requires_exactly_one_wheel_in_directory(tmp_path: Path) -> None:
         verifier.verify_wheel(tmp_path)
 
 
-def test_rejects_missing_generated_proto(tmp_path: Path) -> None:
-    files = _wheel_files() - {"openshell/_proto/sandbox_pb2.py"}
+@pytest.mark.parametrize("stem", ["sandbox", "credential_driver", "extension"])
+@pytest.mark.parametrize("suffix", ["_pb2.py", "_pb2.pyi", "_pb2_grpc.py"])
+def test_rejects_missing_generated_proto(
+    tmp_path: Path, stem: str, suffix: str
+) -> None:
+    files = _wheel_files() - {f"openshell/_proto/{stem}{suffix}"}
     wheel = _write_wheel(tmp_path, files=files)
 
-    with pytest.raises(ValueError, match=r"sandbox_pb2\.py"):
+    with pytest.raises(ValueError, match="wheel is missing required files"):
         verifier.verify_wheel(wheel)
 
 

tasks/scripts/generate_python_proto.py

@@ -14,63 +14,10 @@
     "proto/datamodel.proto",
     "proto/options.proto",
     "proto/sandbox.proto",
+    "proto/extension.proto",
+    "proto/credential_driver.proto",
 ]
 
-LINE_REWRITES = {
-    "python/openshell/_proto/openshell_pb2_grpc.py": [
-        (
-            r"^import openshell_pb2 as openshell__pb2$",
-            "from . import openshell_pb2 as openshell__pb2",
-        ),
-        (
-            r"^import sandbox_pb2 as sandbox__pb2$",
-            "from . import sandbox_pb2 as sandbox__pb2",
-        ),
-    ],
-    "python/openshell/_proto/openshell_pb2.py": [
-        (
-            r"^import datamodel_pb2 as datamodel__pb2$",
-            "from . import datamodel_pb2 as datamodel__pb2",
-        ),
-        (
-            r"^import options_pb2 as options__pb2$",
-            "from . import options_pb2 as options__pb2",
-        ),
-        (
-            r"^import sandbox_pb2 as sandbox__pb2$",
-            "from . import sandbox_pb2 as sandbox__pb2",
-        ),
-    ],
-    "python/openshell/_proto/datamodel_pb2.py": [
-        (
-            r"^import options_pb2 as options__pb2$",
-            "from . import options_pb2 as options__pb2",
-        ),
-        (
-            r"^import sandbox_pb2 as sandbox__pb2$",
-            "from . import sandbox_pb2 as sandbox__pb2",
-        ),
-    ],
-    "python/openshell/_proto/datamodel_pb2_grpc.py": [
-        (
-            r"^import datamodel_pb2 as datamodel__pb2$",
-            "from . import datamodel_pb2 as datamodel__pb2",
-        ),
-    ],
-    "python/openshell/_proto/sandbox_pb2.py": [
-        (
-            r"^import datamodel_pb2 as datamodel__pb2$",
-            "from . import datamodel_pb2 as datamodel__pb2",
-        ),
-    ],
-    "python/openshell/_proto/sandbox_pb2_grpc.py": [
-        (
-            r"^import sandbox_pb2 as sandbox__pb2$",
-            "from . import sandbox_pb2 as sandbox__pb2",
-        ),
-    ],
-}
-
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
@@ -83,6 +30,7 @@ def main() -> None:
         ),
     )
     args = parser.parse_args()
+    output_dir = Path("python/openshell/_proto")
 
     subprocess.run(
         [
@@ -98,13 +46,19 @@ def main() -> None:
         check=True,
     )
 
-    for path, rules in LINE_REWRITES.items():
-        file_path = Path(path)
-        text = file_path.read_text()
-        text = text.replace("from . from . import", "from . import")
-        for pattern, replacement in rules:
-            text = re.sub(pattern, replacement, text, flags=re.MULTILINE)
-        file_path.write_text(text)
+    # Rewrite only imports of our generated modules, leaving google.protobuf
+    # imports intact. The same rule covers runtime modules and type stubs.
+    modules = "|".join(re.escape(f"{Path(path).stem}_pb2") for path in PROTO_FILES)
+    pattern = rf"^import ({modules})( as \w+)?$"
+    suffixes = ["_pb2.py", "_pb2_grpc.py"]
+    if not args.no_pyi:
+        suffixes.append("_pb2.pyi")
+    for proto in PROTO_FILES:
+        for suffix in suffixes:
+            file_path = output_dir / f"{Path(proto).stem}{suffix}"
+            text = file_path.read_text()
+            text = re.sub(pattern, r"from . import \1\2", text, flags=re.MULTILINE)
+            file_path.write_text(text)
 
 
 if __name__ == "__main__":

tasks/scripts/verify-python-wheel.py

@@ -13,7 +13,14 @@
 from pathlib import Path
 from zipfile import BadZipFile, ZipFile
 
-PROTO_STEMS = ("datamodel", "openshell", "options", "sandbox")
+PROTO_STEMS = (
+    "credential_driver",
+    "datamodel",
+    "extension",
+    "openshell",
+    "options",
+    "sandbox",
+)
 BYTECODE_SUFFIXES = (".pyc", ".pyo")
 NATIVE_SUFFIXES = (".dll", ".dylib", ".exe", ".pyd", ".so")