feat(python): ship typed sync and async gRPC stubs

docs/sdk/python.mdx

@@ -148,6 +148,39 @@ 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.
 
+## Call Generated RPCs
+
+The generated gRPC stubs include concrete request and response types for synchronous `grpc.Channel` and asynchronous `grpc.aio.Channel` connections. Install the dependency type stubs in your development environment for editor signatures and type checking:
+
+```shell
+uv add --dev types-grpcio 'types-protobuf>=6.32,<7'
+```
+
+Construct `OpenShellStub` with your channel. The channel determines whether calls return responses directly or awaitable gRPC calls:
+
+```python
+import grpc.aio
+
+from openshell._proto import datamodel_pb2 as dm
+from openshell._proto import openshell_pb2 as pb
+from openshell._proto.openshell_pb2_grpc import OpenShellStub
+
+
+async def get_provider(channel: grpc.aio.Channel, name: str) -> pb.ProviderResponse:
+    stub = OpenShellStub(channel)
+    return await stub.GetProvider(
+        pb.GetProviderRequest(
+            workspace_scope=dm.WorkspaceSelector(workspace="default"),
+            name=name,
+        ),
+        timeout=30,
+    )
+```
+
+Your application configures the channel's TLS and authentication, supplies RPC deadlines, and closes the channel. Create and use async channels on the same event loop. Cancellation and gRPC errors propagate directly. Generated stubs do not reuse `SandboxClient` authentication automatically; `SandboxClient` and `ProviderClient` remain synchronous.
+
+If you annotate a stored async stub explicitly, use a deferred annotation such as `"openshell_pb2_grpc.OpenShellAsyncStub"`. That type exists only in the `.pyi` file; construct `OpenShellStub` at runtime. The same convention applies to credential-driver stubs.
+
 ## Implement a Credential Driver
 
 The SDK includes generated message classes, type stubs, and gRPC bindings for

e2e/mcp-conformance/host-bridge.py

@@ -221,7 +221,7 @@ def do_POST(self) -> None:
                 timeout=TIMEOUT,
             )
             log(f"client exited {result.returncode} for {server_url}")
-            body = {
+            body: dict[str, object] = {
                 "exit_code": result.returncode,
                 "stdout": result.stdout,
                 "stderr": result.stderr,

pyproject.toml

@@ -49,11 +49,14 @@ dev = [
     "pytest-cov>=4.0",
     "pytest-xdist>=3.0",
     "ruff>=0.4",
-    "ty>=0.0.1a6",
+    "ty>=0.0.94",
     "grpcio-tools>=1.60",
     "pyelftools>=0.30",
     "pyyaml==6.0.2",
     "types-pyyaml>=6.0.12.20260518",
+    "mypy-protobuf>=5.1",
+    "types-grpcio>=1.0.0.20251009",
+    "types-protobuf>=6.32,<7",
 ]
 
 [tool.setuptools]

python/generate_python_proto_test.py

@@ -35,7 +35,7 @@ def test_generate_package_relative_bindings(tmp_path: Path, no_pyi: bool) -> Non
     }
     suffixes = ["_pb2.py", "_pb2_grpc.py"]
     if not no_pyi:
-        suffixes.append("_pb2.pyi")
+        suffixes.extend(["_pb2.pyi", "_pb2_grpc.pyi"])
     assert {path.name for path in output.iterdir()} == {
         f"{stem}{suffix}" for stem in stems for suffix in suffixes
     }

python/grpc_transport_test.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Generated async stubs retain native gRPC transport behavior."""
+
+import asyncio
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+import grpc
+import grpc.aio
+import pytest
+
+from openshell._proto import datamodel_pb2 as dm
+from openshell._proto import openshell_pb2 as pb
+from openshell._proto import openshell_pb2_grpc as rpc
+
+
+class Gateway(rpc.OpenShellServicer):
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+        self.cancelled = asyncio.Event()
+
+    async def GetProvider(
+        self, request: pb.GetProviderRequest, context: grpc.aio.ServicerContext
+    ) -> pb.ProviderResponse:
+        assert ("authorization", "Bearer test-token") in (
+            context.invocation_metadata() or ()
+        )
+        if request.name == "missing":
+            await context.abort(grpc.StatusCode.NOT_FOUND, "No provider")
+        if request.name == "blocked":
+            self.started.set()
+            try:
+                await asyncio.Future()
+            finally:
+                self.cancelled.set()
+        return pb.ProviderResponse(
+            provider=dm.Provider(metadata=dm.ObjectMeta(name=request.name))
+        )
+
+
+@asynccontextmanager
+async def gateway() -> AsyncGenerator[tuple[Gateway, grpc.aio.Channel]]:
+    service = Gateway()
+    server = grpc.aio.server()
+    rpc.add_OpenShellServicer_to_server(service, server)
+    port = server.add_insecure_port("127.0.0.1:0")
+    await server.start()
+    try:
+        async with grpc.aio.insecure_channel(f"127.0.0.1:{port}") as channel:
+            yield service, channel
+    finally:
+        await server.stop(0)
+
+
+_METADATA = (("authorization", "Bearer test-token"),)
+
+
+async def test_concurrent_calls_and_task_cancellation() -> None:
+    async with gateway() as (service, channel):
+        stub = rpc.OpenShellStub(channel)
+
+        async def blocked_call() -> pb.ProviderResponse:
+            return await stub.GetProvider(
+                pb.GetProviderRequest(name="blocked"), timeout=5, metadata=_METADATA
+            )
+
+        pending = asyncio.create_task(blocked_call())
+        try:
+            await asyncio.wait_for(service.started.wait(), timeout=5)
+            response = await stub.GetProvider(
+                pb.GetProviderRequest(name="ready"), timeout=5, metadata=_METADATA
+            )
+            assert response.provider.metadata.name == "ready"
+        finally:
+            pending.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await pending
+        await asyncio.wait_for(service.cancelled.wait(), timeout=5)
+        response = await stub.GetProvider(
+            pb.GetProviderRequest(name="still-ready"), timeout=5, metadata=_METADATA
+        )
+        assert response.provider.metadata.name == "still-ready"
+
+
+@pytest.mark.parametrize(
+    ("name", "expected"),
+    [
+        ("missing", grpc.StatusCode.NOT_FOUND),
+        ("blocked", grpc.StatusCode.DEADLINE_EXCEEDED),
+    ],
+)
+async def test_status_and_deadline(name: str, expected: grpc.StatusCode) -> None:
+    async with gateway() as (_, channel):
+        stub = rpc.OpenShellStub(channel)
+        with pytest.raises(grpc.aio.AioRpcError) as error:
+            await stub.GetProvider(
+                pb.GetProviderRequest(name=name), timeout=0.2, metadata=_METADATA
+            )
+        assert error.value.code() == expected

python/grpc_typing_test.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Check the consumer-facing types of generated sync and async RPCs."""
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+def test_generated_rpc_types(tmp_path: Path, asynchronous: bool) -> None:
+    channel = "grpc.aio.Channel" if asynchronous else "grpc.Channel"
+    function = "async def" if asynchronous else "def"
+    await_call = "await " if asynchronous else ""
+    iterate = "async for" if asynchronous else "for"
+    source = f"""
+from typing import assert_type
+import grpc
+import grpc.aio
+from openshell._proto import credential_driver_pb2 as credentials
+from openshell._proto import credential_driver_pb2_grpc as driver
+from openshell._proto import openshell_pb2 as pb
+from openshell._proto import openshell_pb2_grpc as rpc
+
+{function} check(channel: {channel}) -> None:
+    stub = rpc.OpenShellStub(channel)
+    assert_type({await_call}stub.GetProvider(pb.GetProviderRequest()), pb.ProviderResponse)
+    assert_type({await_call}stub.CreateProvider(pb.CreateProviderRequest()), pb.ProviderResponse)
+    assert_type({await_call}stub.UpdateProvider(pb.UpdateProviderRequest()), pb.ProviderResponse)
+    assert_type({await_call}stub.DeleteProvider(pb.DeleteProviderRequest()), pb.DeleteProviderResponse)
+    assert_type({await_call}stub.ListProviders(pb.ListProvidersRequest()), pb.ListProvidersResponse)
+    assert_type({await_call}stub.GetProviderProfile(pb.GetProviderProfileRequest()), pb.ProviderProfileResponse)
+    assert_type({await_call}stub.GetSandbox(pb.GetSandboxRequest()), pb.SandboxResponse)
+    assert_type({await_call}stub.DetachSandboxProvider(pb.DetachSandboxProviderRequest()), pb.DetachSandboxProviderResponse)
+    {iterate} event in stub.ExecSandbox(pb.ExecSandboxRequest()):
+        assert_type(event, pb.ExecSandboxEvent)
+    {iterate} frame in stub.ForwardTcp(iter([pb.TcpForwardFrame()])):
+        assert_type(frame, pb.TcpForwardFrame)
+    credential_stub = driver.CredentialDriverStub(channel)
+    assert_type({await_call}credential_stub.StoreCredential(credentials.StoreCredentialRequest()), credentials.StoreCredentialResponse)
+"""
+    path = tmp_path / "consumer.py"
+    path.write_text(source)
+    result = _check_types(path)
+    assert result.returncode == 0, result.stdout + result.stderr
+
+    # A missing import or Any-typed callable must not make this test pass.
+    path.write_text(
+        source.replace(
+            "stub.GetProvider(pb.GetProviderRequest())",
+            "stub.GetProvider(pb.DeleteProviderRequest())",
+        )
+    )
+    result = _check_types(path)
+    assert result.returncode != 0
+    assert "invalid-argument-type" in result.stdout
+    assert "GetProviderRequest" in result.stdout
+    assert "DeleteProviderRequest" in result.stdout
+
+
+def _check_types(path: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "ty",
+            "check",
+            "--python",
+            sys.executable,
+            "--output-format",
+            "concise",
+            str(path),
+        ],
+        cwd=Path(__file__).resolve().parents[1],
+        capture_output=True,
+        text=True,
+        check=False,
+    )

python/openshell/sandbox_test.py

@@ -15,6 +15,7 @@
 from typing import Any, cast
 
 import pytest
+from google.protobuf.struct_pb2 import Struct
 
 import openshell.sandbox as sandbox_module
 from openshell._proto import openshell_pb2
@@ -427,14 +428,14 @@ def ExecSandbox(
 
 
 def _client_with_fake_stub(stub: object) -> SandboxClient:
-    client = cast("SandboxClient", object.__new__(SandboxClient))
+    client = object.__new__(SandboxClient)
     client._timeout = 30.0
     client._stub = cast("Any", stub)
     return client
 
 
 def _template_client_with_fake_stub(stub: object) -> SandboxTemplateClient:
-    client = cast("SandboxTemplateClient", object.__new__(SandboxTemplateClient))
+    client = object.__new__(SandboxTemplateClient)
     client._timeout = 30.0
     client._stub = cast("Any", stub)
     return client
@@ -2347,7 +2348,9 @@ def test_sandbox_template_create_builds_template_from_public_fields() -> None:
     assert template.spec.workload.resources.cpu == "1"
     assert template.spec.workload.resources.memory == "512Mi"
     assert template.spec.workload.resources.gpu.count == 2
-    assert template.spec.driver_config["kubernetes"]["runtime_class_name"] == "kata"
+    config = template.spec.driver_config["kubernetes"]
+    assert isinstance(config, Struct)
+    assert config["runtime_class_name"] == "kata"
 
 
 def test_sandbox_template_create_materializes_default_workload() -> None:
@@ -2380,7 +2383,9 @@ def test_sandbox_template_create_materializes_workload_with_driver_config_only()
     assert template.HasField("spec")
     assert template.spec.HasField("workload")
     assert template.spec.workload.image == ""
-    assert template.spec.driver_config["kubernetes"]["runtime_class_name"] == "kata"
+    config = template.spec.driver_config["kubernetes"]
+    assert isinstance(config, Struct)
+    assert config["runtime_class_name"] == "kata"
 
 
 def test_sandbox_template_create_rejects_missing_public_name() -> None:
@@ -2486,12 +2491,9 @@ def test_sandbox_template_client_crud_forwards_requests() -> None:
         stub.create_template_request.template.spec.workload.image
         == "ghcr.io/test/gpu-kata:latest"
     )
-    assert (
-        stub.create_template_request.template.spec.driver_config["kubernetes"][
-            "runtime_class_name"
-        ]
-        == "kata"
-    )
+    config = stub.create_template_request.template.spec.driver_config["kubernetes"]
+    assert isinstance(config, Struct)
+    assert config["runtime_class_name"] == "kata"
 
     got = client.get("gpu-kata", workspace="default")
     assert got.metadata.name == "gpu-kata"
@@ -2892,7 +2894,7 @@ def test_returned_labels_are_immutable() -> None:
     ref = _sandbox_ref(proto)
 
     with pytest.raises(TypeError):
-        ref.labels["mutated"] = "nope"  # type: ignore[index]
+        ref.labels["mutated"] = "nope"  # ty: ignore[invalid-assignment]
 
 
 def test_direct_sandbox_ref_construction_defaults_labels() -> None:
@@ -2953,7 +2955,7 @@ def test_direct_sandbox_ref_copies_and_freezes_labels() -> None:
 
     assert dict(ref.labels) == {"aiq": "deep-research"}
     with pytest.raises(TypeError):
-        ref.labels["mutated"] = "nope"  # type: ignore[index]
+        ref.labels["mutated"] = "nope"  # ty: ignore[invalid-assignment]
 
 
 def test_high_level_creation_forwards_name_and_labels(

python/wheel_verifier_test.py

@@ -43,6 +43,7 @@ def _wheel_files() -> set[str]:
         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")
+        files.add(f"openshell/_proto/{stem}_pb2_grpc.pyi")
     return files
 
 
@@ -86,8 +87,10 @@ def test_requires_exactly_one_wheel_in_directory(tmp_path: Path) -> None:
         verifier.verify_wheel(tmp_path)
 
 
-@pytest.mark.parametrize("stem", ["sandbox", "credential_driver", "extension"])
-@pytest.mark.parametrize("suffix", ["_pb2.py", "_pb2.pyi", "_pb2_grpc.py"])
+@pytest.mark.parametrize("stem", verifier.PROTO_STEMS)
+@pytest.mark.parametrize(
+    "suffix", ["_pb2.py", "_pb2.pyi", "_pb2_grpc.py", "_pb2_grpc.pyi"]
+)
 def test_rejects_missing_generated_proto(
     tmp_path: Path, stem: str, suffix: str
 ) -> None:

tasks/scripts/generate_python_proto.py

@@ -39,7 +39,14 @@ def main() -> None:
             "grpc_tools.protoc",
             "-Iproto",
             "--python_out=python/openshell/_proto",
-            *([] if args.no_pyi else ["--pyi_out=python/openshell/_proto"]),
+            *(
+                []
+                if args.no_pyi
+                else [
+                    "--pyi_out=python/openshell/_proto",
+                    "--mypy_grpc_out=generate_concrete_servicer_stubs:python/openshell/_proto",
+                ]
+            ),
             "--grpc_python_out=python/openshell/_proto",
             *PROTO_FILES,
         ],
@@ -52,7 +59,7 @@ def main() -> None:
     pattern = rf"^import ({modules})( as \w+)?$"
     suffixes = ["_pb2.py", "_pb2_grpc.py"]
     if not args.no_pyi:
-        suffixes.append("_pb2.pyi")
+        suffixes.extend(["_pb2.pyi", "_pb2_grpc.pyi"])
     for proto in PROTO_FILES:
         for suffix in suffixes:
             file_path = output_dir / f"{Path(proto).stem}{suffix}"

tasks/scripts/sync_docs_website.py

@@ -354,7 +354,7 @@ def ordered_entries(
 def render_versions(entries: list[VersionEntry]) -> list[YamlMapping]:
     rendered: list[YamlMapping] = []
     for entry in entries:
-        item = {
+        item: YamlMapping = {
             "display-name": entry.display_name,
             "path": entry.path,
             "slug": entry.slug,