feat(api)!: unify sandbox references

Signed-off-by: Drew Newberry <anewberry@nvidia.com>

python/openshell/sandbox.py

@@ -47,6 +47,18 @@ def _all_workspaces_scope() -> datamodel_pb2.WorkspaceSelector:
     return datamodel_pb2.WorkspaceSelector(all_workspaces=datamodel_pb2.AllWorkspaces())
 
 
+def _sandbox_reference_by_name(
+    workspace: str, name: str
+) -> datamodel_pb2.SandboxReference:
+    return datamodel_pb2.SandboxReference(
+        name=name, workspace_scope=_workspace_scope(workspace)
+    )
+
+
+def _sandbox_reference_by_id(sandbox_id: str) -> datamodel_pb2.SandboxReference:
+    return datamodel_pb2.SandboxReference(id=sandbox_id)
+
+
 class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails):
     pass
 
@@ -806,7 +818,7 @@ class SandboxClient:
     def get(self, sandbox_name: str, *, workspace: str) -> SandboxRef:
         response = self._stub.GetSandbox(
             openshell_pb2.GetSandboxRequest(
-                name=sandbox_name, workspace_scope=_workspace_scope(workspace)
+                sandbox_ref=_sandbox_reference_by_name(workspace, sandbox_name)
             ),
             timeout=self._timeout,
         )
@@ -885,7 +897,7 @@ class SandboxClient:
     def delete(self, sandbox_name: str, *, workspace: str) -> bool:
         response = self._stub.DeleteSandbox(
             openshell_pb2.DeleteSandboxRequest(
-                name=sandbox_name, workspace_scope=_workspace_scope(workspace)
+                sandbox_ref=_sandbox_reference_by_name(workspace, sandbox_name)
             ),
             timeout=self._timeout,
         )
@@ -894,7 +906,7 @@ class SandboxClient:
     def stop(self, sandbox_name: str, *, workspace: str) -> SandboxRef:
         response = self._stub.StopSandbox(
             openshell_pb2.StopSandboxRequest(
-                name=sandbox_name, workspace_scope=_workspace_scope(workspace)
+                sandbox_ref=_sandbox_reference_by_name(workspace, sandbox_name)
             ),
             timeout=self._timeout,
         )
@@ -903,7 +915,7 @@ class SandboxClient:
     def start(self, sandbox_name: str, *, workspace: str) -> SandboxRef:
         response = self._stub.StartSandbox(
             openshell_pb2.StartSandboxRequest(
-                name=sandbox_name, workspace_scope=_workspace_scope(workspace)
+                sandbox_ref=_sandbox_reference_by_name(workspace, sandbox_name)
             ),
             timeout=self._timeout,
         )
@@ -996,7 +1008,7 @@ class SandboxClient:
             raise SandboxError("command must not be empty")
 
         request = openshell_pb2.ExecSandboxRequest(
-            sandbox_id=sandbox_id,
+            sandbox_ref=_sandbox_reference_by_id(sandbox_id),
             command=list(command),
             workdir=workdir or "",
             environment=dict(env or {}),

python/openshell/sandbox_test.py

@@ -43,7 +43,12 @@ from openshell.sandbox import (
 
 
 def _request_workspace(request: Any) -> str | None:
-    scope = request.workspace_scope
+    sandbox_ref = getattr(request, "sandbox_ref", None)
+    scope = (
+        sandbox_ref.workspace_scope
+        if sandbox_ref is not None and sandbox_ref.HasField("workspace_scope")
+        else request.workspace_scope
+    )
     if scope.WhichOneof("selection") == "workspace":
         return cast("str", scope.workspace)
     return None
@@ -53,6 +58,12 @@ def _request_selects_all_workspaces(request: Any) -> bool:
     return request.workspace_scope.WhichOneof("selection") == "all_workspaces"
 
 
+def _request_sandbox_name(request: Any) -> str:
+    sandbox_ref = request.sandbox_ref
+    assert sandbox_ref.WhichOneof("identifier") == "name"
+    return cast("str", sandbox_ref.name)
+
+
 def _client_credentials_fixture() -> dict[str, Any]:
     return json.loads(
         (
@@ -1993,7 +2004,7 @@ class _FakeSandboxStub:
         return SimpleNamespace(
             sandbox=_make_sandbox_proto(
                 "sandbox-1",
-                request.name,
+                _request_sandbox_name(request),
                 workspace=_request_workspace(request) or "default",
             )
         )
@@ -2017,7 +2028,7 @@ class _FakeSandboxStub:
         return SimpleNamespace(
             sandbox=_make_sandbox_proto(
                 "sandbox-1",
-                request.name,
+                _request_sandbox_name(request),
                 phase=openshell_pb2.SANDBOX_PHASE_STOPPED,
                 workspace=_request_workspace(request) or "default",
             )
@@ -2033,7 +2044,7 @@ class _FakeSandboxStub:
         return SimpleNamespace(
             sandbox=_make_sandbox_proto(
                 "sandbox-1",
-                request.name,
+                _request_sandbox_name(request),
                 phase=openshell_pb2.SANDBOX_PHASE_STARTING,
                 workspace=_request_workspace(request) or "default",
             )
@@ -2429,13 +2440,13 @@ def test_stop_and_start_forward_workspace_and_return_phase() -> None:
 
     stopped = client.stop("job-1", workspace="team-a")
     assert stub.stop_request is not None
-    assert stub.stop_request.name == "job-1"
+    assert _request_sandbox_name(stub.stop_request) == "job-1"
     assert _request_workspace(stub.stop_request) == "team-a"
     assert stopped.phase == openshell_pb2.SANDBOX_PHASE_STOPPED
 
     starting = client.start("job-1", workspace="team-a")
     assert stub.start_request is not None
-    assert stub.start_request.name == "job-1"
+    assert _request_sandbox_name(stub.start_request) == "job-1"
     assert _request_workspace(stub.start_request) == "team-a"
     assert starting.phase == openshell_pb2.SANDBOX_PHASE_STARTING
 
@@ -2460,7 +2471,7 @@ def test_wait_ready_handles_terminal_main_process_results(
             return SimpleNamespace(
                 sandbox=_make_sandbox_proto(
                     "sandbox-1",
-                    request.name,
+                    _request_sandbox_name(request),
                     phase=phase,
                     workspace=_request_workspace(request) or "default",
                 )