fix(agents): allow create_agent + HostedRuntime for training configs (HUD-2100)

docs/v6/guides/training-agents.mdx

@@ -94,6 +94,11 @@ One job spans the session; each step appends a batch and trains on it:
 - **Open the job** with `group=8` - 8 rollouts per task, so the rewards are comparable (next).
 - **Roll out** the batch, the same eval as the [previous guide](/v6/guides/running-an-eval). The
   [runtime](/v6/reference/runtime) sets where it runs; swap `LocalRuntime` for `HUDRuntime()` unchanged.
+  Prefer **`HUDRuntime`** (or `LocalRuntime`) for TrainingClient loops: the agent stays local so
+  `create_agent(...)` gateway clients, `return_token_ids`, system prompts, and Tinker/Qwen
+  `completion_kwargs` all apply in-process. **`HostedRuntime`** also works with `create_agent`
+  (the gateway client is rebuilt remotely and those kwargs still travel), but a true BYOK /
+  custom `model_client` cannot be serialized — use `HUDRuntime` for that.
 - **Nudge** with `trainer.step` - the one line that learns. It scores each rollout against its group,
   shifts the weights, then **promotes** them so the gateway serves the new ones at once.
 

docs/v6/reference/agents.mdx

@@ -35,7 +35,14 @@ The HUD gateway is an OpenAI-compatible endpoint (`settings.hud_gateway_url`, de
 `https://inference.beta.hud.ai`) that fronts every provider behind your single `HUD_API_KEY`, so you
 switch between Claude, GPT, and Gemini by name alone, with unified tracing. `create_agent` accepts any id
 the gateway knows (`claude-...`, `gpt-...`, `gemini-...`); extra kwargs pass through to the agent's
-config.
+config — including `system_prompt` and `completion_kwargs` (for example
+`extra_body={"return_token_ids": True}` for training).
+
+For local loops (`HUDRuntime` / `LocalRuntime`) the agent keeps a gateway `model_client`. For
+[`HostedRuntime`](/v6/reference/runtime#hostedruntime) that client is stripped and rebuilt on the
+platform from the model name, so the same `create_agent(...)` call works for hosted evals and
+training. A true BYOK / custom `model_client` cannot travel with hosted submission — keep the agent
+loop local with `HUDRuntime` instead.
 
 Built-in agents are **catalog-driven**: each run they read the environment's manifest, open the
 capabilities they support, build the matching provider tools, and loop against `run.prompt_messages`.

docs/v6/reference/runtime.mdx

@@ -173,6 +173,14 @@ Where `HUDRuntime` runs the agent loop locally against a tunneled env, `HostedRu
 runs the agent right next to it. This process only submits the rollout and polls its trace to
 completion. It requires a gateway agent that can serialize its identity (Claude/OpenAI/Gemini).
 
+Agents from [`create_agent`](/v6/reference/agents#create-agent) are supported: the local gateway
+`model_client` is stripped and rebuilt on the platform, while `system_prompt`,
+`completion_kwargs` (including `return_token_ids` and Tinker/Qwen request kwargs), and other
+config travel intact. A true BYOK / custom `model_client` cannot be serialized — use
+`HUDRuntime` (or `LocalRuntime`) so the agent loop stays on your machine. For
+[`TrainingClient`](/v6/reference/training#trainingclient) workflows, prefer `HUDRuntime` unless
+you specifically need co-located hosted execution.
+
 ### `Runtime`
 
 ```python

docs/v6/reference/training.mdx

@@ -38,6 +38,18 @@ Gradients accumulate across `forward_backward` / `forward_backward_custom` / `ba
 | `TrainingClient` | Drives managed training for one model slug. |
 | `CheckpointResponse` | One node in the model's checkpoint tree. |
 
+## Runtime placement
+
+Training reuses the same [runtime](/v6/reference/runtime) placement as eval. Which one you pick
+matters for how the agent is constructed:
+
+| Placement | Agent loop | Use for training when |
+| --- | --- | --- |
+| `LocalRuntime` / `HUDRuntime` | Local (this process) | Default. `create_agent` gateway clients, `system_prompt`, and `completion_kwargs` (`return_token_ids`, Tinker/Qwen kwargs) apply in-process. Required for a true BYOK / custom `model_client`. |
+| `HostedRuntime` | Remote (serialized agent) | Co-located env+agent on a leased box. `create_agent(...)` is supported: the gateway client is rebuilt remotely and config kwargs still travel. Custom / BYOK clients fail early — switch to `HUDRuntime`. |
+
+Cookbooks that load a platform taskset (`Taskset.from_api`) use `HUDRuntime()` for this reason.
+
 ## TrainingClient
 
 A `TrainingClient` drives managed training for one model: it accumulates gradients from rewarded

hud/agents/__init__.py

@@ -30,7 +30,16 @@
 def create_agent(model: str, **kwargs: Any) -> GatewayAgent:
     """Create an agent routed through the HUD gateway.
 
-    For direct API access with provider API keys, instantiate the agent classes directly.
+    Attaches a gateway ``model_client`` for local loops (``HUDRuntime`` /
+    ``LocalRuntime``). For ``HostedRuntime``, that client is stripped by
+    :meth:`~hud.agents.tool_agent.ToolAgent.hosted_spec` and rebuilt on the
+    platform from the model name — ``system_prompt``, ``completion_kwargs``
+    (including ``return_token_ids`` / Tinker kwargs), and other config still
+    travel. True BYOK / custom clients are not serializable; use
+    ``HUDRuntime`` for those TrainingClient workflows.
+
+    For direct API access with provider API keys, instantiate the agent classes
+    directly.
     """
     requested_model = model
     model = normalize_gateway_model_id(model)

hud/agents/tool_agent.py

@@ -90,12 +90,17 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
     def hosted_spec(self) -> dict[str, Any]:
         """HUD-hosted execution runs the agent remotely, so it is
         reconstructed there from this identity (type, model, step budget, system
-        prompt) with the model resolved through the HUD gateway.
+        prompt, provider kwargs) with the model resolved through the HUD gateway.
         """
-        if self.config.model_client is not None:
+        from hud.utils.gateway import is_gateway_client
+
+        if not is_gateway_client(self.config.model_client):
             raise ValueError(
                 "hosted execution cannot serialize a custom model_client; "
-                "set the model by name and let the hosted runner build the gateway client"
+                "use create_agent(model, ...) so the hosted runner rebuilds the "
+                "gateway client, or run the agent loop locally with HUDRuntime() "
+                "/ LocalRuntime (recommended for TrainingClient workflows that "
+                "attach a BYOK client)"
             )
         agent_type = AgentType.of(self)
         if agent_type is None:

hud/eval/tests/test_hosted.py

@@ -92,13 +92,83 @@ def test_hosted_spec_serializes_full_config() -> None:
     assert "hosted_tools" not in config
 
 
+def test_hosted_spec_strips_create_agent_gateway_client(monkeypatch: pytest.MonkeyPatch) -> None:
+    """create_agent attaches a gateway client for local loops; HostedRuntime must
+    still serialize system_prompt + completion_kwargs (HUD-2100)."""
+    from hud.agents import create_agent
+    from hud.utils.gateway import GatewayModelInfo, GatewayProviderInfo, mark_gateway_client
+
+    class _GatewayStub:
+        """Attribute-capable stand-in for a provider SDK client."""
+
+    model = GatewayModelInfo(
+        id="arith-rl",
+        model_name="arith-rl",
+        sdk_agent_type="openai_compatible",
+        provider=GatewayProviderInfo(name="openai"),
+    )
+    monkeypatch.setattr("hud.agents.list_gateway_models", lambda: [model])
+    monkeypatch.setattr(
+        "hud.agents.build_gateway_client",
+        lambda _provider: mark_gateway_client(_GatewayStub()),
+    )
+
+    agent = create_agent(
+        "arith-rl",
+        system_prompt="/no_think",
+        completion_kwargs={
+            "extra_body": {
+                "return_token_ids": True,
+                "chat_template_kwargs": {"enable_thinking": False},
+            }
+        },
+    )
+    assert agent.config.model_client is not None
+
+    spec = agent.hosted_spec()
+    config = spec["config"]
+    assert spec["type"] == "openai_compatible"
+    assert config["model"] == "arith-rl"
+    assert config["system_prompt"] == "/no_think"
+    assert config["completion_kwargs"]["extra_body"]["return_token_ids"] is True
+    assert config["completion_kwargs"]["extra_body"]["chat_template_kwargs"] == {
+        "enable_thinking": False
+    }
+    assert "model_client" not in config
+
+
 def test_hosted_spec_rejects_custom_model_client() -> None:
     agent = _agent()
     agent.config = OpenAIChatConfig(model="m", model_client=object())
-    with pytest.raises(ValueError, match="model_client"):
+    with pytest.raises(ValueError, match="custom model_client"):
+        agent.hosted_spec()
+    with pytest.raises(ValueError, match="HUDRuntime"):
         agent.hosted_spec()
 
 
+def test_mark_gateway_client_round_trip() -> None:
+    from hud.utils.gateway import is_gateway_client, mark_gateway_client
+
+    class _Stub:
+        pass
+
+    tagged = mark_gateway_client(_Stub())
+    assert is_gateway_client(tagged)
+    assert not is_gateway_client(_Stub())
+    assert is_gateway_client(None)
+
+
+def test_mark_gateway_client_works_on_async_openai() -> None:
+    """Provider SDK clients must accept the attribute mark (no id-registry fallback)."""
+    from openai import AsyncOpenAI
+
+    from hud.utils.gateway import is_gateway_client, mark_gateway_client
+
+    client = AsyncOpenAI(api_key="test", base_url="http://localhost")
+    assert is_gateway_client(mark_gateway_client(client))
+    assert not is_gateway_client(AsyncOpenAI(api_key="test", base_url="http://localhost"))
+
+
 @pytest.mark.asyncio
 async def test_run_rejects_non_gateway_agent() -> None:
     """An agent that can't serialize its identity yields a failed Run, not a crash."""

hud/utils/gateway.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 from functools import lru_cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 from openai import AsyncOpenAI
 from pydantic import BaseModel, Field
@@ -25,6 +25,8 @@
 
     GatewayClient: TypeAlias = AsyncAnthropic | AsyncAnthropicBedrock | GenaiClient | AsyncOpenAI
 
+_T = TypeVar("_T")
+
 
 class GatewayProviderInfo(BaseModel):
     name: str | None = None
@@ -88,22 +90,52 @@ def build_gateway_client(provider: str) -> GatewayClient:
     if provider == "anthropic":
         from anthropic import AsyncAnthropic
 
-        return AsyncAnthropic(api_key=settings.api_key, base_url=settings.hud_gateway_url)
+        client: GatewayClient = AsyncAnthropic(
+            api_key=settings.api_key, base_url=settings.hud_gateway_url
+        )
+        return mark_gateway_client(client)
 
     if provider == "gemini":
         from google import genai
         from google.genai.types import HttpOptions
 
-        return genai.Client(
+        client = genai.Client(
             api_key=settings.api_key,
             http_options=HttpOptions(
                 api_version="v1beta",
                 base_url=settings.hud_gateway_url,
             ),
         )
+        return mark_gateway_client(client)
 
     # OpenAI-compatible (openai, azure, together, groq, fireworks, etc.)
-    return AsyncOpenAI(api_key=settings.api_key, base_url=settings.hud_gateway_url)
+    client = AsyncOpenAI(api_key=settings.api_key, base_url=settings.hud_gateway_url)
+    return mark_gateway_client(client)
+
+
+_HUD_GATEWAY_CLIENT_ATTR = "_hud_gateway_client"
+
+
+def mark_gateway_client(client: _T) -> _T:
+    """Tag a client built for the HUD gateway so hosted serialization can drop it.
+
+    The mark lives on the instance (``setattr``). Clients that reject attributes
+    are not supported — real provider SDK clients accept them.
+    """
+    setattr(client, _HUD_GATEWAY_CLIENT_ATTR, True)
+    return client
+
+
+def is_gateway_client(client: object | None) -> bool:
+    """True when ``client`` is missing or was built by :func:`build_gateway_client`.
+
+    Hosted reconstruction rebuilds a gateway client from the model name, so a
+    tagged gateway client is safe to strip from ``hosted_spec``. A true BYOK /
+    custom client is not.
+    """
+    if client is None:
+        return True
+    return bool(getattr(client, _HUD_GATEWAY_CLIENT_ATTR, False))
 
 
 @lru_cache(maxsize=1)