diff --git a/plugins/nemo-agents/openapi/openapi.yaml b/plugins/nemo-agents/openapi/openapi.yaml index 7b3ff0fe55..b4a9cc7848 100644 --- a/plugins/nemo-agents/openapi/openapi.yaml +++ b/plugins/nemo-agents/openapi/openapi.yaml @@ -5452,6 +5452,17 @@ components: type: string title: Input description: Prompt to pass to the agent. + environment: + anyOf: + - type: string + title: Reference + description: A reference to AgentEnvironmentInline. + - $ref: '#/components/schemas/AgentEnvironmentInline' + title: Environment + description: 'AgentEnvironment to run under: a "workspace/name" ref to a + stored AgentEnvironment, an inline environment, or None. Its EnvironmentSpec + is merged into the agent config and its ComputeSpec/secret refs are snapshotted + onto the job step at creation time.' workdir: allOf: - $ref: '#/components/schemas/AgentWorkdir' @@ -5475,6 +5486,13 @@ components: $ref: '#/components/schemas/ResolvedAgentConfig' workdir: $ref: '#/components/schemas/AgentWorkdir' + compute: + $ref: '#/components/schemas/ComputeSpecInline' + secrets: + additionalProperties: + type: string + type: object + title: Secrets type: object required: - request diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/jobs/execute.py b/plugins/nemo-agents/src/nemo_agents_plugin/jobs/execute.py index 0fd065e440..315d68f434 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/jobs/execute.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/jobs/execute.py @@ -15,7 +15,17 @@ from nemo_agents_plugin.agent_config import AgentConfig from nemo_agents_plugin.agent_config_formats import resolve_agent_config_for_deployment from nemo_agents_plugin.config import AgentsConfig -from nemo_agents_plugin.entities import NEMO_AGENTS_SPEC_CONFIG_FORMAT, Agent +from nemo_agents_plugin.entities import ( + NEMO_AGENTS_SPEC_CONFIG_FORMAT, + Agent, + AgentEnvironmentInline, + ComputeSpecInline, +) +from nemo_agents_plugin.environment_resolution import ( + EnvironmentResolutionError, + merge_environment_spec_into_agent_config, + resolve_environment, +) from nemo_agents_plugin.fabric.invocation import ( AgentConfigInvocationRequest, FabricDirectories, @@ -35,9 +45,28 @@ from nemo_platform_plugin.jobs.api_factory import ( ContainerSpec, CPUExecutionProviderSpec, + EnvironmentVariable, + EnvironmentVariableFromSecret, PlatformJobSpec, PlatformJobStep, + ResourcesLimitsSpec, + ResourcesSpec, ) +from nemo_platform_plugin.jobs.constants import ( + CONFIG_TASK_STORAGE_PATH_ENVVAR, + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, +) +from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_platform_plugin.jobs.image import get_qualified_image from nemo_platform_plugin.refs import ENTITY_REF_PATTERN, parse_entity_ref from pydantic import BaseModel, Field @@ -57,12 +86,62 @@ DEFAULT_AGENT_EXECUTION_TIMEOUT_SECONDS = 60 * 60 DEFAULT_AGENT_EXECUTION_IMAGE_NAME = "nmp-api" +# k8s resource key carrying the GPU count. The agents ``ComputeResources`` maps +# express GPUs the Kubernetes way (a ``nvidia.com/gpu`` entry in ``limits``); +# the jobs executor's ``ResourcesSpec`` instead carries a top-level integer +# ``num_gpus``. We translate that one key and pass ``cpu``/``memory`` through. +_GPU_RESOURCE_KEY = "nvidia.com/gpu" +# Resource keys we know how to translate onto the jobs ``ResourcesSpec``. Any +# other k8s resource key would be silently dropped (it has no home on the +# executor spec), so we reject it up front instead. +_SUPPORTED_RESOURCE_KEYS = frozenset({"cpu", "memory", _GPU_RESOURCE_KEY}) + +# Env var names a secret-backed env var must never shadow. Splitting a secret's +# resolved value over one of these would clobber platform-injected job state (the +# jobs substrate sets the ``NEMO_JOB_*``/``NMP_TASK_CONFIG`` family on every +# step) or the agent-container env the execute task relies on to reach the +# platform SDK (mirrors the deployment container's reserved set). Reject the +# collision at compile time so it can never reach the running step. +_RESERVED_ENV_VAR_NAMES = frozenset( + { + # Jobs substrate (nemo_platform_plugin.jobs.constants). + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + CONFIG_TASK_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + # Agent execution env (mirrors the deployment container's reserved set). + "NMP_WORKSPACE", + "NMP_AGENT_NAME", + "NMP_BASE_URL", + "PYTHONPATH", + "AGENT_CONFIG_PATH", + "NAT_CONFIG_PATH", + } +) + class ExecuteAgentJobConfig(BaseModel): model_config = {"json_schema_mode_override": "validation"} agent: str = Field(pattern=ENTITY_REF_PATTERN, description="Agent entity name or workspace/name ref.") input: str = Field(description="Prompt to pass to the agent.") + environment: str | AgentEnvironmentInline | None = Field( + default=None, + description=( + 'AgentEnvironment to run under: a "workspace/name" ref to a stored ' + "AgentEnvironment, an inline environment, or None. Its EnvironmentSpec " + "is merged into the agent config and its ComputeSpec/secret refs are " + "snapshotted onto the job step at creation time." + ), + ) workdir: AgentWorkdir | None = Field( default=None, description="Optional working directory configuration for the execution.", @@ -85,6 +164,14 @@ class ExecuteAgentStepConfig(BaseModel): request: ExecuteAgentJobConfig agent: ResolvedAgentConfig workdir: AgentWorkdir | None = None + # Immutable snapshot of the resolved AgentEnvironment, mirroring + # AgentDeployment. ``agent.config`` already holds the merged config; + # ``compute`` and ``secrets`` are snapshotted for the executor. The raw + # request environment (a ref or inline) is retained for provenance on + # ``request.environment`` — no separate field is needed here. Once created, + # the job is not kept in sync with the underlying environment entities. + compute: ComputeSpecInline | None = None + secrets: dict[str, str] = Field(default_factory=dict) class ExecuteAgentJob(NemoJob): @@ -133,7 +220,27 @@ async def to_spec( # type: ignore[override] workspace=agent.workspace, agent_name=agent.name, ) - _validate_agent_config(resolved_agent_config) + + # Resolve and merge the referenced AgentEnvironment. The EnvironmentSpec is + # merged into the resolved config (EnvironmentSpec-wins precedence); the + # ComputeSpec and secret-env references are snapshotted onto the step for + # the executor. This mirrors ``create_deployment`` so environment errors + # surface on the create request and the referenced entities are captured + # as an immutable snapshot. ``EnvironmentResolutionError`` is surfaced as a + # ``ValueError`` (matching the "Agent not found" pattern) so the jobs + # create path reports it to the caller. + try: + resolved_env = await resolve_environment( + request.environment, workspace=workspace, entity_client=typed_entity_client + ) + merged = merge_environment_spec_into_agent_config(resolved_agent_config, resolved_env.environment_spec) + except EnvironmentResolutionError as exc: + raise ValueError(str(exc)) from exc + + # Validate the merged config: an EnvironmentSpec can override the harness + # provider, so validation must run after the merge to reject e.g. a spec + # that selects a non-local Fabric environment. + _validate_agent_config(merged.config) workdir = None if request.workdir is not None: @@ -145,10 +252,12 @@ async def to_spec( # type: ignore[override] agent=ResolvedAgentConfig( name=agent.name, workspace=agent.workspace, - config=resolved_agent_config, + config=merged.config, config_format=agent.config_format, ), workdir=workdir, + compute=resolved_env.compute_spec, + secrets=merged.secrets, ) @classmethod @@ -165,21 +274,46 @@ async def compile( # type: ignore[override] ) -> PlatformJobSpec: del workspace, entity_client, job_name, async_sdk, options step_config = ExecuteAgentStepConfig.model_validate(spec) + + executor = CPUExecutionProviderSpec( + profile=profile or "default", + provider="cpu", + container=ContainerSpec( + image=cls._execution_image(), + entrypoint=["python", "-m"], + command=["nemo_agents_plugin.tasks.execute"], + ), + ) + # Snapshotted compute -> executor resources, and secret-env refs -> + # secret-backed step env vars. Both validate the snapshot and raise + # ``ValueError`` on a bad shape (an unsupported resource key, or a secret + # env name colliding with a reserved job env var). Surface those as a + # ``PlatformJobCompilationError`` so the jobs create route maps them to a + # descriptive 422 rather than an opaque 500 (the route's compile wrapper + # only translates ``PlatformJobCompilationError``, not bare ``ValueError``). + try: + # Snapshotted compute -> executor resources. Injected only when the + # environment supplied a compute spec, so the default CPU sizing is + # left to the jobs backend otherwise. + resources = _compute_to_resources(step_config.compute) + # Snapshotted secret-env refs -> secret-backed step env vars. The jobs + # substrate materializes each value into the process env under + # ENV_NAME; Fabric and its MCP servers read it by name (env-var-name + # indirection - see environment_resolution). + environment = _secret_environment(step_config.secrets) + except ValueError as exc: + raise PlatformJobCompilationError(str(exc)) from exc + + if resources is not None: + executor["resources"] = resources + return PlatformJobSpec( steps=[ PlatformJobStep( name="execute-agent", - executor=CPUExecutionProviderSpec( - profile=profile or "default", - provider="cpu", - container=ContainerSpec( - image=cls._execution_image(), - entrypoint=["python", "-m"], - command=["nemo_agents_plugin.tasks.execute"], - ), - ), + executor=executor, config=step_config.model_dump(mode="json"), - environment=[], + environment=environment, ) ], ) @@ -279,6 +413,95 @@ def _has_workdir_inputs(workdir: AgentWorkdir) -> bool: return workdir.base_workdir is not None or bool(workdir.artifact_mounts) +def _secret_environment(secrets: dict[str, str]) -> list[EnvironmentVariable]: + """Compile snapshotted secret-env refs into secret-backed step env vars. + + Each ``ENV_NAME -> "workspace/secret"`` entry becomes an + ``EnvironmentVariable`` whose value is populated from the referenced Secret. + A secret env name must never shadow a platform-injected env var (the jobs + substrate or agent-execution env - see ``_RESERVED_ENV_VAR_NAMES``): the + resolved value would clobber platform state, so reject the collision here. + """ + if not secrets: + return [] + + reserved = sorted(name for name in secrets if name in _RESERVED_ENV_VAR_NAMES) + if reserved: + raise ValueError( + f"Secret env var(s) {', '.join(reserved)} collide with reserved job env var name(s). " + f"Reserved names: {', '.join(sorted(_RESERVED_ENV_VAR_NAMES))}." + ) + + return [ + EnvironmentVariable(name=env_name, from_secret=EnvironmentVariableFromSecret(name=ref)) + for env_name, ref in secrets.items() + ] + + +def _compute_to_resources(compute: ComputeSpecInline | None) -> ResourcesSpec | None: + """Map a snapshotted agents ComputeSpec onto the jobs executor ResourcesSpec. + + Agents ``ComputeResources`` express requests/limits the Kubernetes way - + ``cpu``/``memory`` scalars plus a ``nvidia.com/gpu`` GPU count. The jobs + executor's ``ResourcesSpec`` instead carries ``cpu``/``memory`` under + ``limits``/``requests`` and a top-level integer ``num_gpus``. We translate + ``cpu``/``memory`` through and lift ``nvidia.com/gpu`` into ``num_gpus``, + preferring the limits count and falling back to requests. Any other k8s + resource key has no home on ``ResourcesSpec`` and would be silently dropped, + so it is rejected instead. + """ + if compute is None: + return None + + resources = compute.resources + _reject_unsupported_resource_keys(resources.limits, "limits") + _reject_unsupported_resource_keys(resources.requests, "requests") + + spec: ResourcesSpec = {} + limits = _cpu_memory_spec(resources.limits) + if limits: + spec["limits"] = limits + requests = _cpu_memory_spec(resources.requests) + if requests: + spec["requests"] = requests + + num_gpus = _gpu_count(resources.limits, resources.requests) + if num_gpus is not None: + spec["num_gpus"] = num_gpus + + return spec or None + + +def _reject_unsupported_resource_keys(resource_map: dict[str, str], where: str) -> None: + unsupported = sorted(key for key in resource_map if key not in _SUPPORTED_RESOURCE_KEYS) + if unsupported: + raise ValueError( + f"Unsupported compute resource key(s) in {where}: {', '.join(unsupported)}. " + f"agents.execute jobs support only {', '.join(sorted(_SUPPORTED_RESOURCE_KEYS))}." + ) + + +def _cpu_memory_spec(resource_map: dict[str, str]) -> ResourcesLimitsSpec: + # ``ResourcesLimitsSpec`` and ``ResourcesRequestsSpec`` are the same TypedDict + # (``ComputeResourceSpecParam``); one builder covers both sides. + spec: ResourcesLimitsSpec = {} + if "cpu" in resource_map: + spec["cpu"] = resource_map["cpu"] + if "memory" in resource_map: + spec["memory"] = resource_map["memory"] + return spec + + +def _gpu_count(limits: dict[str, str], requests: dict[str, str]) -> int | None: + raw = limits.get(_GPU_RESOURCE_KEY, requests.get(_GPU_RESOURCE_KEY)) + if raw is None: + return None + try: + return int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid {_GPU_RESOURCE_KEY!r} value {raw!r}; expected an integer GPU count.") from exc + + def _save_json_result(ctx: JobContext, name: str, path: Path, payload: dict[str, Any]) -> ResultRef: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") diff --git a/plugins/nemo-agents/tests/unit/test_execute_job.py b/plugins/nemo-agents/tests/unit/test_execute_job.py index 2b1b59dd9e..48b1b66317 100644 --- a/plugins/nemo-agents/tests/unit/test_execute_job.py +++ b/plugins/nemo-agents/tests/unit/test_execute_job.py @@ -12,7 +12,17 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from nemo_agents_plugin.entities import Agent +from nemo_agents_plugin.entities import ( + Agent, + AgentComputeSpec, + AgentEnvironment, + AgentEnvironmentInline, + AgentEnvironmentSpec, + ComputeResources, + ComputeSpecInline, + EnvironmentSpecInline, + McpFulfillment, +) from nemo_agents_plugin.fabric.runtime import FabricRuntimeResult from nemo_agents_plugin.jobs.execute import ( DEFAULT_AGENT_EXECUTION_TIMEOUT_SECONDS, @@ -36,6 +46,7 @@ from nemo_platform_plugin.dependencies import get_entity_client, get_sdk_client from nemo_platform_plugin.entity_client import NemoEntityNotFoundError from nemo_platform_plugin.job_context import JobContext +from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_platform_plugin.jobs.routes import add_job_routes @@ -450,6 +461,315 @@ async def test_compile_falls_back_to_qualified_api_image() -> None: assert step["executor"]["container"]["image"] == "qualified/nmp-api:dev" +# --- environment / compute / secrets wiring ------------------------------------ + + +def _env_ref_entity_client() -> AsyncMock: + """Entity client that dereferences an AgentEnvironment ref chain. + + Dispatches ``entity_client.get`` by entity type: the Agent, then the + AgentEnvironment and its EnvironmentSpec/ComputeSpec, all resolved by name. + """ + entity_client = AsyncMock() + + env = AgentEnvironment( + name="prod", + workspace="default", + environment_spec="default/prod-spec", + compute_spec="default/prod-compute", + ) + env_spec = AgentEnvironmentSpec( + name="prod-spec", + workspace="default", + secrets={"OPENAI_API_KEY": "default/openai-key"}, + ) + compute_spec = AgentComputeSpec( + name="prod-compute", + workspace="default", + resources=ComputeResources(limits={"cpu": "2", "memory": "4Gi", "nvidia.com/gpu": "1"}), + ) + + async def _get(entity_type: type, *, name: str, workspace: str) -> object: + if entity_type is Agent: + return _agent() + if entity_type is AgentEnvironment: + return env + if entity_type is AgentEnvironmentSpec: + return env_spec + if entity_type is AgentComputeSpec: + return compute_spec + raise AssertionError(f"unexpected entity type {entity_type!r}") + + entity_client.get.side_effect = _get + return entity_client + + +@pytest.mark.asyncio +async def test_to_spec_snapshots_resolved_environment_ref() -> None: + entity_client = _env_ref_entity_client() + + spec = await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello", environment="default/prod"), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + step_config = ExecuteAgentStepConfig.model_validate(spec) + # Raw environment retained for provenance on the stored request. + assert step_config.request.environment == "default/prod" + # Secret refs collected from the EnvironmentSpec top-level secrets. + assert step_config.secrets == {"OPENAI_API_KEY": "default/openai-key"} + # Compute spec snapshotted (k8s-style resource maps preserved as-is). + assert step_config.compute is not None + assert step_config.compute.resources.limits == {"cpu": "2", "memory": "4Gi", "nvidia.com/gpu": "1"} + + +@pytest.mark.asyncio +async def test_to_spec_merges_inline_environment_into_config() -> None: + entity_client = AsyncMock() + entity_client.get.return_value = _agent() + + environment = AgentEnvironmentInline( + environment_spec=EnvironmentSpecInline( + env={"MY_FLAG": "on"}, + secrets={"OPENAI_API_KEY": "default/openai-key"}, + ), + compute_spec=ComputeSpecInline(resources=ComputeResources(limits={"cpu": "500m"})), + ) + + spec = await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello", environment=environment), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + step_config = ExecuteAgentStepConfig.model_validate(spec) + # Inline env resolution never touches the entity store beyond the Agent. + entity_client.get.assert_awaited_once_with(Agent, name="calc", workspace="default") + # Plaintext env merged into the config's environment block. + assert step_config.agent.config["environment"]["env"] == {"MY_FLAG": "on"} + assert step_config.secrets == {"OPENAI_API_KEY": "default/openai-key"} + assert step_config.compute is not None + assert step_config.compute.resources.limits == {"cpu": "500m"} + + +@pytest.mark.asyncio +async def test_to_spec_mcp_secret_indirection_through_environment() -> None: + """An Agent-declared MCP server is fulfilled with a secret ref via env-name + indirection: the ref is collected into ``secrets`` (never written into the + config), and the server's env references the value by name so the running + step reads it from the process env the substrate populates. + """ + mcp_agent_config = { + "config_format": "nemo-agents-spec-v1", + "name": "calc", + "default_harness": "hermes", + "harnesses": {"hermes": {"kind": "hermes"}}, + "models": {"default": {"provider": "openai", "model": "openai/gpt-5.4"}}, + "mcp": {"servers": {"search": {"transport": "streamable-http", "url": "http://agent-url"}}}, + } + entity_client = AsyncMock() + entity_client.get.return_value = Agent( + name="calc", + workspace="default", + config=mcp_agent_config, + config_format="nemo-agents-spec-v1", + ) + + environment = AgentEnvironmentInline( + environment_spec=EnvironmentSpecInline( + mcp={ + "search": McpFulfillment( + url="http://env-url", + env={"SEARCH_MODE": "fast"}, + secrets={"SEARCH_TOKEN": "default/search-token"}, + ), + }, + ), + ) + + spec = await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello", environment=environment), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + step_config = ExecuteAgentStepConfig.model_validate(spec) + server = step_config.agent.config["mcp"]["servers"]["search"] + # Fulfillment url wins; non-secret env merged into the server config. + assert server["url"] == "http://env-url" + assert server["env"] == {"SEARCH_MODE": "fast"} + # Secret ref collected for injection, NOT written into the server config. + assert step_config.secrets == {"SEARCH_TOKEN": "default/search-token"} + assert "SEARCH_TOKEN" not in server.get("env", {}) + + +@pytest.mark.asyncio +async def test_to_spec_with_no_environment_leaves_snapshot_empty() -> None: + entity_client = AsyncMock() + entity_client.get.return_value = _agent() + + spec = await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello"), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + step_config = ExecuteAgentStepConfig.model_validate(spec) + assert step_config.request.environment is None + assert step_config.compute is None + assert step_config.secrets == {} + + +@pytest.mark.asyncio +async def test_to_spec_rejects_missing_environment_ref() -> None: + entity_client = AsyncMock() + + async def _get(entity_type: type, *, name: str, workspace: str) -> object: + if entity_type is Agent: + return _agent() + raise NemoEntityNotFoundError("missing") + + entity_client.get.side_effect = _get + + with pytest.raises(ValueError, match="AgentEnvironment 'prod' not found"): + await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello", environment="default/prod"), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + +@pytest.mark.asyncio +async def test_to_spec_rejects_environment_spec_selecting_non_local_provider() -> None: + entity_client = AsyncMock() + entity_client.get.return_value = _agent() + + environment = AgentEnvironmentInline( + environment_spec=EnvironmentSpecInline(provider="docker"), + ) + + with pytest.raises(ValueError, match="only support local Fabric environments"): + await ExecuteAgentJob.to_spec( + ExecuteAgentJobConfig(agent="calc", input="hello", environment=environment), + workspace="default", + entity_client=entity_client, + async_sdk=_sdk_with_files(), + is_local=False, + ) + + +@pytest.mark.asyncio +async def test_compile_injects_secret_env_and_compute_resources() -> None: + spec = ExecuteAgentStepConfig( + request=ExecuteAgentJobConfig(agent="calc", input="hello"), + agent=_resolved_agent(), + compute=ComputeSpecInline( + resources=ComputeResources( + limits={"cpu": "2", "memory": "4Gi", "nvidia.com/gpu": "2"}, + requests={"cpu": "1", "memory": "2Gi"}, + ) + ), + secrets={"OPENAI_API_KEY": "default/openai-key"}, + ) + + with patch("nemo_agents_plugin.jobs.execute.AgentsConfig.get") as get_config: + get_config.return_value.deployments.default_image = "registry.example/nmp-api:test" + platform_spec = await ExecuteAgentJob.compile( + workspace="default", + spec=spec, + entity_client=MagicMock(), + job_name=None, + async_sdk=MagicMock(), + ) + + step = list(platform_spec["steps"])[0] + # Secret ref -> secret-backed env var. + assert step["environment"] == [{"name": "OPENAI_API_KEY", "from_secret": {"name": "default/openai-key"}}] + # Compute -> executor resources: cpu/memory pass through, gpu -> num_gpus. + executor = cast(dict[str, Any], step["executor"]) + resources = executor["resources"] + assert resources["limits"] == {"cpu": "2", "memory": "4Gi"} + assert resources["requests"] == {"cpu": "1", "memory": "2Gi"} + assert resources["num_gpus"] == 2 + + +@pytest.mark.asyncio +async def test_compile_without_compute_omits_executor_resources() -> None: + spec = ExecuteAgentStepConfig( + request=ExecuteAgentJobConfig(agent="calc", input="hello"), + agent=_resolved_agent(), + ) + + with patch("nemo_agents_plugin.jobs.execute.AgentsConfig.get") as get_config: + get_config.return_value.deployments.default_image = "registry.example/nmp-api:test" + platform_spec = await ExecuteAgentJob.compile( + workspace="default", + spec=spec, + entity_client=MagicMock(), + job_name=None, + async_sdk=MagicMock(), + ) + + step = list(platform_spec["steps"])[0] + assert "resources" not in cast(dict[str, Any], step["executor"]) + assert step["environment"] == [] + + +@pytest.mark.asyncio +async def test_compile_rejects_unsupported_compute_resource_key() -> None: + spec = ExecuteAgentStepConfig( + request=ExecuteAgentJobConfig(agent="calc", input="hello"), + agent=_resolved_agent(), + compute=ComputeSpecInline(resources=ComputeResources(limits={"cpu": "1", "ephemeral-storage": "1Gi"})), + ) + + with ( + patch("nemo_agents_plugin.jobs.execute.AgentsConfig.get") as get_config, + pytest.raises(PlatformJobCompilationError, match="Unsupported compute resource key"), + ): + get_config.return_value.deployments.default_image = "registry.example/nmp-api:test" + await ExecuteAgentJob.compile( + workspace="default", + spec=spec, + entity_client=MagicMock(), + job_name=None, + async_sdk=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_compile_rejects_secret_env_colliding_with_reserved_name() -> None: + spec = ExecuteAgentStepConfig( + request=ExecuteAgentJobConfig(agent="calc", input="hello"), + agent=_resolved_agent(), + secrets={"NMP_BASE_URL": "default/some-secret"}, + ) + + with ( + patch("nemo_agents_plugin.jobs.execute.AgentsConfig.get") as get_config, + pytest.raises(PlatformJobCompilationError, match="reserved job env var name"), + ): + get_config.return_value.deployments.default_image = "registry.example/nmp-api:test" + await ExecuteAgentJob.compile( + workspace="default", + spec=spec, + entity_client=MagicMock(), + job_name=None, + async_sdk=MagicMock(), + ) + + def test_run_without_input_workdir_saves_empty_input_snapshot(ctx: JobContext) -> None: spec = ExecuteAgentStepConfig( request=ExecuteAgentJobConfig(agent="calc", input="hello"), @@ -776,6 +1096,9 @@ async def _create_job(*, workspace: str, body: object) -> MagicMock: "provider": "local", "workspace": "./workspace", "artifacts": "./artifacts", + "connection": {}, + "env": {}, + "metadata": {}, "settings": {}, } assert ( @@ -785,9 +1108,55 @@ async def _create_job(*, workspace: str, body: object) -> MagicMock: assert body.spec["request"] == { "agent": "calc", "input": "hello", + "environment": None, "workdir": {"base_workdir": "source#project", "artifact_mounts": []}, "timeout_seconds": DEFAULT_AGENT_EXECUTION_TIMEOUT_SECONDS, } assert body.spec["workdir"] == {"base_workdir": "default/source#project/", "artifact_mounts": []} assert body.platform_spec.steps[0].name == "execute-agent" assert response.json()["spec"]["workdir"]["base_workdir"] == "default/source#project/" + + +def test_execute_job_create_route_maps_reserved_secret_env_to_422() -> None: + """A reserved-name secret-env collision must surface as a 422, not a 500. + + The collision is detected in ``compile`` (``_secret_environment``). The jobs + create route only translates ``PlatformJobCompilationError`` into a 422 - a + bare ``ValueError`` escaping ``compile`` would fall through to the global + handler as an opaque 500. This guards that the error reaches the client as a + descriptive 422 at the HTTP boundary (the unit test on ``_secret_environment`` + only checks the raised exception, not the mapped status code). + """ + app = FastAPI() + app.include_router(add_job_routes(ExecuteAgentJob), prefix="/apis/agents/v2/workspaces/{workspace}") + + # Agent resolves, plus an environment whose EnvironmentSpec maps a secret env + # var onto a reserved job env var name (NMP_BASE_URL). + env = AgentEnvironment(name="prod", workspace="default", environment_spec="default/prod-spec") + env_spec = AgentEnvironmentSpec( + name="prod-spec", workspace="default", secrets={"NMP_BASE_URL": "default/some-secret"} + ) + + async def _get(entity_type: type, *, name: str, workspace: str) -> object: + if entity_type is Agent: + return _agent() + if entity_type is AgentEnvironment: + return env + if entity_type is AgentEnvironmentSpec: + return env_spec + raise AssertionError(f"unexpected entity type {entity_type!r}") + + entity_client = AsyncMock() + entity_client.get.side_effect = _get + app.dependency_overrides[get_entity_client] = lambda: entity_client + app.dependency_overrides[get_sdk_client] = lambda: _sdk_with_files() + + with patch("nemo_agents_plugin.jobs.execute.AgentsConfig.get") as get_config: + get_config.return_value.deployments.default_image = "registry.example/nmp-api:test" + response = TestClient(app, raise_server_exceptions=False).post( + "/apis/agents/v2/workspaces/default/jobs/execute", + json={"name": "execute-1", "spec": {"agent": "calc", "input": "hello", "environment": "default/prod"}}, + ) + + assert response.status_code == 422, response.text + assert "reserved job env var name" in response.json()["detail"]