diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index 64b1d9cb54..035e5e73e2 100644 --- a/docs/sdk/python.mdx +++ b/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. diff --git a/python/credential_driver_proto_test.py b/python/credential_driver_proto_test.py new file mode 100644 index 0000000000..f3bd96289d --- /dev/null +++ b/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) diff --git a/python/generate_python_proto_test.py b/python/generate_python_proto_test.py new file mode 100644 index 0000000000..ad57378ba7 --- /dev/null +++ b/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" diff --git a/python/wheel_verifier_test.py b/python/wheel_verifier_test.py index 873fb75fbc..edbafb5568 100644 --- a/python/wheel_verifier_test.py +++ b/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) diff --git a/tasks/scripts/generate_python_proto.py b/tasks/scripts/generate_python_proto.py index 4a58031ff8..b9c6cec01f 100644 --- a/tasks/scripts/generate_python_proto.py +++ b/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__": diff --git a/tasks/scripts/verify-python-wheel.py b/tasks/scripts/verify-python-wheel.py index 0539bd88a0..d694d410fd 100644 --- a/tasks/scripts/verify-python-wheel.py +++ b/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")