From f10a63c3acf871eaa070b85593c083c44c604064 Mon Sep 17 00:00:00 2001 From: James Graham Date: Thu, 2 Jul 2026 09:34:53 +0100 Subject: [PATCH] Pass agent inputs through a file rather than environment variables To pass inputs ot agents in production, upload a file to a GCS bucket, and pass a URL to that file into the agent environment. This enables having inputs which don't fit inside the 32kB limits of enironment variables for GCS. For local runs using docker, data is still pased via environment variables. --- .../autowebcompat_repro/__main__.py | 9 +-- .../hackbot_agents/bug_fix/__main__.py | 9 +-- .../hackbot_agents/build_repair/__main__.py | 9 +-- .../frontend_triage/__main__.py | 9 +-- .../test_plan_generator/__main__.py | 9 +-- .../hackbot_runtime/__init__.py | 3 +- .../hackbot_runtime/context.py | 38 +++++++++- .../hackbot_runtime/remote_config.py | 25 +++++++ libs/hackbot-runtime/tests/test_context.py | 55 ++++++++++++++- services/hackbot-api/app/agents.py | 28 -------- services/hackbot-api/app/gcs.py | 33 +++++++++ services/hackbot-api/app/routers/runs.py | 8 ++- services/hackbot-api/tests/test_agents.py | 69 ++----------------- 13 files changed, 175 insertions(+), 129 deletions(-) create mode 100644 libs/hackbot-runtime/hackbot_runtime/remote_config.py diff --git a/agents/autowebcompat-repro/hackbot_agents/autowebcompat_repro/__main__.py b/agents/autowebcompat-repro/hackbot_agents/autowebcompat_repro/__main__.py index bdce00e69c..f397ee7bb9 100644 --- a/agents/autowebcompat-repro/hackbot_agents/autowebcompat_repro/__main__.py +++ b/agents/autowebcompat-repro/hackbot_agents/autowebcompat_repro/__main__.py @@ -1,12 +1,11 @@ -from hackbot_runtime import HackbotContext, run_async -from pydantic_settings import BaseSettings, SettingsConfigDict +from hackbot_runtime import BaseAgentInputs, HackbotContext, run_async from .agent import AutowebcompatReproResult, run_autowebcompat_repro from .firefox_install import install_firefox_nightly from .setup_profile import setup_profile -class AgentInputs(BaseSettings): +class AgentInputs(BaseAgentInputs): bugzilla_mcp_url: str bug_data: str | None = None bug_id: int | None = None @@ -14,11 +13,9 @@ class AgentInputs(BaseSettings): max_turns: int | None = None effort: str | None = None - model_config = SettingsConfigDict(extra="ignore") - async def main(ctx: HackbotContext) -> AutowebcompatReproResult: - inputs = AgentInputs() + inputs = ctx.load_inputs(AgentInputs) # Provision a fresh Nightly at startup so each run reproduces against a # current build; drive the binary the install reports back. diff --git a/agents/bug-fix/hackbot_agents/bug_fix/__main__.py b/agents/bug-fix/hackbot_agents/bug_fix/__main__.py index c1423e3d67..dad341e1d9 100644 --- a/agents/bug-fix/hackbot_agents/bug_fix/__main__.py +++ b/agents/bug-fix/hackbot_agents/bug_fix/__main__.py @@ -1,21 +1,18 @@ -from hackbot_runtime import HackbotContext, run_async -from pydantic_settings import BaseSettings, SettingsConfigDict +from hackbot_runtime import BaseAgentInputs, HackbotContext, run_async from .agent import BugFixResult, run_bug_fix -class AgentInputs(BaseSettings): +class AgentInputs(BaseAgentInputs): bug_id: int bugzilla_mcp_url: str model: str | None = None max_turns: int | None = None effort: str | None = None - model_config = SettingsConfigDict(extra="ignore") - async def main(ctx: HackbotContext) -> BugFixResult: - inputs = AgentInputs() + inputs = ctx.load_inputs(AgentInputs) return await run_bug_fix( task="Triage and fix the bug, and verify the fix", diff --git a/agents/build-repair/hackbot_agents/build_repair/__main__.py b/agents/build-repair/hackbot_agents/build_repair/__main__.py index bef34de4ef..e64194ff51 100644 --- a/agents/build-repair/hackbot_agents/build_repair/__main__.py +++ b/agents/build-repair/hackbot_agents/build_repair/__main__.py @@ -1,12 +1,11 @@ import os -from hackbot_runtime import HackbotContext, run_async -from pydantic_settings import BaseSettings, SettingsConfigDict +from hackbot_runtime import BaseAgentInputs, HackbotContext, run_async from .agent import BuildRepairResult, run_build_repair -class AgentInputs(BaseSettings): +class AgentInputs(BaseAgentInputs): bug_id: int | None = None git_commit: str failure_tasks: dict[str, str] @@ -15,11 +14,9 @@ class AgentInputs(BaseSettings): model: str | None = None max_turns: int | None = None - model_config = SettingsConfigDict(extra="ignore") - async def main(ctx: HackbotContext) -> BuildRepairResult: - inputs = AgentInputs() + inputs = ctx.load_inputs(AgentInputs) # The build failure lives at this commit; pin the checkout there before the # runtime prepares the source tree (consumed in HackbotContext.source_repo). diff --git a/agents/frontend-triage/hackbot_agents/frontend_triage/__main__.py b/agents/frontend-triage/hackbot_agents/frontend_triage/__main__.py index aa24b74b05..8254e90149 100644 --- a/agents/frontend-triage/hackbot_agents/frontend_triage/__main__.py +++ b/agents/frontend-triage/hackbot_agents/frontend_triage/__main__.py @@ -1,5 +1,4 @@ -from hackbot_runtime import HackbotContext, run_async -from pydantic_settings import BaseSettings, SettingsConfigDict +from hackbot_runtime import BaseAgentInputs, HackbotContext, run_async from .agent import FrontendTriageResult, run_frontend_triage @@ -13,18 +12,16 @@ ) -class AgentInputs(BaseSettings): +class AgentInputs(BaseAgentInputs): bug_id: int bugzilla_mcp_url: str model: str | None = None max_turns: int | None = None effort: str | None = None - model_config = SettingsConfigDict(extra="ignore") - async def main(ctx: HackbotContext) -> FrontendTriageResult: - inputs = AgentInputs() + inputs = ctx.load_inputs(AgentInputs) return await run_frontend_triage( task=TRIAGE_TASK, diff --git a/agents/test-plan-generator/hackbot_agents/test_plan_generator/__main__.py b/agents/test-plan-generator/hackbot_agents/test_plan_generator/__main__.py index af67eac60b..348e3a76b0 100644 --- a/agents/test-plan-generator/hackbot_agents/test_plan_generator/__main__.py +++ b/agents/test-plan-generator/hackbot_agents/test_plan_generator/__main__.py @@ -1,11 +1,10 @@ -from hackbot_runtime import HackbotContext, run_async -from pydantic_settings import BaseSettings, SettingsConfigDict +from hackbot_runtime import BaseAgentInputs, HackbotContext, run_async from .agent import TestPlanGeneratorResult, run_test_plan_generator from .firefox_install import install_firefox_nightly -class AgentInputs(BaseSettings): +class AgentInputs(BaseAgentInputs): feature_name: str feature_description: str test_scope: str @@ -13,11 +12,9 @@ class AgentInputs(BaseSettings): max_turns: int | None = None effort: str | None = None - model_config = SettingsConfigDict(extra="ignore") - async def main(ctx: HackbotContext) -> TestPlanGeneratorResult: - inputs = AgentInputs() + inputs = ctx.load_inputs(AgentInputs) firefox_path = str(install_firefox_nightly()) diff --git a/libs/hackbot-runtime/hackbot_runtime/__init__.py b/libs/hackbot-runtime/hackbot_runtime/__init__.py index 277d2084f0..750a42671e 100644 --- a/libs/hackbot-runtime/hackbot_runtime/__init__.py +++ b/libs/hackbot-runtime/hackbot_runtime/__init__.py @@ -1,6 +1,6 @@ from hackbot_runtime.actions.recorder import ActionsRecorder from hackbot_runtime.config import HackbotConfig -from hackbot_runtime.context import HackbotContext +from hackbot_runtime.context import BaseAgentInputs, HackbotContext from hackbot_runtime.errors import AgentError from hackbot_runtime.results import HackbotAgentResult from hackbot_runtime.runtime import run, run_async @@ -13,6 +13,7 @@ "HackbotAgentResult", "HackbotConfig", "HackbotContext", + "BaseAgentInputs", "SignedPolicyUploader", "ensure_source_repo", "run", diff --git a/libs/hackbot-runtime/hackbot_runtime/context.py b/libs/hackbot-runtime/hackbot_runtime/context.py index 8269de8fd1..dc29c2202a 100644 --- a/libs/hackbot-runtime/hackbot_runtime/context.py +++ b/libs/hackbot-runtime/hackbot_runtime/context.py @@ -19,15 +19,20 @@ import uuid from functools import cached_property from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeVar from pydantic import Field, PrivateAttr -from pydantic_settings import BaseSettings, SettingsConfigDict +from pydantic_settings import ( + BaseSettings, + PydanticBaseSettingsSource, + SettingsConfigDict, +) from hackbot_runtime import artifacts, changes from hackbot_runtime.actions.recorder import ActionsRecorder from hackbot_runtime.config import HackbotConfig, load_config from hackbot_runtime.providers import AnthropicAuth +from hackbot_runtime.remote_config import load_remote_config from hackbot_runtime.source import ensure_source_repo from hackbot_runtime.uploader import SignedPolicyUploader @@ -37,6 +42,26 @@ log = logging.getLogger("hackbot_runtime.context") +class BaseAgentInputs(BaseSettings): + model_config = SettingsConfigDict(extra="ignore") + + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + # Environment variables override settings passed through + # as init parameters + return (env_settings, dotenv_settings, init_settings, file_secret_settings) + + +InputsType = TypeVar("InputsType", bound=BaseAgentInputs) + + def _default_run_id() -> str: """A unique, sortable id for runs that don't get one from the platform. @@ -62,6 +87,7 @@ class HackbotContext(BaseSettings): results_prefix: str = "" results_policy_url: str | None = None results_policy_fields: dict[str, str] = {} + run_inputs_url: str | None = None # Base for locally-persisted artifacts when no uploader is configured # (compose/direct runs). Each run is namespaced under it by run_id (see # `run_artifacts_dir`). Overridable via ARTIFACTS_DIR — compose points this @@ -177,6 +203,14 @@ def log_path(self) -> Path: def actions(self) -> ActionsRecorder: return ActionsRecorder(self.uploader, artifacts_dir=self.run_artifacts_dir) + def load_inputs(self, inputs_cls: type[InputsType]) -> InputsType: + remote_config = load_remote_config(self.run_inputs_url) + if remote_config: + kwargs = remote_config + else: + kwargs = {} + return inputs_cls(**kwargs) + def publish_file( self, key: str, path: Path, content_type: str | None = None ) -> str: diff --git a/libs/hackbot-runtime/hackbot_runtime/remote_config.py b/libs/hackbot-runtime/hackbot_runtime/remote_config.py new file mode 100644 index 0000000000..58d24063e2 --- /dev/null +++ b/libs/hackbot-runtime/hackbot_runtime/remote_config.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import logging +from typing import Mapping + +import requests +from pydantic.types import Json + +log = logging.getLogger("hackbot_runtime.remote_config") + + +def load_remote_config(config_url: str | None) -> Mapping[str, Json] | None: + if config_url is None: + return None + + response = requests.get(config_url, timeout=30) + response.raise_for_status() + config = response.json() + + if not isinstance(config, dict): + raise ValueError( + f"Config fetched from {config_url} was not a JSON object; got {type(config).__name__}" + ) + + return config diff --git a/libs/hackbot-runtime/tests/test_context.py b/libs/hackbot-runtime/tests/test_context.py index 26ef0b4d1d..29b0c10347 100644 --- a/libs/hackbot-runtime/tests/test_context.py +++ b/libs/hackbot-runtime/tests/test_context.py @@ -3,7 +3,7 @@ from pathlib import Path import pytest -from hackbot_runtime import HackbotContext +from hackbot_runtime import BaseAgentInputs, HackbotContext from hackbot_runtime.config import FirefoxConfig, HackbotConfig, SourceConfig @@ -13,6 +13,11 @@ def _hb(tmp_path, config: HackbotConfig) -> HackbotContext: return hb +class _SampleInputs(BaseAgentInputs): + bug_id: int + model: str | None = None + + def test_source_repo_without_declaration_raises(tmp_path): hb = _hb(tmp_path, HackbotConfig()) with pytest.raises(RuntimeError, match="\\[source\\]"): @@ -101,3 +106,51 @@ def test_results_plumbing(tmp_path): hb.actions.record("bugzilla.update_bug", {"bug_id": 1}, reasoning="r") assert hb.actions.actions[0]["type"] == "bugzilla.update_bug" + + +def test_load_inputs_without_url_reads_env(tmp_path, monkeypatch): + # Local/docker path: no RUN_INPUTS_URL, so inputs come from the environment. + monkeypatch.setenv("BUG_ID", "42") + monkeypatch.setenv("MODEL", "claude-opus") + hb = _hb(tmp_path, HackbotConfig()) + assert hb.run_inputs_url is None + + inputs = hb.load_inputs(_SampleInputs) + + assert inputs.bug_id == 42 + assert inputs.model == "claude-opus" + + +def test_load_inputs_uses_remote_config(tmp_path, monkeypatch): + # Production path: the required field is supplied by the fetched file, not env. + monkeypatch.delenv("BUG_ID", raising=False) + monkeypatch.delenv("MODEL", raising=False) + monkeypatch.setattr( + "hackbot_runtime.context.load_remote_config", + lambda url: {"bug_id": 7, "model": "from-config"}, + ) + hb = _hb(tmp_path, HackbotConfig()) + hb.run_inputs_url = "https://signed.example/inputs.json" + + inputs = hb.load_inputs(_SampleInputs) + + assert inputs.bug_id == 7 + assert inputs.model == "from-config" + + +def test_load_inputs_env_overrides_remote_config(tmp_path, monkeypatch): + # An env var wins over the same key in the config, while keys absent from the + # environment still fall through to the config. + monkeypatch.setenv("MODEL", "from-env") + monkeypatch.delenv("BUG_ID", raising=False) + monkeypatch.setattr( + "hackbot_runtime.context.load_remote_config", + lambda url: {"bug_id": 7, "model": "from-config"}, + ) + hb = _hb(tmp_path, HackbotConfig()) + hb.run_inputs_url = "https://signed.example/inputs.json" + + inputs = hb.load_inputs(_SampleInputs) + + assert inputs.model == "from-env" # env overrides config + assert inputs.bug_id == 7 # config supplies what env doesn't diff --git a/services/hackbot-api/app/agents.py b/services/hackbot-api/app/agents.py index 6ad1c5c00d..dfbcfb87a6 100644 --- a/services/hackbot-api/app/agents.py +++ b/services/hackbot-api/app/agents.py @@ -1,5 +1,3 @@ -import json -from collections.abc import Callable from dataclasses import dataclass from pydantic import BaseModel @@ -19,32 +17,6 @@ class AgentSpec: description: str job_name: str input_schema: type[BaseModel] - # Optional override for the rare agent whose env vars don't map 1:1 from - # its input schema. Defaults to ``model_to_env`` (field -> UPPER_SNAKE env). - build_env: Callable[[BaseModel], dict[str, str]] | None = None - - -def model_to_env(inputs: BaseModel) -> dict[str, str]: - """Serialise validated inputs into Cloud Run Job env overrides. - - Each schema field maps to an upper-cased env var (``bug_id`` -> ``BUG_ID``); - ``None`` fields are skipped, and the agent reads them back via - ``pydantic_settings.BaseSettings`` (which upper-cases field names by - default). Lists/dicts are JSON-encoded. Deploy-time constants (e.g. the - broker loopback URL) are NOT inputs — they belong in the Job's static env - config, not here. - """ - env: dict[str, str] = {} - for name, value in inputs.model_dump(mode="json").items(): - if value is None: - continue - if isinstance(value, str): - env[name.upper()] = value - elif isinstance(value, (list, dict)): - env[name.upper()] = json.dumps(value) - else: - env[name.upper()] = str(value) - return env AGENT_REGISTRY: dict[str, AgentSpec] = { diff --git a/services/hackbot-api/app/gcs.py b/services/hackbot-api/app/gcs.py index 49e720eff8..7ce4c048fc 100644 --- a/services/hackbot-api/app/gcs.py +++ b/services/hackbot-api/app/gcs.py @@ -11,6 +11,7 @@ from google.auth import impersonated_credentials from google.auth.transport.requests import Request as AuthRequest from google.cloud import storage +from pydantic import Json from app.config import settings from app.schemas import ArtifactRef, RunSummary @@ -26,6 +27,10 @@ def summary_blob_name(run_id: str) -> str: return f"{run_prefix(run_id)}summary.json" +def inputs_blob_name(run_id: str) -> str: + return f"{run_prefix(run_id)}inputs.json" + + @lru_cache(maxsize=1) def _signing_credentials() -> impersonated_credentials.Credentials: """Impersonate-self credentials so we can `sign_bytes` on Cloud Run. @@ -126,6 +131,34 @@ async def generate_results_policy(run_id: str) -> dict[str, Any]: return await asyncio.to_thread(_generate_post_policy_sync, run_id) +def _put_run_inputs_sync(run_id: str, config: dict[str, Json]) -> str: + """Upload the run's inputs and return the URL to fetch them. + + This reuses the GCS bucket configured for the agent results. + """ + bucket = _client().bucket(settings.results_bucket) + blob = bucket.blob(inputs_blob_name(run_id)) + blob.upload_from_string( + json.dumps(config), + content_type="application/json", + ) + # Valid for the whole run: the agent fetches config at startup, but matching + # the upload policy's window keeps a retried/slow-starting task covered. + expiration_seconds = ( + settings.job_execution_timeout_seconds + settings.signed_policy_grace_seconds + ) + return blob.generate_signed_url( + version="v4", + expiration=datetime.timedelta(seconds=expiration_seconds), + method="GET", + credentials=_signing_credentials(), + ) + + +async def put_run_inputs(run_id: str, config: dict[str, Json]) -> str: + return await asyncio.to_thread(_put_run_inputs_sync, run_id, config) + + def _read_summary_sync(run_id: str) -> RunSummary | None: bucket = _client().bucket(settings.results_bucket) blob = bucket.blob(summary_blob_name(run_id)) diff --git a/services/hackbot-api/app/routers/runs.py b/services/hackbot-api/app/routers/runs.py index ff55d466e8..5d0a2430a0 100644 --- a/services/hackbot-api/app/routers/runs.py +++ b/services/hackbot-api/app/routers/runs.py @@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app import gcs, jobs -from app.agents import AGENT_REGISTRY, AgentSpec, model_to_env +from app.agents import AGENT_REGISTRY, AgentSpec from app.auth import require_api_key from app.config import settings from app.database.connection import get_db @@ -77,13 +77,17 @@ async def create_run( db.add(run) await db.flush() + run_inputs_url = await gcs.put_run_inputs( + str(run_id), inputs.model_dump(mode="json") + ) + env_overrides: dict[str, str] = { "RUN_ID": str(run_id), "RESULTS_BUCKET": settings.results_bucket, "RESULTS_PREFIX": results_prefix, "RESULTS_POLICY_URL": policy["url"], "RESULTS_POLICY_FIELDS": json.dumps(policy["fields"]), - **(agent.build_env or model_to_env)(inputs), + "RUN_INPUTS_URL": run_inputs_url, } try: diff --git a/services/hackbot-api/tests/test_agents.py b/services/hackbot-api/tests/test_agents.py index 1e782b9963..0fc86cf9a6 100644 --- a/services/hackbot-api/tests/test_agents.py +++ b/services/hackbot-api/tests/test_agents.py @@ -1,9 +1,7 @@ -"""Tests for the agent registry and generic env serialization.""" - -import json +"""Tests for the agent registry.""" import pytest -from app.agents import AGENT_REGISTRY, model_to_env +from app.agents import AGENT_REGISTRY from app.schemas import ( BugFixInputs, BuildRepairInputs, @@ -14,80 +12,21 @@ from pydantic import ValidationError -def test_model_to_env_uppercases_and_stringifies(): - env = model_to_env(BugFixInputs(bug_id=12345, model="claude-opus", max_turns=8)) - assert env["BUG_ID"] == "12345" - assert env["MODEL"] == "claude-opus" - assert env["MAX_TURNS"] == "8" - - -def test_model_to_env_skips_none_fields(): - env = model_to_env(BugFixInputs(bug_id=1)) - assert env == {"BUG_ID": "1"} - # Optional fields left unset must not leak as empty/"None" env vars. - assert "MODEL" not in env - assert "EFFORT" not in env - - -def test_model_to_env_does_not_emit_deploy_constants(): - # The broker loopback URL is static Job config, not a per-run input. - env = model_to_env(BugFixInputs(bug_id=1, model="x", max_turns=2, effort="high")) - assert "BUGZILLA_MCP_URL" not in env - - -def test_bug_fix_registry_uses_default_env_serializer(): +def test_bug_fix_registry_entry(): spec = AGENT_REGISTRY["bug-fix"] - # No hand-written build_env: the router falls back to model_to_env. - assert spec.build_env is None assert spec.input_schema is BugFixInputs + assert spec.job_name == "hackbot-agent-bug-fix" def test_build_repair_registry_entry(): spec = AGENT_REGISTRY["build-repair"] - assert spec.build_env is None assert spec.input_schema is BuildRepairInputs assert spec.job_name == "hackbot-agent-build-repair" -def test_model_to_env_json_encodes_failure_tasks_and_bool(): - tasks = {"build-linux64/opt": "OyF95j0oQ-CF_YuBM1b7vg"} - env = model_to_env( - BuildRepairInputs( - bug_id=1, git_commit="deadbeef", failure_tasks=tasks, run_try_push=True - ) - ) - assert env["GIT_COMMIT"] == "deadbeef" - assert json.loads(env["FAILURE_TASKS"]) == tasks - assert env["RUN_TRY_PUSH"] == "True" - - def test_test_plan_generator_inputs_require_feature_description(): with pytest.raises(ValidationError): PlanGeneratorInputs( feature_name="Bookmarks and History", test_scope="Bookmarks toolbar behavior.", ) - - -def test_test_plan_generator_env_serialization(): - env = model_to_env( - PlanGeneratorInputs( - feature_name="Bookmarks and History", - feature_description="Bookmarks and history controls in Firefox.", - test_scope="Bookmarks toolbar behavior.", - ) - ) - - assert env == { - "FEATURE_NAME": "Bookmarks and History", - "FEATURE_DESCRIPTION": "Bookmarks and history controls in Firefox.", - "TEST_SCOPE": "Bookmarks toolbar behavior.", - } - - -def test_test_plan_generator_registry_uses_default_env_serializer(): - spec = AGENT_REGISTRY["test-plan-generator"] - - assert spec.build_env is None - assert spec.job_name == "hackbot-agent-test-plan-generator" - assert spec.input_schema is PlanGeneratorInputs