diff --git a/sdk/demos/09_jungle_grid_gpu_execution/IMPLEMENTATION_DECISION.md b/sdk/demos/09_jungle_grid_gpu_execution/IMPLEMENTATION_DECISION.md new file mode 100644 index 000000000..05857178a --- /dev/null +++ b/sdk/demos/09_jungle_grid_gpu_execution/IMPLEMENTATION_DECISION.md @@ -0,0 +1,65 @@ +# Jungle Grid Integration Decision + +## Selected Extension Point + +This contribution is a runnable demo network with a Python `WorkerAgent`. The agent +uses OpenAgents' project mod for the long-running workflow, project messages for +estimate and lifecycle updates, and project artifacts for logs and Jungle Grid +artifact metadata. + +Jungle Grid is an external agentic AI workload execution and GPU orchestration +layer, not an OpenAgents transport, launcher agent type, or network mod. A demo +keeps the integration provider-specific while showing a reusable OpenAgents +pattern: an agent delegates asynchronous compute, waits for human approval +before billable work, and returns results to a shared project. + +## Rejected Alternatives + +- **Launcher agent type:** Jungle Grid executes workloads; it is not an interactive + coding-agent runtime managed by the launcher. +- **Core provider integration:** No OpenAgents core abstraction requires a + provider-specific compute backend. +- **Jungle Grid mod:** The integration does not add network-wide event semantics or + shared infrastructure. Existing project events already cover the workflow. +- **Hosted MCP entry:** Jungle Grid's hosted Streamable HTTP endpoint uses OAuth, + while local stdio uses an API key. The direct REST integration keeps approval + and project state inside OpenAgents without requiring an MCP auth change. +- **Local stdio MCP dependency:** The Jungle Grid stdio MCP package is supported, + but a direct Python API client is easier to validate, test, and constrain around + mandatory human approval. It also avoids requiring Node.js for a Python demo. + +## Jungle Grid Contract Used + +The demo uses the documented public execution API: + +- `POST /v1/jobs/estimate` +- `POST /v1/jobs` +- `GET /v1/jobs/{job_id}` +- `GET /v1/jobs/{job_id}/runtime` +- `GET /v1/jobs/{job_id}/logs` +- `POST /v1/jobs/{job_id}/cancel` +- `GET /v1/jobs/{job_id}/artifacts` +- `POST /v1/jobs/{job_id}/artifacts/{artifact_id}/download` + +Authentication is a scoped server-side API key in `JUNGLE_GRID_API_KEY`; the +REST base can be overridden with `JUNGLE_GRID_API`. The +documented lifecycle includes `pending`, `queued`, `assigned`, `running`, +`completed`, `failed`, `rejected`, and `cancelled`. + +The current REST request shape includes `model_size_gb`. Estimate responses +describe classification, routing, capacity, rates, cost ranges, queue waits, +start windows, warnings, and screening without starting compute. Managed +workloads can publish regular files from `/workspace/artifacts`; temporary +signed artifact download URLs are treated as secrets and are not stored in the +OpenAgents project. + +Workload environment values are not accepted in project goals. A goal may use +`environment_from_env` to reference variables available only in the executor +process; those values are resolved after human approval and are excluded from +the estimate request and project-visible output. + +## Contribution Workflow + +OpenAgents' contributing guide asks contributors to create an issue for feature +suggestions before submitting a pull request. This demo should be proposed in an +issue and held for maintainer direction before a PR is opened. diff --git a/sdk/demos/09_jungle_grid_gpu_execution/README.md b/sdk/demos/09_jungle_grid_gpu_execution/README.md new file mode 100644 index 000000000..599cf77ab --- /dev/null +++ b/sdk/demos/09_jungle_grid_gpu_execution/README.md @@ -0,0 +1,202 @@ +# Jungle Grid GPU Execution Demo + +This demo shows an OpenAgents execution agent delegating long-running AI and GPU +workloads to [Jungle Grid](https://junglegrid.dev), an agentic AI workload +execution and GPU orchestration layer that classifies intent, resolves capacity, +and places workloads without requiring agents to manage GPU servers. + +The workflow fits OpenAgents because the workload is asynchronous and +collaborative: an agent estimates the job, a human approves spending in the +shared project, and the agent returns lifecycle updates, logs, and artifact +metadata to the same workspace. + +## Security And Billing Warning + +Jungle Grid jobs may consume credits or incur charges. The executor never submits +a workload when a project starts. It requires an exact approval command from a +human identity after posting the estimate. Keep API keys in environment variables +and do not paste secrets into project goals, messages, logs, metadata, or +committed files. Workloads that need environment values must use +`environment_from_env`; the executor resolves those references only after human +approval, immediately before submission. + +## Prerequisites + +- Python with the OpenAgents development package installed. +- A Jungle Grid account and a scoped API key that can estimate, submit, read, and + cancel jobs. +- A public container image suitable for the requested workload. + +## Environment Variables + +- `JUNGLE_GRID_API_KEY` is required. The agent reads this server-side API key and + sends it only as a Bearer token to Jungle Grid. +- `JUNGLE_GRID_API` optionally overrides the default REST API base, + `https://api.junglegrid.dev`. +- Any workload-specific variables referenced by `environment_from_env` must also + be exported in the executor process. Their values are never placed in the + project goal or estimate request. + +## Setup + +From the repository root, install OpenAgents with SDK and development +dependencies so the network, agent, and test commands are available: + +```bash +pip install -e ".[sdk,dev]" +``` + +Export the Jungle Grid API key in the shell that will run the executor. This +keeps the credential out of the repository and network configuration: + +```bash +export JUNGLE_GRID_API_KEY="jg_..." +``` + +## Run The Demo + +The current demo assumes exactly one executor. Run one +`jungle-grid-executor` process so a project is estimated and submitted at most +once. + +Start the OpenAgents network from this demo directory. The network enables the +project mod and exposes the `Jungle Grid GPU Execution` project template: + +```bash +cd sdk/demos/09_jungle_grid_gpu_execution +openagents network start network.yaml +``` + +In a second terminal, start the deterministic Python executor. It does not need +an LLM provider key: + +```bash +cd sdk/demos/09_jungle_grid_gpu_execution +python agents/jungle_grid_executor.py +``` + +The script connects with the password hash configured for the `executors` +group. OpenAgents records that connection in +`network.topology.agent_group_membership`, which is the runtime source used by +the project mod. The optional `metadata.agents` list in an agent-group +configuration does not assign runtime membership and is intentionally not used +by this demo. + +Open Studio at `http://localhost:8700/studio`, create a project with the +`Jungle Grid GPU Execution` template, and use a JSON object as the project goal. +For example: + +```json +{ + "name": "openagents-batch-demo", + "workload_type": "batch", + "image": "python:3.11-slim", + "model_size_gb": 1, + "command": "python", + "args": ["-c", "print('hello from Jungle Grid')"], + "optimize_for": "cost" +} +``` + +The agent validates the request and calls the read-only +`POST /v1/jobs/estimate` endpoint. Current estimates include workload +classification, routing and capacity signals, hourly and total cost ranges, +queue-wait ranges, estimated start windows, warnings, and screening details. +The executor posts that structured estimate and stores it as project artifact +`jungle_grid_estimate`. No compute has been submitted at this point. + +For a workload that needs a credential or other environment value, export it in +the executor shell and reference only its local variable name in the goal: + +```bash +export MODEL_TOKEN="..." +``` + +```json +{ + "name": "openagents-inference-demo", + "workload_type": "inference", + "image": "example/model-server:latest", + "model_size_gb": 7, + "environment_from_env": { + "MODEL_TOKEN": "MODEL_TOKEN" + }, + "optimize_for": "cost" +} +``` + +The mapping key is the variable sent to the workload, and the mapping value is +the local executor variable to resolve. Literal `environment` values, API keys, +Bearer tokens, and secret-like metadata keys are rejected. + +Review the estimate, then reply in the project with the exact command shown by +the agent. Estimates that explicitly report `available: false` or +`can_submit: false` cannot be approved: + +```text +APPROVE +``` + +After approval, the agent submits with `POST /v1/jobs`, polls +`GET /v1/jobs/{job_id}`, and posts public lifecycle changes: pending, queued, +assigned, running, completed, failed, rejected, or cancelled. On a terminal +state it retrieves the runtime surface, the latest 100 stored log entries, and +the managed artifact list. Regular files written by managed workloads under +`/workspace/artifacts` are eligible for automatic upload. + +Artifact download requests mint temporary signed URLs. The executor requests +download metadata but redacts the URL before storing `jungle_grid_result`; do +not log or share signed URLs. + +To cancel a submitted job, reply with the exact job ID: + +```text +CANCEL +``` + +Cancellation is explicit and only applies when the job ID matches the project. +Only a human identity can request cancellation. The agent reports cancellation +failures without exposing the API key. + +## Failure Behavior + +Invalid workload JSON, missing required fields, missing API keys, timeouts, +invalid Jungle Grid responses, and API errors are posted to the project in +sanitized form. Failed, rejected, or cancelled jobs stop the OpenAgents project. +Completed jobs complete the project. + +The API key needs `jobs:estimate`, `jobs:submit`, `jobs:read`, and `logs:read` +capabilities for the complete flow. + +## Jungle Grid Interfaces + +This demo calls the REST API directly so OpenAgents can enforce project-based +human approval. Jungle Grid also provides the `jungle` CLI, whose `submit` +command estimates and asks for confirmation before queuing, and a hosted MCP +endpoint at `https://mcp.junglegrid.dev/mcp`. Hosted MCP uses OAuth; local stdio +MCP uses `JUNGLE_GRID_API_KEY`. The current MCP tools are `estimate_job`, +`submit_job`, `list_jobs`, `get_job`, `get_job_logs`, `cancel_job`, +`list_artifacts`, and `get_artifact`. + +## Tests + +Run the focused mocked tests. They do not contact Jungle Grid or submit paid +work: + +```bash +pytest tests/agents/test_jungle_grid_executor.py +``` + +Run the repository formatter and linter checks used by the Python project: + +```bash +ruff format --check sdk/demos/09_jungle_grid_gpu_execution tests/agents/test_jungle_grid_executor.py +ruff check sdk/demos/09_jungle_grid_gpu_execution tests/agents/test_jungle_grid_executor.py +``` + +## Optional Live Estimate + +The normal demo performs a live estimate when a project starts, but it never +automatically submits a job. Use a low-cost workload goal, review the estimate in +the project, and do not send the approval command unless you explicitly intend +to start billable compute. diff --git a/sdk/demos/09_jungle_grid_gpu_execution/agents/jungle_grid_executor.py b/sdk/demos/09_jungle_grid_gpu_execution/agents/jungle_grid_executor.py new file mode 100644 index 000000000..23348120b --- /dev/null +++ b/sdk/demos/09_jungle_grid_gpu_execution/agents/jungle_grid_executor.py @@ -0,0 +1,605 @@ +#!/usr/bin/env python3 +"""Jungle Grid execution agent for the OpenAgents project workflow demo.""" + +import asyncio +import json +import logging +import os +import re +import uuid +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Optional +from urllib.parse import quote + +import aiohttp + +from openagents.agents.worker_agent import WorkerAgent, on_event +from openagents.models.event_context import EventContext +from openagents.mods.workspace.project import DefaultProjectAgentAdapter + +logger = logging.getLogger(__name__) + +DEFAULT_API_BASE = "https://api.junglegrid.dev" +EXECUTORS_GROUP_PASSWORD_HASH = "8fba13dab71d6fdd8a9b9db1f06e81315dfbfd69167b6097f724604db3c91cdf" +TERMINAL_STATUSES = {"completed", "failed", "rejected", "cancelled"} +VALID_WORKLOAD_TYPES = {"inference", "training", "batch"} +VALID_OPTIMIZE_FOR = {"balanced", "cost", "speed"} +SUBMIT_FIELDS = { + "name", + "workload_type", + "image", + "model_size_gb", + "command", + "args", + "environment_from_env", + "optimize_for", + "constraints", + "template", + "metadata", +} +ESTIMATE_FIELDS = { + "name", + "workload_type", + "image", + "model_size_gb", + "command", + "args", + "optimize_for", + "constraints", + "template", +} +SENSITIVE_PATTERN = re.compile(r"(?i)(bearer\s+)[^\s,;]+|jg_[A-Za-z0-9_-]+") +SENSITIVE_KEY_PATTERN = re.compile(r"(?i)(api[_-]?key|authorization|password|secret|token)") + + +class JungleGridError(Exception): + """Sanitized Jungle Grid client error.""" + + def __init__(self, code: str, message: str, status: Optional[int] = None): + super().__init__(message) + self.code = code + self.status = status + + +def redact_sensitive(value: Any, secret: Optional[str] = None) -> str: + """Return a log-safe string with credentials removed.""" + text = str(value) + if secret: + text = text.replace(secret, "[REDACTED]") + return SENSITIVE_PATTERN.sub(lambda match: f"{match.group(1) or ''}[REDACTED]", text) + + +def _collect_string_values(value: Any) -> list[str]: + """Collect nested string values that must not be exposed in project output.""" + if isinstance(value, str): + return [value] if value else [] + if isinstance(value, dict): + strings = [] + for nested in value.values(): + strings.extend(_collect_string_values(nested)) + return strings + if isinstance(value, list): + strings = [] + for nested in value: + strings.extend(_collect_string_values(nested)) + return strings + return [] + + +def _contains_sensitive_key(value: Any) -> bool: + """Return whether nested data uses a key commonly associated with credentials.""" + if isinstance(value, dict): + return any( + SENSITIVE_KEY_PATTERN.search(str(key)) or _contains_sensitive_key(nested) for key, nested in value.items() + ) + if isinstance(value, list): + return any(_contains_sensitive_key(nested) for nested in value) + return False + + +def sanitize_project_data(value: Any, secrets: Iterable[str]) -> Any: + """Recursively redact credentials and workload-provided secret values.""" + secret_values = [secret for secret in secrets if secret] + if isinstance(value, str): + result = value + for secret in secret_values: + result = result.replace(secret, "[REDACTED]") + return redact_sensitive(result) + if isinstance(value, dict): + return {key: sanitize_project_data(nested, secret_values) for key, nested in value.items()} + if isinstance(value, list): + return [sanitize_project_data(nested, secret_values) for nested in value] + return value + + +def _unwrap_response(data: Any) -> Any: + if isinstance(data, dict) and data.get("ok") is True and "data" in data: + return data["data"] + return data + + +def _error_detail(data: Any, status: int) -> tuple[str, str]: + if isinstance(data, dict): + nested = data.get("error") + if isinstance(nested, dict): + return ( + redact_sensitive(nested.get("code") or "API_ERROR"), + redact_sensitive(nested.get("message") or f"HTTP {status}"), + ) + return ( + redact_sensitive(data.get("code") or "API_ERROR"), + redact_sensitive(data.get("message") or f"HTTP {status}"), + ) + return "API_ERROR", f"HTTP {status}" + + +class JungleGridClient: + """Small async client for Jungle Grid's documented public execution API.""" + + def __init__( + self, + api_base: Optional[str] = None, + timeout_seconds: float = 30.0, + ): + raw_api_base = api_base if api_base is not None else os.getenv("JUNGLE_GRID_API", DEFAULT_API_BASE) + self.api_key = os.getenv("JUNGLE_GRID_API_KEY", "").strip() + self.api_base = raw_api_base.rstrip("/") + self.timeout_seconds = timeout_seconds + + def _require_api_key(self) -> str: + if not self.api_key: + raise JungleGridError("MISSING_API_KEY", "JUNGLE_GRID_API_KEY is required.") + return self.api_key + + async def _request(self, method: str, path: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + api_key = self._require_api_key() + timeout = aiohttp.ClientTimeout(total=self.timeout_seconds) + headers = { + "Accept": "application/json", + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + try: + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.request(method, f"{self.api_base}{path}", headers=headers, json=payload) as response: + text = await response.text() + try: + data = json.loads(text) if text.strip() else {} + except json.JSONDecodeError as exc: + raise JungleGridError( + "INVALID_API_RESPONSE", "Jungle Grid returned invalid JSON.", response.status + ) from exc + if response.status < 200 or response.status >= 300: + code, message = _error_detail(data, response.status) + raise JungleGridError( + redact_sensitive(code, api_key), + redact_sensitive(message, api_key), + response.status, + ) + except asyncio.TimeoutError as exc: + raise JungleGridError("NETWORK_TIMEOUT", "Jungle Grid request timed out.") from exc + except aiohttp.ClientError as exc: + raise JungleGridError("NETWORK_ERROR", redact_sensitive(exc, api_key)) from exc + + result = _unwrap_response(data) + if not isinstance(result, dict): + raise JungleGridError("INVALID_API_RESPONSE", "Jungle Grid returned an unexpected response shape.") + return result + + async def estimate_job(self, workload: Dict[str, Any]) -> Dict[str, Any]: + return await self._request("POST", "/v1/jobs/estimate", workload) + + async def submit_job(self, workload: Dict[str, Any]) -> Dict[str, Any]: + return await self._request("POST", "/v1/jobs", workload) + + async def get_job(self, job_id: str) -> Dict[str, Any]: + return await self._request("GET", f"/v1/jobs/{quote(job_id, safe='')}") + + async def get_job_runtime(self, job_id: str) -> Dict[str, Any]: + return await self._request("GET", f"/v1/jobs/{quote(job_id, safe='')}/runtime") + + async def get_job_logs(self, job_id: str) -> Dict[str, Any]: + return await self._request("GET", f"/v1/jobs/{quote(job_id, safe='')}/logs?tail=100") + + async def cancel_job(self, job_id: str, reason: str) -> Dict[str, Any]: + return await self._request("POST", f"/v1/jobs/{quote(job_id, safe='')}/cancel", {"reason": reason}) + + async def list_artifacts(self, job_id: str) -> Dict[str, Any]: + return await self._request("GET", f"/v1/jobs/{quote(job_id, safe='')}/artifacts") + + async def get_artifact(self, job_id: str, artifact_id: str) -> Dict[str, Any]: + return await self._request( + "POST", + f"/v1/jobs/{quote(job_id, safe='')}/artifacts/{quote(artifact_id, safe='')}/download", + ) + + +def parse_workload_goal(goal: str) -> Dict[str, Any]: + """Parse and validate a project goal containing a Jungle Grid workload JSON object.""" + text = goal.strip() + if text.startswith("```"): + text = re.sub(r"^```(?:json)?\s*", "", text) + text = re.sub(r"\s*```$", "", text) + try: + workload = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError("Project goal must be a JSON object describing the Jungle Grid workload.") from exc + if not isinstance(workload, dict): + raise ValueError("Project goal must be a JSON object.") + if SENSITIVE_PATTERN.search(json.dumps(workload)): + raise ValueError("Workload must not contain API keys or Bearer tokens.") + + unsupported = sorted(set(workload) - SUBMIT_FIELDS) + if unsupported: + raise ValueError(f"Unsupported workload fields: {', '.join(unsupported)}.") + required = {"name", "workload_type", "image"} + missing = sorted(key for key in required if not isinstance(workload.get(key), str) or not workload[key].strip()) + if missing: + raise ValueError(f"Missing required workload fields: {', '.join(missing)}.") + model_size_gb = workload.get("model_size_gb") + if not isinstance(model_size_gb, (int, float)) or isinstance(model_size_gb, bool) or model_size_gb <= 0: + raise ValueError("model_size_gb must be a positive number.") + if workload["workload_type"] not in VALID_WORKLOAD_TYPES: + raise ValueError(f"workload_type must be one of: {', '.join(sorted(VALID_WORKLOAD_TYPES))}.") + if "optimize_for" in workload and workload["optimize_for"] not in VALID_OPTIMIZE_FOR: + raise ValueError(f"optimize_for must be one of: {', '.join(sorted(VALID_OPTIMIZE_FOR))}.") + if "args" in workload and not ( + isinstance(workload["args"], list) and all(isinstance(item, str) for item in workload["args"]) + ): + raise ValueError("args must be an array of strings.") + if "environment_from_env" in workload and not ( + isinstance(workload["environment_from_env"], dict) + and all( + isinstance(key, str) and key.strip() and isinstance(value, str) and value.strip() + for key, value in workload["environment_from_env"].items() + ) + ): + raise ValueError("environment_from_env must map workload variable names to local environment variable names.") + if _contains_sensitive_key(workload.get("metadata")): + raise ValueError("metadata must not contain secret-like keys.") + return workload + + +def build_estimate_payload(workload: Dict[str, Any]) -> Dict[str, Any]: + """Build an estimate-only payload without submit-only or secret-bearing fields.""" + return {key: value for key, value in workload.items() if key in ESTIMATE_FIELDS} + + +def build_submit_payload(workload: Dict[str, Any]) -> Dict[str, Any]: + """Build a submit payload, resolving secret environment values only at submission time.""" + payload = {key: value for key, value in workload.items() if key != "environment_from_env"} + references = workload.get("environment_from_env") + if not references: + return payload + + missing = sorted(env_name for env_name in references.values() if not os.getenv(env_name)) + if missing: + raise ValueError(f"Missing required local environment variables: {', '.join(missing)}.") + payload["environment"] = {name: os.environ[env_name] for name, env_name in references.items()} + return payload + + +def public_workload(workload: Dict[str, Any]) -> Dict[str, Any]: + """Return workload metadata safe to share in a project message or artifact.""" + result = dict(workload) + if "metadata" in result: + metadata = result["metadata"] + result["metadata"] = {key: "[REDACTED]" for key in metadata} if isinstance(metadata, dict) else "[REDACTED]" + return result + + +def lifecycle_label(status: str) -> str: + """Map Jungle Grid status to a user-facing lifecycle label.""" + if status == "assigned": + return "assigned (provisioning)" + return status + + +def estimate_can_submit(estimate: Dict[str, Any]) -> bool: + """Return whether an estimate explicitly permits submission.""" + return estimate.get("available") is not False and estimate.get("can_submit") is not False + + +@dataclass +class ProjectExecution: + """State tracked between estimate, approval, submission, and completion.""" + + project_id: str + workload: Dict[str, Any] + estimate_id: str + estimate: Dict[str, Any] + job_id: Optional[str] = None + last_status: Optional[str] = None + approved_by: Optional[str] = None + submission_started: bool = False + submit_payload: Optional[Dict[str, Any]] = None + secret_values: Optional[list[str]] = None + + +class JungleGridExecutorAgent(WorkerAgent): + """Execute approved Jungle Grid workloads and report results to an OpenAgents project.""" + + default_agent_id = "jungle-grid-executor" + + def __init__( + self, + jungle_grid_client: Optional[JungleGridClient] = None, + poll_interval_seconds: float = 10.0, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.jungle_grid = jungle_grid_client or JungleGridClient() + self.poll_interval_seconds = poll_interval_seconds + self.project_adapter = DefaultProjectAgentAdapter() + self.executions: Dict[str, ProjectExecution] = {} + self.monitor_tasks: Dict[str, asyncio.Task] = {} + + async def on_startup(self): + """Bind the project adapter after the OpenAgents client is connected.""" + self.project_adapter.bind_client(self.client) + self.project_adapter.bind_connector(self.client.connector) + self.project_adapter.bind_agent(self.agent_id) + logger.info("Jungle Grid executor is ready") + + async def on_shutdown(self): + """Stop local monitor tasks without cancelling remote jobs.""" + for task in self.monitor_tasks.values(): + task.cancel() + if self.monitor_tasks: + await asyncio.gather(*self.monitor_tasks.values(), return_exceptions=True) + + async def _post(self, project_id: str, text: str): + await self.project_adapter.send_project_message(project_id=project_id, content={"text": text}) + + async def _set_artifact(self, project_id: str, key: str, value: Dict[str, Any]): + await self.project_adapter.set_project_artifact( + project_id=project_id, key=key, value=json.dumps(value, indent=2) + ) + + def _project_secrets(self, execution: ProjectExecution) -> list[str]: + return [ + self.jungle_grid.api_key, + *(execution.secret_values or []), + *_collect_string_values(execution.workload.get("metadata")), + ] + + def _sanitize_for_project(self, value: Any, execution: ProjectExecution) -> Any: + return sanitize_project_data(value, self._project_secrets(execution)) + + def _is_human_approver(self, sender_id: str) -> bool: + return sender_id.startswith("human:") + + @on_event("project.notification.started") + async def handle_project_started(self, context: EventContext): + """Estimate a workload and request human approval without submitting it.""" + payload = context.incoming_event.payload + project_id = payload.get("project_id") + goal = payload.get("goal", "") + if not project_id: + return + try: + workload = parse_workload_goal(goal) + estimate = await self.jungle_grid.estimate_job(build_estimate_payload(workload)) + estimate_id = uuid.uuid4().hex[:12] + execution = ProjectExecution(project_id, workload, estimate_id, estimate) + self.executions[project_id] = execution + shared_workload = self._sanitize_for_project(public_workload(workload), execution) + shared_estimate = self._sanitize_for_project(estimate, execution) + await self._set_artifact( + project_id, + "jungle_grid_estimate", + {"estimate_id": estimate_id, "workload": shared_workload, "estimate": shared_estimate}, + ) + if not estimate_can_submit(estimate): + await self._post( + project_id, + "Jungle Grid estimate is not currently eligible for submission.\n\n" + f"```json\n{json.dumps({'estimate_id': estimate_id, 'workload': shared_workload, 'estimate': shared_estimate}, indent=2)}\n```", + ) + await self.project_adapter.stop_project( + project_id=project_id, reason="Jungle Grid estimate is not eligible for submission" + ) + return + await self._post( + project_id, + "Jungle Grid estimate ready. No job has been submitted.\n\n" + f"```json\n{json.dumps({'estimate_id': estimate_id, 'workload': shared_workload, 'estimate': shared_estimate}, indent=2)}\n```\n\n" + f"A human must reply exactly `APPROVE {estimate_id}` before billable compute can start.", + ) + except (ValueError, JungleGridError) as exc: + await self._post( + project_id, f"Jungle Grid estimate failed: {redact_sensitive(exc, self.jungle_grid.api_key)}" + ) + await self.project_adapter.stop_project(project_id=project_id, reason="Jungle Grid estimate failed") + + @on_event("project.notification.message_received") + async def handle_project_message(self, context: EventContext): + """Handle explicit approval and cancellation commands.""" + payload = context.incoming_event.payload + project_id = payload.get("project_id") + sender_id = str(payload.get("sender_id", "")) + content = payload.get("content", {}) + text = content.get("text", "") if isinstance(content, dict) else "" + if not project_id or not isinstance(text, str): + return + command = text + execution = self.executions.get(project_id) + + if command.startswith("APPROVE "): + if not execution: + await self._post(project_id, "There is no pending Jungle Grid estimate for this project.") + return + if not self._is_human_approver(sender_id): + await self._post( + project_id, "Approval rejected: billable Jungle Grid submission requires a human approver." + ) + return + if command != f"APPROVE {execution.estimate_id}": + await self._post(project_id, "Approval rejected: estimate id does not match the pending estimate.") + return + if execution.submission_started: + suffix = f" as job `{execution.job_id}`" if execution.job_id else "" + await self._post(project_id, f"Jungle Grid submission has already been requested{suffix}.") + return + await self._submit_and_monitor(execution, sender_id) + return + + if command.startswith("CANCEL "): + if not execution or not execution.job_id: + await self._post(project_id, "There is no submitted Jungle Grid job to cancel for this project.") + return + if command != f"CANCEL {execution.job_id}": + await self._post(project_id, "Cancellation rejected: job id does not match this project.") + return + if not self._is_human_approver(sender_id): + await self._post( + project_id, "Cancellation rejected: Jungle Grid cancellation requires a human approver." + ) + return + try: + result = await self.jungle_grid.cancel_job( + execution.job_id, f"Requested from OpenAgents by {sender_id}" + ) + shared_result = self._sanitize_for_project(result, execution) + await self._post( + project_id, + f"Cancellation requested for Jungle Grid job `{execution.job_id}`.\n\n```json\n{json.dumps(shared_result, indent=2)}\n```", + ) + except JungleGridError as exc: + await self._post( + project_id, f"Jungle Grid cancellation failed: {redact_sensitive(exc, self.jungle_grid.api_key)}" + ) + + async def _submit_and_monitor(self, execution: ProjectExecution, approved_by: str): + execution.submission_started = True + execution.approved_by = approved_by + try: + execution.submit_payload = build_submit_payload(execution.workload) + execution.secret_values = _collect_string_values(execution.submit_payload.get("environment")) + result = await self.jungle_grid.submit_job(execution.submit_payload) + job_id = str(result.get("job_id") or result.get("id") or "").strip() + if not job_id: + raise JungleGridError("INVALID_API_RESPONSE", "Jungle Grid submit response did not include a job id.") + execution.job_id = job_id + execution.last_status = str(result.get("status") or "submitted") + await self._set_artifact( + execution.project_id, + "jungle_grid_submission", + { + "approved_by": approved_by, + "estimate_id": execution.estimate_id, + "submission": self._sanitize_for_project(result, execution), + }, + ) + await self._post( + execution.project_id, + f"Jungle Grid job submitted after approval by `{approved_by}`: `{job_id}` " + f"(status: `{lifecycle_label(execution.last_status)}`).", + ) + task = asyncio.create_task(self._monitor(execution)) + self.monitor_tasks[execution.project_id] = task + except (ValueError, JungleGridError) as exc: + await self._post( + execution.project_id, + f"Jungle Grid submission failed: {redact_sensitive(exc, self.jungle_grid.api_key)}", + ) + await self.project_adapter.stop_project( + project_id=execution.project_id, reason="Jungle Grid submission failed" + ) + + async def _monitor(self, execution: ProjectExecution): + assert execution.job_id + try: + while True: + job = await self.jungle_grid.get_job(execution.job_id) + status = str(job.get("status") or "unknown") + if status != execution.last_status: + execution.last_status = status + await self._post( + execution.project_id, + f"Jungle Grid job `{execution.job_id}` is now `{lifecycle_label(status)}`.", + ) + if status in TERMINAL_STATUSES: + await self._finalize(execution, job) + return + await asyncio.sleep(self.poll_interval_seconds) + except JungleGridError as exc: + await self._post( + execution.project_id, + f"Jungle Grid monitoring failed: {redact_sensitive(exc, self.jungle_grid.api_key)}", + ) + await self.project_adapter.stop_project( + project_id=execution.project_id, reason="Jungle Grid monitoring failed" + ) + finally: + self.monitor_tasks.pop(execution.project_id, None) + + async def _finalize(self, execution: ProjectExecution, job: Dict[str, Any]): + assert execution.job_id + runtime: Dict[str, Any] = {} + logs: Dict[str, Any] = {} + artifacts: Dict[str, Any] = {} + downloads = [] + try: + runtime = await self.jungle_grid.get_job_runtime(execution.job_id) + except JungleGridError as exc: + runtime = {"error": redact_sensitive(exc, self.jungle_grid.api_key)} + try: + logs = await self.jungle_grid.get_job_logs(execution.job_id) + except JungleGridError as exc: + logs = {"error": redact_sensitive(exc, self.jungle_grid.api_key)} + try: + artifacts = await self.jungle_grid.list_artifacts(execution.job_id) + for artifact in artifacts.get("artifacts", []): + if not isinstance(artifact, dict): + continue + artifact_id = str(artifact.get("artifact_id") or artifact.get("id") or "").strip() + if artifact_id: + download = await self.jungle_grid.get_artifact(execution.job_id, artifact_id) + if "url" in download: + download = {**download, "url": "[REDACTED]"} + downloads.append(download) + except JungleGridError as exc: + artifacts = {"error": redact_sensitive(exc, self.jungle_grid.api_key)} + + result = self._sanitize_for_project( + {"job": job, "runtime": runtime, "logs": logs, "artifacts": artifacts, "downloads": downloads}, + execution, + ) + await self._set_artifact(execution.project_id, "jungle_grid_result", result) + status = str(job.get("status") or "unknown") + await self._post( + execution.project_id, + f"Jungle Grid job `{execution.job_id}` finished with status `{status}`. " + "Logs and artifact metadata are stored in project artifact `jungle_grid_result`.", + ) + if status == "completed": + await self.project_adapter.complete_project( + project_id=execution.project_id, + summary=f"Jungle Grid job {execution.job_id} completed successfully.", + ) + else: + await self.project_adapter.stop_project( + project_id=execution.project_id, + reason=f"Jungle Grid job {execution.job_id} finished with status {status}.", + ) + + +async def main(): + """Run the Jungle Grid executor agent.""" + logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") + agent = JungleGridExecutorAgent() + try: + await agent.async_start( + network_host="localhost", + network_port=8700, + password_hash=EXECUTORS_GROUP_PASSWORD_HASH, + ) + while True: + await asyncio.sleep(3600) + finally: + await agent.async_stop() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/demos/09_jungle_grid_gpu_execution/network.yaml b/sdk/demos/09_jungle_grid_gpu_execution/network.yaml new file mode 100644 index 000000000..30c0f5012 --- /dev/null +++ b/sdk/demos/09_jungle_grid_gpu_execution/network.yaml @@ -0,0 +1,88 @@ +network: + name: JungleGridGPUExecution + mode: centralized + node_id: jungle-grid-gpu-execution-1 + initialized: true + transports: + - type: http + config: + port: 8700 + serve_studio: true + serve_mcp: true + - type: grpc + config: + port: 8600 + manifest_transport: http + recommended_transport: grpc + encryption_enabled: false + default_agent_group: guest + requires_password: false + agent_groups: + executors: + description: Agents allowed to execute Jungle Grid project workflows + password_hash: 8fba13dab71d6fdd8a9b9db1f06e81315dfbfd69167b6097f724604db3c91cdf + metadata: + permissions: + - execute_external_compute + mods: + - name: openagents.mods.workspace.default + enabled: true + config: + custom_events_enabled: true + - name: openagents.mods.workspace.project + enabled: true + config: + max_concurrent_projects: 5 + project_templates: + jungle_grid_execution: + name: Jungle Grid GPU Execution + description: Estimate, approve, execute, and monitor an AI workload on Jungle Grid + expose_as_tool: true + tool_name: run_jungle_grid_workload + tool_description: Start a Jungle Grid workload project. The task must be a JSON object with name, workload_type, image, and model_size_gb; use environment_from_env for workload environment values. + tool_mode: async + agent_groups: + - executors + context: | + This project delegates a long-running AI or GPU workload to Jungle Grid. + The executor estimates cost first and will not submit a job until a human + replies with the exact approval command shown in the project. Do not put + credentials in the goal; use environment_from_env to reference variables + available only in the executor process. + created_by_version: 0.9.3 + +network_profile: + discoverable: true + name: Jungle Grid GPU Execution + description: A demo of human-approved asynchronous AI and GPU workload delegation through Jungle Grid. + tags: + - demo + - jungle-grid + - gpu + - execution + - project + categories: + - demo + - workflow + country: Worldwide + required_openagents_version: 0.9.3 + capacity: 10 + authentication: + type: none + host: 0.0.0.0 + port: 8700 + +log_level: INFO +data_dir: ./data/jungle-grid-gpu-execution +runtime_limit: null +shutdown_timeout: 30 + +external_access: + default_agent_group: guest + auth_token: null + auth_token_env: null + instruction: null + exposed_tools: + - start_run_jungle_grid_workload + - get_result_run_jungle_grid_workload + excluded_tools: [] diff --git a/tests/agents/test_jungle_grid_executor.py b/tests/agents/test_jungle_grid_executor.py new file mode 100644 index 000000000..aa9bf248e --- /dev/null +++ b/tests/agents/test_jungle_grid_executor.py @@ -0,0 +1,622 @@ +"""Mocked tests for the Jungle Grid GPU execution demo agent.""" + +import asyncio +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import yaml + +from openagents.core.network import AgentNetwork +from openagents.models.event import Event +from openagents.models.event_context import EventContext +from openagents.models.network_config import AgentGroupConfig, NetworkConfig +from openagents.models.transport import TransportType +from openagents.mods.workspace.project.mod import DefaultProjectNetworkMod + +MODULE_PATH = ( + Path(__file__).parent.parent.parent + / "sdk" + / "demos" + / "09_jungle_grid_gpu_execution" + / "agents" + / "jungle_grid_executor.py" +) +NETWORK_CONFIG_PATH = MODULE_PATH.parent.parent / "network.yaml" +SPEC = importlib.util.spec_from_file_location("jungle_grid_executor", MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +assert SPEC and SPEC.loader +SPEC.loader.exec_module(MODULE) + +JungleGridClient = MODULE.JungleGridClient +JungleGridError = MODULE.JungleGridError +JungleGridExecutorAgent = MODULE.JungleGridExecutorAgent +ProjectExecution = MODULE.ProjectExecution +EXECUTORS_GROUP_PASSWORD_HASH = MODULE.EXECUTORS_GROUP_PASSWORD_HASH +build_estimate_payload = MODULE.build_estimate_payload +build_submit_payload = MODULE.build_submit_payload +estimate_can_submit = MODULE.estimate_can_submit +lifecycle_label = MODULE.lifecycle_label +parse_workload_goal = MODULE.parse_workload_goal +public_workload = MODULE.public_workload +redact_sensitive = MODULE.redact_sensitive +sanitize_project_data = MODULE.sanitize_project_data + + +def context(event_name, payload): + return EventContext( + incoming_event=Event(event_name=event_name, source_id="system", payload=payload), + event_threads={}, + incoming_thread_id="thread-1", + ) + + +def workload(): + return { + "name": "batch-demo", + "workload_type": "batch", + "image": "python:3.11-slim", + "model_size_gb": 1, + "command": "python", + "args": ["-c", "print(42)"], + "optimize_for": "cost", + } + + +class FakeJungleGridClient: + def __init__(self): + self.api_key = "test-api-key" + self.estimate_job = AsyncMock(return_value={"available": True, "estimated_cost_usd": {"min": 0.1, "max": 0.2}}) + self.submit_job = AsyncMock(return_value={"job_id": "job_123", "status": "queued"}) + self.get_job = AsyncMock(return_value={"job_id": "job_123", "status": "completed"}) + self.get_job_runtime = AsyncMock(return_value={"exit_code": 0, "stdout_tail": "done"}) + self.get_job_logs = AsyncMock(return_value={"items": [{"message": "done"}]}) + self.cancel_job = AsyncMock(return_value={"job_id": "job_123", "status": "cancelled", "cancelled": True}) + self.list_artifacts = AsyncMock( + return_value={"artifacts": [{"artifact_id": "artifact_1", "filename": "output.json"}]} + ) + self.get_artifact = AsyncMock( + return_value={ + "artifact": {"artifact_id": "artifact_1", "filename": "output.json"}, + "url": "https://example.test/file", + } + ) + + +def agent_with_mocks(fake=None): + agent = JungleGridExecutorAgent(jungle_grid_client=fake or FakeJungleGridClient(), poll_interval_seconds=0) + agent.project_adapter = AsyncMock() + agent.project_adapter.send_project_message = AsyncMock(return_value={"success": True}) + agent.project_adapter.set_project_artifact = AsyncMock(return_value={"success": True}) + agent.project_adapter.complete_project = AsyncMock(return_value={"success": True}) + agent.project_adapter.stop_project = AsyncMock(return_value={"success": True}) + return agent + + +@pytest.mark.asyncio +async def test_executor_group_membership_delivers_project_start_and_returns_estimate(): + network_yaml = yaml.safe_load(NETWORK_CONFIG_PATH.read_text()) + executor_group = network_yaml["network"]["agent_groups"]["executors"] + assert executor_group["password_hash"] == EXECUTORS_GROUP_PASSWORD_HASH + assert "agents" not in executor_group.get("metadata", {}) + + config = NetworkConfig( + name="JungleGridGroupTest", + default_agent_group="guest", + requires_password=False, + agent_groups={"executors": AgentGroupConfig(**executor_group)}, + ) + network = AgentNetwork.create_from_config(config) + registration = await network.register_agent( + agent_id="jungle-grid-executor", + transport_type=TransportType.HTTP, + metadata={"name": "Jungle Grid Executor"}, + certificate=None, + password_hash=EXECUTORS_GROUP_PASSWORD_HASH, + ) + assert registration.success + assert network.topology.agent_group_membership["jungle-grid-executor"] == "executors" + + project_mod = DefaultProjectNetworkMod() + project_mod.update_config( + { + "project_templates": { + "jungle_grid_execution": { + "name": "Jungle Grid GPU Execution", + "agent_groups": ["executors"], + } + } + } + ) + project_mod.initialize() + project_mod.bind_network(network) + assert project_mod._get_agents_in_group("executors") == ["jungle-grid-executor"] + + fake = FakeJungleGridClient() + executor = agent_with_mocks(fake) + delivered = [] + + async def deliver(event): + delivered.append(event) + if event.destination_id == "jungle-grid-executor": + await executor.handle_project_started( + EventContext( + incoming_event=event, + event_threads={}, + incoming_thread_id="project-start", + ) + ) + return SimpleNamespace(success=True) + + project_mod.send_event = AsyncMock(side_effect=deliver) + response = await project_mod.process_system_message( + Event( + event_name="project.start", + source_id="human:project-owner", + payload={ + "template_id": "jungle_grid_execution", + "goal": json.dumps(workload()), + "name": "Jungle Grid test", + }, + ) + ) + + assert response.success + assert "jungle-grid-executor" in response.data["authorized_agents"] + assert any( + event.event_name == "project.notification.started" + and event.destination_id == "jungle-grid-executor" + and event.payload["initiator_agent_id"] == "human:project-owner" + for event in delivered + ) + fake.estimate_job.assert_awaited_once_with(build_estimate_payload(workload())) + estimate_message = executor.project_adapter.send_project_message.await_args.kwargs["content"]["text"] + assert "Jungle Grid estimate ready" in estimate_message + assert "APPROVE" in estimate_message + + +@pytest.mark.asyncio +async def test_successful_estimate_flow_posts_estimate_and_requires_approval(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + + await agent.handle_project_started( + context("project.notification.started", {"project_id": "project-1", "goal": json.dumps(workload())}) + ) + + fake.estimate_job.assert_awaited_once_with(build_estimate_payload(workload())) + fake.submit_job.assert_not_awaited() + assert "project-1" in agent.executions + message = agent.project_adapter.send_project_message.await_args.kwargs["content"]["text"] + assert "No job has been submitted" in message + assert "APPROVE" in message + + +@pytest.mark.asyncio +async def test_unavailable_estimate_never_requests_approval_or_submits(): + fake = FakeJungleGridClient() + fake.estimate_job = AsyncMock(return_value={"available": False, "can_submit": False}) + agent = agent_with_mocks(fake) + + await agent.handle_project_started( + context("project.notification.started", {"project_id": "project-1", "goal": json.dumps(workload())}) + ) + + fake.submit_job.assert_not_awaited() + message = agent.project_adapter.send_project_message.await_args.kwargs["content"]["text"] + assert "not currently eligible for submission" in message + assert "APPROVE" not in message + agent.project_adapter.stop_project.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_approval_required_before_submit_and_non_human_is_rejected(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {"available": True}) + agent.executions["project-1"] = execution + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "agent:other", "content": {"text": "APPROVE estimate-1"}}, + ) + ) + + fake.submit_job.assert_not_awaited() + assert ( + "requires a human approver" in agent.project_adapter.send_project_message.await_args.kwargs["content"]["text"] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command", ["APPROVE estimate-2", " APPROVE estimate-1", "APPROVE estimate-1\n"]) +async def test_approval_requires_exact_command(command): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {"available": True}) + agent.executions["project-1"] = execution + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "human:user", "content": {"text": command}}, + ) + ) + + fake.submit_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_approved_submit_flow_starts_monitor(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {"available": True}) + agent.executions["project-1"] = execution + agent._monitor = AsyncMock() + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "human:user", "content": {"text": "APPROVE estimate-1"}}, + ) + ) + await asyncio.sleep(0) + + fake.submit_job.assert_awaited_once_with(workload()) + assert execution.job_id == "job_123" + agent._monitor.assert_awaited_once_with(execution) + + +@pytest.mark.asyncio +async def test_concurrent_matching_approvals_submit_only_once(): + fake = FakeJungleGridClient() + submit_started = asyncio.Event() + release_submit = asyncio.Event() + + async def delayed_submit(_workload): + submit_started.set() + await release_submit.wait() + return {"job_id": "job_123", "status": "queued"} + + fake.submit_job = AsyncMock(side_effect=delayed_submit) + agent = agent_with_mocks(fake) + agent._monitor = AsyncMock() + execution = ProjectExecution("project-1", workload(), "estimate-1", {"available": True}) + agent.executions["project-1"] = execution + approval = context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "human:user", "content": {"text": "APPROVE estimate-1"}}, + ) + + first = asyncio.create_task(agent.handle_project_message(approval)) + await submit_started.wait() + await agent.handle_project_message(approval) + release_submit.set() + await first + await asyncio.sleep(0) + + fake.submit_job.assert_awaited_once_with(workload()) + + +@pytest.mark.asyncio +async def test_status_polling_posts_updates_and_completes(): + fake = FakeJungleGridClient() + fake.get_job = AsyncMock( + side_effect=[ + {"job_id": "job_123", "status": "running"}, + {"job_id": "job_123", "status": "completed"}, + ] + ) + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123", last_status="queued") + + await agent._monitor(execution) + + texts = [call.kwargs["content"]["text"] for call in agent.project_adapter.send_project_message.await_args_list] + assert any("`running`" in text for text in texts) + assert any("`completed`" in text for text in texts) + agent.project_adapter.complete_project.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_workload_stops_project(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123") + + await agent._finalize(execution, {"job_id": "job_123", "status": "failed"}) + + agent.project_adapter.stop_project.assert_awaited_once() + agent.project_adapter.complete_project.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_logs_and_artifacts_are_stored_in_project_artifact(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + execution = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123") + + await agent._finalize(execution, {"job_id": "job_123", "status": "completed"}) + + fake.get_job_runtime.assert_awaited_once_with("job_123") + fake.get_job_logs.assert_awaited_once_with("job_123") + fake.list_artifacts.assert_awaited_once_with("job_123") + fake.get_artifact.assert_awaited_once_with("job_123", "artifact_1") + artifact_call = agent.project_adapter.set_project_artifact.await_args + assert artifact_call.kwargs["key"] == "jungle_grid_result" + assert "output.json" in artifact_call.kwargs["value"] + assert "stdout_tail" in artifact_call.kwargs["value"] + assert "https://example.test/file" not in artifact_call.kwargs["value"] + assert "[REDACTED]" in artifact_call.kwargs["value"] + + +@pytest.mark.asyncio +async def test_resolved_environment_values_are_redacted_from_results(monkeypatch): + monkeypatch.setenv("MODEL_TOKEN", "secret-value") + fake = FakeJungleGridClient() + fake.get_job_logs = AsyncMock(return_value={"items": [{"message": "token=secret-value"}]}) + agent = agent_with_mocks(fake) + requested = {**workload(), "environment_from_env": {"MODEL_TOKEN": "MODEL_TOKEN"}} + execution = ProjectExecution( + "project-1", + requested, + "estimate-1", + {}, + job_id="job_123", + submit_payload=build_submit_payload(requested), + secret_values=["secret-value"], + ) + + await agent._finalize(execution, {"job_id": "job_123", "status": "completed"}) + + artifact_value = agent.project_adapter.set_project_artifact.await_args.kwargs["value"] + assert "secret-value" not in artifact_value + assert "[REDACTED]" in artifact_value + + +@pytest.mark.asyncio +async def test_cancellation_uses_matching_job_id(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + agent.executions["project-1"] = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123") + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "human:user", "content": {"text": "CANCEL job_123"}}, + ) + ) + + fake.cancel_job.assert_awaited_once_with("job_123", "Requested from OpenAgents by human:user") + + +@pytest.mark.asyncio +async def test_non_human_cancellation_is_rejected(): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + agent.executions["project-1"] = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123") + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "agent:other", "content": {"text": "CANCEL job_123"}}, + ) + ) + + fake.cancel_job.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("command", ["CANCEL job_456", " CANCEL job_123", "CANCEL job_123\n"]) +async def test_cancellation_requires_exact_command(command): + fake = FakeJungleGridClient() + agent = agent_with_mocks(fake) + agent.executions["project-1"] = ProjectExecution("project-1", workload(), "estimate-1", {}, job_id="job_123") + + await agent.handle_project_message( + context( + "project.notification.message_received", + {"project_id": "project-1", "sender_id": "human:user", "content": {"text": command}}, + ) + ) + + fake.cancel_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_missing_api_key_is_reported_without_network_call(monkeypatch): + monkeypatch.delenv("JUNGLE_GRID_API_KEY", raising=False) + client = JungleGridClient() + with pytest.raises(JungleGridError, match="JUNGLE_GRID_API_KEY is required"): + await client.estimate_job(workload()) + + +def test_invalid_workload_is_rejected(): + with pytest.raises(ValueError, match="Missing required workload fields"): + parse_workload_goal('{"workload_type": "batch"}') + + +def test_workload_requires_positive_model_size(): + with pytest.raises(ValueError, match="model_size_gb"): + parse_workload_goal(json.dumps({**workload(), "model_size_gb": 0})) + + +def test_estimate_payload_matches_current_draft_job_fields(): + requested = { + **workload(), + "constraints": {"max_price_per_hour": 2.5, "preferred_gpu_family": "l4"}, + } + + assert build_estimate_payload(requested) == requested + + +def test_workload_rejects_literal_credentials_and_secret_like_metadata(): + with pytest.raises(ValueError, match="must not contain API keys"): + parse_workload_goal(json.dumps({**workload(), "command": "curl -H 'Bearer secret-value'"})) + with pytest.raises(ValueError, match="secret-like keys"): + parse_workload_goal(json.dumps({**workload(), "metadata": {"api_token": "secret-value"}})) + + +def test_build_submit_payload_resolves_environment_only_at_submission(monkeypatch): + monkeypatch.setenv("MODEL_TOKEN", "secret-value") + requested = {**workload(), "environment_from_env": {"MODEL_TOKEN": "MODEL_TOKEN"}} + + assert "environment_from_env" not in build_estimate_payload(requested) + assert build_submit_payload(requested)["environment"] == {"MODEL_TOKEN": "secret-value"} + assert public_workload(requested)["environment_from_env"] == {"MODEL_TOKEN": "MODEL_TOKEN"} + + +def test_build_submit_payload_rejects_missing_local_environment(monkeypatch): + monkeypatch.delenv("MISSING_MODEL_TOKEN", raising=False) + requested = {**workload(), "environment_from_env": {"MODEL_TOKEN": "MISSING_MODEL_TOKEN"}} + + with pytest.raises(ValueError, match="MISSING_MODEL_TOKEN"): + build_submit_payload(requested) + + +def test_secret_redaction_removes_api_keys_and_bearer_tokens(): + text = redact_sensitive("failed with Bearer abc123 and jg_super_secret", "jg_super_secret") + assert "abc123" not in text + assert "jg_super_secret" not in text + assert "[REDACTED]" in text + + +def test_public_workload_redacts_metadata_values(): + shared = public_workload({**workload(), "metadata": {"nested": {"value": "secret"}}}) + assert shared["metadata"] == {"nested": "[REDACTED]"} + assert "secret" not in json.dumps(shared) + + +def test_project_data_redaction_removes_nested_workload_secrets(): + result = sanitize_project_data( + {"logs": [{"message": "token=secret-value"}], "error": "Bearer test-api-key"}, + ["secret-value", "test-api-key"], + ) + assert "secret-value" not in json.dumps(result) + assert "test-api-key" not in json.dumps(result) + + +def test_estimate_can_submit_honors_explicit_unavailability(): + assert estimate_can_submit({"available": True, "can_submit": True}) + assert not estimate_can_submit({"available": False}) + assert not estimate_can_submit({"can_submit": False}) + + +@pytest.mark.parametrize( + ("status", "label"), + [ + ("submitted", "submitted"), + ("queued", "queued"), + ("assigned", "assigned (provisioning)"), + ("running", "running"), + ("completed", "completed"), + ("failed", "failed"), + ("rejected", "rejected"), + ("cancelled", "cancelled"), + ], +) +def test_lifecycle_labels(status, label): + assert lifecycle_label(status) == label + + +class FakeResponse: + def __init__(self, status, text): + self.status = status + self._text = text + + async def text(self): + return self._text + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + +class FakeSession: + def __init__(self, response=None, error=None, **kwargs): + self.response = response + self.error = error + + def request(self, *args, **kwargs): + if self.error: + raise self.error + return self.response + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + +@pytest.mark.asyncio +async def test_invalid_jungle_grid_response(monkeypatch): + monkeypatch.setenv("JUNGLE_GRID_API_KEY", "test-api-key") + monkeypatch.setattr(MODULE.aiohttp, "ClientSession", lambda **kwargs: FakeSession(FakeResponse(200, "not-json"))) + client = JungleGridClient() + + with pytest.raises(JungleGridError, match="invalid JSON"): + await client.get_job("job_123") + + +@pytest.mark.asyncio +async def test_network_timeout_is_sanitized(monkeypatch): + monkeypatch.setenv("JUNGLE_GRID_API_KEY", "test-api-key") + monkeypatch.setattr( + MODULE.aiohttp, + "ClientSession", + lambda **kwargs: FakeSession(error=asyncio.TimeoutError()), + ) + client = JungleGridClient() + + with pytest.raises(JungleGridError, match="timed out"): + await client.get_job("job_123") + + +@pytest.mark.asyncio +async def test_api_error_is_sanitized(monkeypatch): + monkeypatch.setenv("JUNGLE_GRID_API_KEY", "test-api-key") + body = json.dumps( + { + "error": { + "code": "provider_jg_private_backend", + "message": "Bearer test-api-key is not allowed", + } + } + ) + monkeypatch.setattr(MODULE.aiohttp, "ClientSession", lambda **kwargs: FakeSession(FakeResponse(403, body))) + client = JungleGridClient() + + with pytest.raises(JungleGridError) as exc_info: + await client.get_job("job_123") + assert "jg_private_backend" not in exc_info.value.code + assert "[REDACTED]" in exc_info.value.code + assert "test-api-key" not in str(exc_info.value) + + +def test_client_uses_documented_rest_api_environment(monkeypatch): + monkeypatch.setenv("JUNGLE_GRID_API", "https://orchestrator.example.test/") + monkeypatch.setenv("JUNGLE_GRID_API_KEY", "test-api-key") + + client = JungleGridClient() + + assert client.api_base == "https://orchestrator.example.test" + + +@pytest.mark.asyncio +async def test_client_uses_documented_runtime_and_log_routes(monkeypatch): + monkeypatch.setenv("JUNGLE_GRID_API_KEY", "test-api-key") + client = JungleGridClient() + client._request = AsyncMock(return_value={}) + + await client.get_job_runtime("job_123") + await client.get_job_logs("job_123") + + assert client._request.await_args_list[0].args == ("GET", "/v1/jobs/job_123/runtime") + assert client._request.await_args_list[1].args == ("GET", "/v1/jobs/job_123/logs?tail=100")