feat(sdk): add typed argument contracts

hud/environment/__init__.py

@@ -22,6 +22,7 @@
 from hud.capabilities import Capability
 from hud.utils.modules import iter_modules
 
+from .arguments import DataFileArg, DataFileRef, DataFilesArg, GradingArg, PromptArg
 from .egress import Peer
 from .env import Answer, Environment
 from .workspace import DEFAULT_SYSTEM_MOUNTS, Mount, MountKind, Workspace
@@ -89,10 +90,15 @@ def load_environment(
     "DEFAULT_SYSTEM_MOUNTS",
     "Answer",
     "Capability",
+    "DataFileArg",
+    "DataFileRef",
+    "DataFilesArg",
     "Environment",
+    "GradingArg",
     "Mount",
     "MountKind",
     "Peer",
+    "PromptArg",
     "Workspace",
     "load_environment",
 ]

hud/environment/arguments.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from typing import Annotated, TypeAlias, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+T = TypeVar("T")
+
+
+class DataFileRef(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    file_id: str = Field(description="HUD data-file id")
+    path: str | None = Field(default=None, description="Environment-owned destination path")
+
+
+PromptArg: TypeAlias = Annotated[
+    str,
+    Field(json_schema_extra={"x-hud-hint": "prompt"}),
+]
+
+DataFileArg: TypeAlias = Annotated[
+    T,
+    Field(json_schema_extra={"x-hud-hint": "data-file"}),
+]
+
+DataFilesArg: TypeAlias = Annotated[
+    list[T],
+    Field(json_schema_extra={"x-hud-hint": "data-files"}),
+]
+
+GradingArg: TypeAlias = Annotated[
+    list[T],
+    Field(json_schema_extra={"x-hud-hint": "grading"}),
+]
+
+
+__all__ = ["DataFileArg", "DataFileRef", "DataFilesArg", "GradingArg", "PromptArg"]

hud/environment/server.py

@@ -72,23 +72,24 @@ def _jsonable(value: Any) -> Any:
 
 
 def _coerce_args(sig: inspect.Signature, args: dict[str, Any]) -> dict[str, Any]:
-    """Coerce string wire args into the task fn's annotated param types.
-
-    JSON-RPC sends args as JSON scalars/strings; a param annotated with a richer
-    type (Pydantic model, list, etc.) is validated via a ``TypeAdapter``. Values
-    that already match (or fail to validate) are passed through unchanged.
-    """
     coerced: dict[str, Any] = {}
     for name, value in args.items():
         param = sig.parameters.get(name)
         annotation = param.annotation if param is not None else inspect.Parameter.empty
-        if annotation in (inspect.Parameter.empty, str, Any) or not isinstance(value, str):
+        if annotation in (inspect.Parameter.empty, Any):
             coerced[name] = value
             continue
+        adapter = TypeAdapter(annotation)
         try:
-            coerced[name] = TypeAdapter(annotation).validate_json(value)
+            coerced[name] = adapter.validate_python(value)
         except ValidationError:
-            coerced[name] = value
+            if not isinstance(value, str):
+                coerced[name] = value
+                continue
+            try:
+                coerced[name] = adapter.validate_json(value)
+            except ValidationError:
+                coerced[name] = value
     return coerced
 
 

hud/environment/tests/test_arguments.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+from hud.environment import (
+    DataFileArg,
+    DataFileRef,
+    DataFilesArg,
+    Environment,
+    GradingArg,
+    PromptArg,
+)
+from hud.eval import Run
+
+from .conftest import served
+
+
+class _Attachment(DataFileRef):
+    expand: bool = False
+
+
+class _Fixture(DataFileRef):
+    channel: Literal["mail", "calendar"]
+
+
+class _Criterion(BaseModel):
+    requirement: str
+    weight: float = 1.0
+    guidance: str | None = None
+
+
+def test_argument_types_publish_editor_hints_and_model_schemas() -> None:
+    env = Environment("typed-args")
+
+    @env.template()
+    async def task(
+        prompt: PromptArg,
+        attachment: DataFileArg[_Attachment],
+        fixtures: DataFilesArg[_Fixture],
+        criteria: GradingArg[_Criterion],
+    ):
+        yield prompt
+        yield 1.0
+
+    schema = task.manifest_entry()["args"]
+    properties = schema["properties"]
+
+    assert properties["prompt"]["x-hud-hint"] == "prompt"
+    assert properties["attachment"]["x-hud-hint"] == "data-file"
+    assert properties["fixtures"]["x-hud-hint"] == "data-files"
+    assert properties["criteria"]["x-hud-hint"] == "grading"
+    assert properties["fixtures"]["items"]["$ref"] == "#/$defs/_Fixture"
+    assert properties["criteria"]["items"]["$ref"] == "#/$defs/_Criterion"
+    assert "channel" in schema["$defs"]["_Fixture"]["properties"]
+    assert "guidance" in schema["$defs"]["_Criterion"]["properties"]
+
+
+async def test_typed_argument_values_reach_task_as_models() -> None:
+    env = Environment("typed-args")
+
+    @env.template()
+    async def task(
+        attachments: DataFilesArg[_Attachment],
+        criteria: GradingArg[_Criterion],
+    ):
+        assert isinstance(attachments[0], _Attachment)
+        assert isinstance(criteria[0], _Criterion)
+        yield f"{attachments[0].path}:{criteria[0].requirement}"
+        yield 1.0
+
+    async with (
+        served(env) as client,
+        Run(
+            client,
+            "task",
+            {
+                "attachments": [{"file_id": "file-1", "path": "brief.pdf"}],
+                "criteria": [{"requirement": "Answer the question", "weight": 2.0}],
+            },
+        ) as run,
+    ):
+        assert run.prompt_text == "brief.pdf:Answer the question"
+        run.trace.content = "done"