Model resources as pydantic models

jumpstarter/drivers/base.py

@@ -22,6 +22,7 @@
     MARKER_STREAMCALL,
     MARKER_STREAMING_DRIVERCALL,
 )
+from jumpstarter.drivers.resources import ClientStreamResource, PresignedRequestResource, Resource
 from jumpstarter.drivers.streams import DriverStreamRequest, ResourceStreamRequest, StreamRequest
 from jumpstarter.v1 import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
@@ -136,8 +137,13 @@ def items(self, parent=None):
 
         return [(self.uuid, parent.uuid if parent else None, self)]
 
-    def resource(self, uuid: str):
-        return self.resources[UUID(uuid)]
+    def resource(self, handle: str):
+        handle = Resource.validate_python(handle)
+        match handle:
+            case ClientStreamResource(uuid=uuid):
+                return self.resources[uuid]
+            case PresignedRequestResource():
+                pass
 
     async def __lookup_drivercall(self, name, context, marker):
         """Lookup drivercall by method name

jumpstarter/drivers/core.py

@@ -4,6 +4,7 @@
 
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from uuid import UUID
 
 from anyio import create_task_group, sleep_forever
 from anyio.streams.stapled import StapledObjectStream
@@ -17,6 +18,7 @@
     create_memory_stream,
     forward_client_stream,
 )
+from jumpstarter.drivers.resources import ClientStreamResource
 from jumpstarter.drivers.streams import DriverStreamRequest, ResourceStreamRequest
 from jumpstarter.v1 import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
@@ -108,7 +110,7 @@ async def resource_async(
                 combined,
                 {"request": ResourceStreamRequest(uuid=self.uuid).model_dump_json()}.items(),
             ):
-                yield (await rx.receive()).decode()
+                yield ClientStreamResource(uuid=UUID((await rx.receive()).decode())).model_dump(mode="json")
 
     @asynccontextmanager
     async def file_async(

jumpstarter/drivers/resources.py

@@ -0,0 +1,24 @@
+from typing import Annotated, Literal, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, TypeAdapter
+
+
+class ClientStreamResource(BaseModel):
+    kind: Literal["client_stream"] = "client_stream"
+    uuid: UUID
+
+
+class PresignedRequestResource(BaseModel):
+    kind: Literal["presigned_request"] = "presigned_request"
+    headers: dict[str, str]
+    url: str
+    method: str
+
+
+Resource = TypeAdapter(
+    Annotated[
+        Union[ClientStreamResource, PresignedRequestResource],
+        Field(discriminator="kind"),
+    ]
+)