From a6ac044c18849207f7e69f6dba0c2d5e021912d3 Mon Sep 17 00:00:00 2001
From: solvemproblr <gg.asadullo@gmail.com>
Date: Wed, 1 Jul 2026 06:17:11 +0500
Subject: [PATCH] fix: update signature handling in _TaskFactory to allow
 forward refs

---
 hud/environment/env.py                 |  2 +-
 hud/environment/tests/test_manifest.py | 26 ++++++++++++++++
 hud/environment/tests/test_server.py   | 41 ++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/hud/environment/env.py b/hud/environment/env.py
index 19f2d9d0..07e1afd8 100644
--- a/hud/environment/env.py
+++ b/hud/environment/env.py
@@ -98,7 +98,7 @@ def __init__(
         #: Type the agent must produce (``None`` = plain text). Drives answer
         #: deserialization into ``Answer[T]``.
         self.return_type = returns
-        self.sig = inspect.signature(func)
+        self.sig = inspect.signature(func, eval_str=True)
         functools.update_wrapper(self, func)
 
     def manifest_entry(self) -> dict[str, Any]:
diff --git a/hud/environment/tests/test_manifest.py b/hud/environment/tests/test_manifest.py
index 8b0f6f3c..8b7c9e2b 100644
--- a/hud/environment/tests/test_manifest.py
+++ b/hud/environment/tests/test_manifest.py
@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 from pydantic import BaseModel
 
 from hud.environment import Environment
@@ -17,6 +19,14 @@ class _Point(BaseModel):
     y: int
 
 
+class _Payload(BaseModel):
+    text: str
+
+
+_Mode = Literal["easy", "hard"]
+_Retries = int | None
+
+
 def test_args_schema_captures_params_defaults_and_required() -> None:
     env = Environment("manifests")
 
@@ -86,3 +96,19 @@ async def typed():
     assert entry["input"]["properties"]["x"]["type"] == "integer"
     assert entry["returns"]["properties"]["y"]["type"] == "integer"
     assert entry["args"]["properties"] == {}
+
+
+def test_args_schema_resolves_postponed_rich_annotations() -> None:
+    env = Environment("manifests")
+
+    @env.template()
+    async def rich(mode: _Mode, payload: _Payload, retries: _Retries = None):
+        yield "go"
+        yield 1.0
+
+    assert callable(rich)
+    schema = env.tasks["rich"].manifest_entry()["args"]
+    assert schema["properties"]["mode"]["enum"] == ["easy", "hard"]
+    assert schema["properties"]["retries"]["default"] is None
+    assert "$ref" in schema["properties"]["payload"]
+    assert set(schema["required"]) == {"mode", "payload"}
diff --git a/hud/environment/tests/test_server.py b/hud/environment/tests/test_server.py
index 3166de59..26efec18 100644
--- a/hud/environment/tests/test_server.py
+++ b/hud/environment/tests/test_server.py
@@ -7,7 +7,10 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 import pytest
+from pydantic import BaseModel
 
 from hud.clients import HudProtocolError
 from hud.environment import Answer, Environment
@@ -17,6 +20,13 @@
 from .conftest import served
 
 
+class _Payload(BaseModel):
+    text: str
+
+
+_Mode = Literal["upper", "lower"]
+
+
 async def test_dict_grade_without_numeric_score_errors_loudly() -> None:
     env = Environment("badgrade")
 
@@ -79,3 +89,34 @@ def test_answer_holds_parsed_content_and_raw_string() -> None:
     answer = Answer(content={"final": "42"}, raw='{"final": "42"}')
     assert answer.content == {"final": "42"}
     assert answer.raw == '{"final": "42"}'
+
+
+async def test_start_coerces_postponed_rich_annotations() -> None:
+    env = Environment("coerce")
+
+    @env.template()
+    async def typed(mode: _Mode, payload: _Payload, retries: int | None = None):
+        if mode == "upper":
+            prompt = payload.text.upper()
+        elif mode == "lower":
+            prompt = payload.text.lower()
+        else:
+            raise ValueError(f"unexpected mode: {mode!r}")
+        if retries is not None:
+            prompt += "!" * retries
+        yield prompt
+        yield 1.0
+
+    assert callable(typed)
+    async with served(env) as client:
+        async with Run(
+            client,
+            "typed",
+            {
+                "mode": '"upper"',
+                "payload": '{"text":"hello"}',
+                "retries": "3",
+            },
+        ) as run:
+            run.trace.content = "x"
+        assert run.prompt == "HELLO!!!"