diff --git a/docs/WORKFLOWS.md b/docs/WORKFLOWS.md index 10f3c57e..404013e4 100644 --- a/docs/WORKFLOWS.md +++ b/docs/WORKFLOWS.md @@ -69,6 +69,30 @@ parent output substitution and validation. See `src/worker/executors/utils/graph_templates.py` for the templating contract. +## API task + +`taskType: api` performs a single HTTP request. By default it routes to the Nebula endpoint and authenticates with the worker's `NEBULA_API_TOKEN`. + +`spec.api.url` overrides the endpoint; when absent, the executor uses `NEBULA_API_BASE_URL` (appending `/v1/chat/completions`). `spec.api.headers` may supply an `Authorization` header directly. + +Credential handling: a caller-supplied `Authorization` header is always used as-is and never overwritten. With no header, `NEBULA_API_TOKEN` is injected only when the call is on the Nebula url (no custom `spec.api.url`) — the Nebula token is never sent to a custom endpoint. A Nebula-path call with no token available fails closed. + +```yaml +spec: + taskType: api + api: + method: POST + headers: + Content-Type: application/json + body: + model: gpt-4o + messages: + - role: user + content: Hello + response: + parse_json: true +``` + ## data_retrieval: type lumid `type: lumid` routes the retrieval through lumid-data-app (HTTP). Three diff --git a/examples/templates/api_two_stage.yaml b/examples/templates/api_two_stage.yaml index 70a0c770..a8a83f81 100644 --- a/examples/templates/api_two_stage.yaml +++ b/examples/templates/api_two_stage.yaml @@ -4,8 +4,9 @@ # Stage 1 calls a chat completion endpoint and returns raw text. # Stage 2 sends Stage 1's returned text as the next prompt. # -# NOTE: Configure NEBULA_API_BASE_URL and NEBULA_API_TOKEN on the worker before -# submitting. APIExecutor injects the Authorization header from NEBULA_API_TOKEN. +# Stage 1 names a custom endpoint via spec.api.url and supplies its own +# Authorization header. Stage 2 omits url and header, so it uses the Nebula +# defaults (NEBULA_API_BASE_URL + NEBULA_API_TOKEN). apiVersion: flowmesh/v1 kind: APITask @@ -19,8 +20,10 @@ spec: - name: stage-1 spec: api: + url: https://api.example.com/v1/chat/completions method: POST headers: + Authorization: Bearer Content-Type: application/json body: model: gpt-4o diff --git a/src/server/dispatcher/base.py b/src/server/dispatcher/base.py index b6db7f21..797f8dd8 100644 --- a/src/server/dispatcher/base.py +++ b/src/server/dispatcher/base.py @@ -25,6 +25,7 @@ ) from shared.tasks.placeholders import PLACEHOLDER_PATTERN from shared.tasks.specs import ( + ApiSpecStrict, ConditionSpec, SSHSpecStrict, SSHSpecTemplate, @@ -36,6 +37,7 @@ from ..services.metrics import MetricsRecorder from ..task.metadata import extract_model_dataset_names from ..task.models import TaskRecord, TaskStatus +from ..task.redact import REDACTED from ..task.runtime import TaskRuntime from ..utils.time import now_iso from .worker_selector import DEFAULT_WORKER_SELECTION, select_worker @@ -393,6 +395,20 @@ def dispatch_once(self, task_id: str) -> bool: ) return True + if self._has_redacted_credential(rendered_task.spec): + self._runtime.release_merge(task_id) + self.fail_task( + task_id, + "credential_not_retained", + payload={ + "error": ( + "the API credential was not retained across the server " + "restart; resubmit the workflow with the credential" + ) + }, + ) + return True + # Conditional execution: skip dispatch if condition not met if self._evaluate_condition_skip(task_id, rendered_task, record): return True @@ -1055,6 +1071,18 @@ def _collect_upstream_results( results[name] = envelope.result return results + def _has_redacted_credential(self, spec: TaskSpecStrict) -> bool: + """Whether an api spec carries a redacted credential placeholder.""" + if not isinstance(spec, ApiSpecStrict): + return False + api = spec.api + if api is None: + return False + headers = api.get("headers") + if not isinstance(headers, dict): + return False + return any(value == REDACTED for value in headers.values()) + def _resolve_upstream_task_ids( self, record: TaskRecord, spec: TaskSpecStrict ) -> dict[str, str] | None: diff --git a/src/server/task/models.py b/src/server/task/models.py index 0c64a20d..722f835a 100644 --- a/src/server/task/models.py +++ b/src/server/task/models.py @@ -1,12 +1,19 @@ import time from typing import Any -from pydantic import BaseModel, Field, computed_field +from pydantic import ( + BaseModel, + Field, + SerializerFunctionWrapHandler, + computed_field, + model_serializer, +) from shared.tasks import TaskEnvelopeTemplate from shared.tasks.worker_message import HardwareUsage from ..utils.time import now_iso +from .redact import redact_api, redact_raw_yaml TRAINING_TASK_TYPES = { "sft", @@ -161,6 +168,16 @@ def last_failed_worker(self) -> str | None: """The most recent worker to have failed this task.""" return self.failed_workers[-1] if self.failed_workers else None + @model_serializer(mode="wrap") + def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + data = handler(self) + data["raw_yaml"] = redact_raw_yaml(self.raw_yaml) + spec = self.task.spec + api = getattr(spec, "api", None) + if isinstance(api, dict): + data["task"]["spec"]["api"] = redact_api(api) + return data + class TaskInfo(TaskRecord): depends_on: list[str] = Field(description="Dependency task IDs.") diff --git a/src/server/task/redact.py b/src/server/task/redact.py new file mode 100644 index 00000000..d2da3c92 --- /dev/null +++ b/src/server/task/redact.py @@ -0,0 +1,84 @@ +"""Redaction of API credentials before a task record is serialized. + +The in-memory ``TaskRecord`` keeps the real credential so dispatch works; the +serializer applies this module so that every dump to Redis is redacted. + +``raw_yaml`` is re-emitted via ``yaml.safe_dump``, which does not preserve +comments, key order or original formatting. That is an accepted cost: the field +is a stored record and is never re-parsed, so losing formatting is fine. If the +YAML cannot be parsed, the whole field is redacted rather than storing text +that might contain a key. +""" + +from typing import Any + +import yaml + +REDACTED = "[REDACTED]" + +# Whole-key matches; a bare substring test would over-redact (e.g. "monkey"). +_SENSITIVE_KEYS = frozenset( + { + "authorization", + "token", + "api-key", + "api_key", + "apikey", + "secret", + "access_token", + "bearer", + "x-api-key", + } +) +_SENSITIVE_SUFFIXES = ("_key", "-key", "_token", "-token") + + +def _is_sensitive_key(name: str) -> bool: + lowered = name.lower() + if lowered in _SENSITIVE_KEYS: + return True + return any(lowered.endswith(suffix) for suffix in _SENSITIVE_SUFFIXES) + + +def _redact_value(value: Any) -> Any: + """Recursively redact credential values by key name at any depth.""" + if isinstance(value, dict): + return { + key: (REDACTED if _is_sensitive_key(str(key)) else _redact_value(val)) + for key, val in value.items() + } + if isinstance(value, list): + return [_redact_value(item) for item in value] + return value + + +def redact_api(api: dict[str, Any] | None) -> dict[str, Any] | None: + """Return a copy of an api spec with credential values replaced. + + Only the five credential-bearing locations are touched; the rest of the + spec is returned unchanged. The original mapping is never mutated. + """ + if not isinstance(api, dict): + return api + redacted = dict(api) + for field in ("headers", "params", "body", "json", "data"): + value = redacted.get(field) + if isinstance(value, (dict, list)): + redacted[field] = _redact_value(value) + return redacted + + +def redact_raw_yaml(raw_yaml: str) -> str: + """Redact credential values from the original workflow YAML text. + + The YAML is parsed and redacted with the same recursive key rule used for + the parsed spec, then re-emitted. If parsing fails, the whole field is + replaced with a marker rather than storing un-analysed text. + """ + try: + tree = yaml.safe_load(raw_yaml) + except yaml.YAMLError: + return REDACTED + if tree is None: + return raw_yaml + return yaml.safe_dump(_redact_value(tree), sort_keys=False) diff --git a/src/worker/executors/api_executor.py b/src/worker/executors/api_executor.py index 1a56c7b1..907216d9 100644 --- a/src/worker/executors/api_executor.py +++ b/src/worker/executors/api_executor.py @@ -19,12 +19,13 @@ class APIExecutor(Executor): - """Executor that performs a single HTTP request defined by task YAML. + """Performs a single HTTP request defined by task YAML. - Uses a class-level connection pool keyed by (base_url, timeout, verify_tls, - follow_redirects) so that repeated calls to the same endpoint (e.g. a trading - bot hitting QuantArena every few seconds) reuse the underlying TCP/TLS - connection instead of paying the handshake cost on every request. + Defaults to the Nebula endpoint via ``NEBULA_API_BASE_URL`` and authenticates + with ``NEBULA_API_TOKEN``. ``spec.api.url`` overrides the endpoint and + ``spec.api.headers`` may supply an ``Authorization`` header directly. A + custom ``spec.api.url`` requires its own credential: the Nebula token is + never sent to an endpoint the caller chose. """ name = "api" @@ -99,6 +100,7 @@ def run(self, task: ExecutorTask, out_dir: Path) -> APIResult: raise ExecutionError("spec.api must be a mapping") url = api_cfg.get("url") + custom_url = url is not None if url is None: url = os.getenv("NEBULA_API_BASE_URL") if not url: @@ -110,8 +112,21 @@ def run(self, task: ExecutorTask, out_dir: Path) -> APIResult: if not isinstance(headers, dict): raise ExecutionError("spec.api.headers must be a mapping") - token = os.getenv("NEBULA_API_TOKEN") - if token and not any(k.lower() == "authorization" for k in headers): + has_credential = any(k.lower() == "authorization" for k in headers) + if custom_url: + if not has_credential: + raise ExecutionError( + "spec.api.url names a custom endpoint but no credential was " + "supplied; set an Authorization header. The Nebula token is " + "never sent to an endpoint the caller chose." + ) + elif not has_credential: + token = os.getenv("NEBULA_API_TOKEN") + if not token: + raise ExecutionError( + "no credential configured: set an Authorization header or " + "NEBULA_API_TOKEN" + ) headers["Authorization"] = f"Bearer {token}" params = api_cfg.get("params") diff --git a/tests/server/test_redact.py b/tests/server/test_redact.py new file mode 100644 index 00000000..83fbba0b --- /dev/null +++ b/tests/server/test_redact.py @@ -0,0 +1,163 @@ +"""Tests for API credential redaction at serialization time.""" + +import json + +from server.task.models import TaskRecord +from server.task.redact import REDACTED, _is_sensitive_key, redact_api, redact_raw_yaml +from shared.tasks import TaskEnvelopeTemplate +from shared.tasks.specs import ApiSpecTemplate + + +def _api_task(api: dict) -> TaskEnvelopeTemplate: + return TaskEnvelopeTemplate.model_validate( + { + "apiVersion": "mloc/v1", + "kind": "Task", + "metadata": {"name": "t"}, + "spec": {"taskType": "api", "api": api}, + } + ) + + +def _record(api: dict, raw_yaml: str = "") -> TaskRecord: + return TaskRecord( + task_id="tsk-1", + workflow_id="wfl-1", + owner_id="owner", + raw_yaml=raw_yaml, + task=_api_task(api), + ) + + +class TestSensitiveKey: + def test_sensitive_keys_match(self) -> None: + for key in ( + "Authorization", + "authorization", + "token", + "api-key", + "api_key", + "apikey", + "secret", + "access_token", + "bearer", + "x-api-key", + "my_token", + "my_key", + ): + assert _is_sensitive_key(key), key + + def test_innocent_keys_do_not_match(self) -> None: + for key in ("monkey", "turkey", "keyword", "keys", "model", "messages"): + assert not _is_sensitive_key(key), key + + +class TestRedactApi: + def test_headers_redacted(self) -> None: + out = redact_api({"headers": {"Authorization": "Bearer SECRET"}}) + assert out is not None + assert out["headers"]["Authorization"] == REDACTED + + def test_nested_in_dict_redacted(self) -> None: + out = redact_api({"json": {"auth": {"token": "SECRET-NESTED"}}}) + assert out is not None + assert out["json"]["auth"]["token"] == REDACTED + + def test_nested_in_list_redacted(self) -> None: + out = redact_api({"json": [{"token": "SECRET"}]}) + assert out is not None + assert out["json"][0]["token"] == REDACTED + + def test_all_five_locations_redacted(self) -> None: + api = { + "headers": {"Authorization": "Bearer H"}, + "params": {"api_key": "P"}, + "body": {"secret": "B"}, + "json": {"token": "J"}, + "data": {"access_token": "D"}, + } + out = redact_api(api) + assert out is not None + for field in ("headers", "params", "body", "json", "data"): + assert list(out[field].values()) == [REDACTED], field + + def test_innocent_values_preserved(self) -> None: + out = redact_api({"json": {"model": "gpt", "monkey": "x"}}) + assert out is not None + assert out["json"]["model"] == "gpt" + assert out["json"]["monkey"] == "x" + + def test_original_not_mutated(self) -> None: + api = {"headers": {"Authorization": "Bearer SECRET"}} + redact_api(api) + assert api["headers"]["Authorization"] == "Bearer SECRET" + + +class TestRedactRawYaml: + def test_plain_scalar_redacted(self) -> None: + out = redact_raw_yaml("Authorization: Bearer SECRET\n") + assert "SECRET" not in out + assert REDACTED in out + + def test_block_scalar_redacted(self) -> None: + out = redact_raw_yaml("Authorization: |\n Bearer SECRET\n") + assert "SECRET" not in out + assert REDACTED in out + + def test_folded_scalar_redacted(self) -> None: + out = redact_raw_yaml("Authorization: >-\n Bearer SECRET\n") + assert "SECRET" not in out + assert REDACTED in out + + def test_nested_yaml_redacted(self) -> None: + out = redact_raw_yaml("api:\n json:\n token: SECRET\n") + assert "SECRET" not in out + + def test_innocent_yaml_preserved(self) -> None: + out = redact_raw_yaml("model: gpt-4o\nmonkey: x\n") + assert "gpt-4o" in out + assert "monkey: x" in out + + def test_unparseable_yaml_fails_closed(self) -> None: + assert redact_raw_yaml("a: [unclosed") == REDACTED + + +class TestTaskRecordSerializer: + def test_dump_redacts_all_five_locations(self) -> None: + api = { + "headers": {"Authorization": "Bearer H"}, + "params": {"api_key": "P"}, + "body": {"secret": "B"}, + "json": {"token": "J"}, + "data": {"access_token": "D"}, + } + rec = _record(api) + dumped = rec.model_dump() + dumped_api = dumped["task"]["spec"]["api"] + for field in ("headers", "params", "body", "json", "data"): + assert list(dumped_api[field].values()) == [REDACTED], field + + def test_dump_json_redacts(self) -> None: + rec = _record({"headers": {"Authorization": "Bearer SECRET"}}) + dumped = json.loads(rec.model_dump_json()) + assert dumped["task"]["spec"]["api"]["headers"]["Authorization"] == REDACTED + + def test_raw_yaml_redacted_in_dump(self) -> None: + rec = _record( + {"headers": {"Authorization": "Bearer SECRET"}}, + raw_yaml="api:\n headers:\n Authorization: Bearer SECRET\n", + ) + dumped = rec.model_dump() + assert "SECRET" not in dumped["raw_yaml"] + + def test_in_memory_keeps_real_credential(self) -> None: + rec = _record({"headers": {"Authorization": "Bearer SECRET"}}) + assert isinstance(rec.task.spec, ApiSpecTemplate) + assert rec.task.spec.api is not None + assert rec.task.spec.api["headers"]["Authorization"] == "Bearer SECRET" + + def test_no_credential_unchanged(self) -> None: + api = {"url": "http://x", "json": {"model": "gpt"}} + rec = _record(api) + dumped = rec.model_dump() + assert dumped["task"]["spec"]["api"] == api diff --git a/tests/worker/test_api_executor.py b/tests/worker/test_api_executor.py new file mode 100644 index 00000000..4dafbcfc --- /dev/null +++ b/tests/worker/test_api_executor.py @@ -0,0 +1,127 @@ +"""Tests for the API executor url override and Nebula credential handling.""" + +from pathlib import Path +from unittest.mock import patch + +import httpx +import pytest + +from shared.tasks.worker_message import WorkerTaskMessage +from worker.executors.api_executor import APIExecutor +from worker.executors.base_executor import ExecutionError + + +def _task_message(**spec_updates: object) -> WorkerTaskMessage: + payload = { + "task_id": "task-api", + "workflow_id": "wf-1", + "owner_id": "owner", + "assigned_worker": "worker-1", + "dispatched_at": "2026-03-22T00:00:00Z", + "task": { + "apiVersion": "mloc/v1", + "kind": "Task", + "metadata": {"name": "wf:api"}, + "spec": { + "taskType": "api", + "api": { + "method": "POST", + "body": {"messages": [{"role": "user", "content": "hi"}]}, + **spec_updates, + }, + }, + }, + } + return WorkerTaskMessage.model_validate(payload) + + +class _RecordingTransport(httpx.MockTransport): + """MockTransport that records the request it served.""" + + def __init__(self) -> None: + self.request: httpx.Request | None = None + super().__init__(self._handler) + + def _handler(self, request: httpx.Request) -> httpx.Response: + self.request = request + return httpx.Response( + 200, + json={ + "choices": [{"message": {"content": "hello"}}], + "usage": {"total_tokens": 3}, + }, + ) + + +def _run( + executor: APIExecutor, task: WorkerTaskMessage, transport: _RecordingTransport +) -> None: + with patch.object( + APIExecutor, "_get_client", return_value=httpx.Client(transport=transport) + ): + executor.run(task, Path("/tmp/out")) + + +class TestNebulaPath: + def test_no_url_no_header_uses_nebula_url_and_token( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("NEBULA_API_BASE_URL", "https://nebula.example.com") + monkeypatch.setenv("NEBULA_API_TOKEN", "nebula-token") + task = _task_message() + transport = _RecordingTransport() + _run(APIExecutor.__new__(APIExecutor), task, transport) + assert transport.request is not None + assert transport.request.url == "https://nebula.example.com/v1/chat/completions" + assert transport.request.headers["Authorization"] == "Bearer nebula-token" + + def test_no_url_with_header_preserves_header_and_skips_token( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("NEBULA_API_BASE_URL", "https://nebula.example.com") + monkeypatch.setenv("NEBULA_API_TOKEN", "nebula-token") + task = _task_message(headers={"Authorization": "Bearer custom"}) + transport = _RecordingTransport() + _run(APIExecutor.__new__(APIExecutor), task, transport) + assert transport.request is not None + assert transport.request.url == "https://nebula.example.com/v1/chat/completions" + assert transport.request.headers["Authorization"] == "Bearer custom" + + def test_neither_url_nor_base_url_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("NEBULA_API_BASE_URL", raising=False) + task = _task_message() + with pytest.raises(ExecutionError, match="spec.api.url or NEBULA_API_BASE_URL"): + _run(APIExecutor.__new__(APIExecutor), task, _RecordingTransport()) + + +class TestCustomUrl: + def test_custom_url_without_credential_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A custom endpoint must carry its own credential. + + The Nebula token is available here, so the failure proves it is withheld + rather than merely absent: a caller-chosen endpoint never receives it. + """ + monkeypatch.setenv("NEBULA_API_TOKEN", "nebula-token") + task = _task_message(url="https://custom.example.com/v1/chat/completions") + transport = _RecordingTransport() + with pytest.raises(ExecutionError, match="custom endpoint"): + _run(APIExecutor.__new__(APIExecutor), task, transport) + assert transport.request is None + + def test_custom_url_with_header_preserves_header( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("NEBULA_API_TOKEN", "nebula-token") + task = _task_message( + url="https://custom.example.com/v1/chat/completions", + headers={"Authorization": "Bearer custom"}, + ) + transport = _RecordingTransport() + _run(APIExecutor.__new__(APIExecutor), task, transport) + assert transport.request is not None + assert transport.request.url == "https://custom.example.com/v1/chat/completions" + assert transport.request.headers["Authorization"] == "Bearer custom"