diff --git a/temporalio/contrib/openai_agents/README.md b/temporalio/contrib/openai_agents/README.md index 83539044c..89a9cc970 100644 --- a/temporalio/contrib/openai_agents/README.md +++ b/temporalio/contrib/openai_agents/README.md @@ -487,6 +487,42 @@ A stateless factory that declares no parameters — like the `lambda: MCPServerS For network-accessible MCP servers, you can also use `HostedMCPTool` from the OpenAI Agents SDK, which uses an MCP client hosted by OpenAI. +## Secrets for Hosted Tools + +⚠️ **Experimental** - This functionality is subject to change prior to General Availability. + +Use `temporal_worker_env_ref()` for a hosted tool credential that should come from the worker's environment rather than being written into your workflow. Pass it the *name of an environment variable*, in place of the credential itself: + +```python +from agents import HostedMCPTool +from temporalio.contrib.openai_agents import temporal_worker_env_ref + +tool = HostedMCPTool( + tool_config={ + "type": "mcp", + "server_label": "my_server", + "server_url": "https://example.com/mcp", + "authorization": temporal_worker_env_ref("MY_MCP_TOKEN"), + } +) +``` + +Every worker that runs model activities must both set `MY_MCP_TOKEN` and name it as resolvable: + +```python +plugin = OpenAIAgentsPlugin(resolvable_worker_env_vars=["MY_MCP_TOKEN"]) +``` + +Names are matched exactly, with no globbing, and `"*"` anywhere in the list allows every environment variable on the worker. + +A reference can sit inside a larger value: in `"Bearer " + temporal_worker_env_ref("MY_MCP_TOKEN")`, the reference is replaced in place and the rest of the string is sent unchanged. + +The environment variable's value is substituted in these fields and no others: + +- `authorization`, and the value of each entry in `headers`, in a `HostedMCPTool`'s `tool_config` +- `value` in each entry of `network_policy.domain_secrets` under a hosted `ShellTool`'s `environment` +- `value` in each entry of `network_policy.domain_secrets` under the `container` in a `CodeInterpreterTool`'s `tool_config` + ## Sandbox Support ⚠️ **Pre-release** - This functionality is subject to change prior to General Availability. @@ -719,6 +755,7 @@ Certain tools are not suitable for a distributed computing environment, so these | HostedMCPTool | Yes | | ImageGenerationTool | Yes | | CodeInterpreterTool | Yes | +| ShellTool | Yes | | ComputerTool | No | #### Tool Context diff --git a/temporalio/contrib/openai_agents/__init__.py b/temporalio/contrib/openai_agents/__init__.py index 3976f633c..3305eaf95 100644 --- a/temporalio/contrib/openai_agents/__init__.py +++ b/temporalio/contrib/openai_agents/__init__.py @@ -13,6 +13,9 @@ OpenAIAgentsPlugin, OpenAIPayloadConverter, ) +from temporalio.contrib.openai_agents._temporal_worker_env_ref import ( + temporal_worker_env_ref, +) from temporalio.contrib.openai_agents.sandbox._sandbox_client_provider import ( SandboxClientProvider, ) @@ -28,6 +31,7 @@ "SandboxClientProvider", "StatelessMCPServerProvider", "StatefulMCPServerProvider", + "temporal_worker_env_ref", "testing", "workflow", ] diff --git a/temporalio/contrib/openai_agents/_invoke_model_activity.py b/temporalio/contrib/openai_agents/_invoke_model_activity.py index 5435b6369..b94d48beb 100644 --- a/temporalio/contrib/openai_agents/_invoke_model_activity.py +++ b/temporalio/contrib/openai_agents/_invoke_model_activity.py @@ -4,6 +4,7 @@ """ import enum +from collections.abc import Collection from dataclasses import dataclass from datetime import timedelta from typing import Any, NoReturn @@ -46,6 +47,9 @@ from temporalio import activity from temporalio.contrib.openai_agents._heartbeat_decorator import auto_heartbeater +from temporalio.contrib.openai_agents._temporal_worker_env_ref import ( + _WorkerEnvRefResolver, +) from temporalio.contrib.workflow_streams import WorkflowStreamClient from temporalio.exceptions import ApplicationError @@ -223,7 +227,7 @@ async def _noop_shell_executor(*_a: Any, **_kw: Any) -> str: return "" -def _build_tool(tool: ToolInput) -> Tool: +def _build_tool(tool: ToolInput, env_refs: _WorkerEnvRefResolver) -> Tool: """Reconstruct a Tool from its data-conversion-friendly input form.""" if isinstance( tool, @@ -231,22 +235,29 @@ def _build_tool(tool: ToolInput) -> Tool: FileSearchTool, WebSearchTool, ImageGenerationTool, - CodeInterpreterTool, LocalShellTool, ToolSearchTool, ), ): return tool + elif isinstance(tool, CodeInterpreterTool): + return CodeInterpreterTool( + tool_config=env_refs.resolve_code_interpreter_tool_config(tool.tool_config) + ) elif isinstance(tool, ShellToolInput): + environment = env_refs.resolve_shell_tool_environment(tool.environment) + # Only a local environment takes an executor. return ShellTool( name=tool.name, - environment=tool.environment, - executor=_noop_shell_executor, + environment=environment, + executor=_noop_shell_executor if environment["type"] == "local" else None, ) elif isinstance(tool, ApplyPatchToolInput): return ApplyPatchTool(name=tool.name, editor=_NoopApplyPatchEditor()) elif isinstance(tool, HostedMCPToolInput): - return HostedMCPTool(tool_config=tool.tool_config) + return HostedMCPTool( + tool_config=env_refs.resolve_mcp_tool_config(tool.tool_config) + ) elif isinstance(tool, CustomToolInput): return CustomTool( name=tool.tool_config["name"], @@ -269,8 +280,9 @@ def _build_tool(tool: ToolInput) -> Tool: def _build_tools_and_handoffs( input: ActivityModelInput, + env_refs: _WorkerEnvRefResolver, ) -> tuple[list[Tool], list[Handoff[Any, Any]]]: - tools = [_build_tool(x) for x in input.get("tools", [])] + tools = [_build_tool(x, env_refs) for x in input.get("tools", [])] handoffs: list[Handoff[Any, Any]] = [ Handoff( tool_name=x.tool_name, @@ -327,18 +339,23 @@ class ModelActivity: Disabling retries in your model of choice is recommended to allow activity retries to define the retry model. """ - def __init__(self, model_provider: ModelProvider | None = None): + def __init__( + self, + model_provider: ModelProvider | None = None, + resolvable_worker_env_vars: Collection[str] = (), + ): """Initialize the activity with a model provider.""" self._model_provider = model_provider or OpenAIProvider( openai_client=AsyncOpenAI(max_retries=0) ) + self._env_refs = _WorkerEnvRefResolver(resolvable_worker_env_vars) @activity.defn @auto_heartbeater async def invoke_model_activity(self, input: ActivityModelInput) -> ModelResponse: """Activity that invokes a model with the given input.""" model = self._model_provider.get_model(input.get("model_name")) - tools, handoffs = _build_tools_and_handoffs(input) + tools, handoffs = _build_tools_and_handoffs(input, self._env_refs) try: return await model.get_response( @@ -382,7 +399,7 @@ async def invoke_model_activity_streaming( trip ``heartbeat_timeout``. """ model = self._model_provider.get_model(input.get("model_name")) - tools, handoffs = _build_tools_and_handoffs(input) + tools, handoffs = _build_tools_and_handoffs(input, self._env_refs) topic = input["streaming_topic"] batch_interval = input.get( diff --git a/temporalio/contrib/openai_agents/_temporal_openai_agents.py b/temporalio/contrib/openai_agents/_temporal_openai_agents.py index 43594657f..54beb9d75 100644 --- a/temporalio/contrib/openai_agents/_temporal_openai_agents.py +++ b/temporalio/contrib/openai_agents/_temporal_openai_agents.py @@ -3,7 +3,7 @@ import dataclasses import json import typing -from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from collections.abc import AsyncIterator, Callable, Collection, Iterator, Sequence from contextlib import asynccontextmanager, contextmanager from datetime import timedelta @@ -219,6 +219,7 @@ def __init__( register_activities: bool = True, add_temporal_spans: bool = True, use_otel_instrumentation: bool = False, + resolvable_worker_env_vars: Collection[str] = (), ) -> None: """Initialize the OpenAI agents plugin. @@ -246,6 +247,12 @@ def __init__( use_otel_instrumentation: If set to true, enable open telemetry instrumentation. Warning: use_otel_instrumentation is experimental and behavior may change in future versions. Use with caution in production environments. + resolvable_worker_env_vars: Names of the environment variables that + ``temporal_worker_env_ref()`` may read on this worker. Names are + matched exactly, with no globbing; ``"*"`` + anywhere in the collection allows every name. + Warning: resolvable_worker_env_vars is experimental and behavior may change in future versions. + Use with caution in production environments. """ if model_params is None: @@ -275,7 +282,9 @@ def add_activities( if not register_activities: return activities or [] - model_activity = ModelActivity(model_provider) + model_activity = ModelActivity( + model_provider, resolvable_worker_env_vars=resolvable_worker_env_vars + ) new_activities = [ model_activity.invoke_model_activity, model_activity.invoke_model_activity_streaming, diff --git a/temporalio/contrib/openai_agents/_temporal_worker_env_ref.py b/temporalio/contrib/openai_agents/_temporal_worker_env_ref.py new file mode 100644 index 000000000..503db3587 --- /dev/null +++ b/temporalio/contrib/openai_agents/_temporal_worker_env_ref.py @@ -0,0 +1,136 @@ +"""References to secrets held in the Temporal Worker's environment.""" + +from __future__ import annotations + +import os +import re +from collections.abc import Collection, Mapping, MutableMapping +from typing import Any, cast + +from agents.tool import ShellToolContainerAutoEnvironment, ShellToolEnvironment +from openai.types.responses.tool_param import CodeInterpreter, Mcp + +_REF_PREFIX = "temporal.worker_env_ref:" + +_REF_PATTERN = re.compile(re.escape(_REF_PREFIX) + r"\{([^}{]*)\}") + +_ANY_ENV_VAR = "*" + + +def temporal_worker_env_ref(name: str) -> str: + """Refer to a secret held in the Temporal Worker's environment. + + .. warning:: + This function is experimental and may change in future versions. + Use with caution in production environments. + + Use it for a hosted tool credential that should come from the worker's + environment rather than being written into your workflow. Put the returned + reference in a ``HostedMCPTool``'s ``authorization`` or header value, or in + the ``value`` of a ``domain_secrets`` entry under a ``ShellTool`` or + ``CodeInterpreterTool``. The reference carries only the variable's name, + never its value. + + Every worker that runs model activities must set the variable and name it in + ``OpenAIAgentsPlugin(resolvable_worker_env_vars=[...])``. A name the worker + allows but has not set resolves to an empty value. + + Args: + name: Name of the environment variable to read on the worker. + + Returns: + A reference string to use in place of the secret. + """ + return f"{_REF_PREFIX}{{{name}}}" + + +class _WorkerEnvRefResolver: # type:ignore[reportUnusedClass] + def __init__(self, resolvable_worker_env_vars: Collection[str]) -> None: + if isinstance(resolvable_worker_env_vars, str): + raise TypeError( + "resolvable_worker_env_vars takes a collection of environment variable " + 'names, such as ["MY_MCP_TOKEN"]. A single string is read as the ' + "collection of its characters, so pass a list even for one name." + ) + self._allowed = frozenset(resolvable_worker_env_vars) + + def _resolve_ref(self, value: str) -> str: + def substitute(match: re.Match[str]) -> str: + name = match.group(1) + if _ANY_ENV_VAR not in self._allowed and name not in self._allowed: + return match.group(0) + return os.environ.get(name, "") + + return _REF_PATTERN.sub(substitute, value) + + def _resolve_domain_secret(self, secret: Mapping[str, Any]) -> dict[str, Any]: + return {**secret, "value": self._resolve_ref(secret["value"])} + + def _resolve_network_policy(self, network_policy: Any) -> Any: + policy = cast(MutableMapping[str, Any], network_policy) + domain_secrets = policy.get("domain_secrets") + if domain_secrets is None: + return network_policy + unresolved = list(domain_secrets) + # On the code interpreter path pydantic deserializes domain_secrets into a + # single-pass iterator, so the entries read here go back onto the input. + policy["domain_secrets"] = unresolved + return { + **policy, + "domain_secrets": [ + self._resolve_domain_secret(secret) for secret in unresolved + ], + } + + def resolve_mcp_tool_config(self, tool_config: Mcp) -> Mcp: + resolved: Mcp = tool_config + if "authorization" in resolved: + resolved = { + **resolved, + "authorization": self._resolve_ref(resolved["authorization"]), + } + headers = resolved.get("headers") + if headers is not None: + resolved = { + **resolved, + "headers": { + name: self._resolve_ref(value) for name, value in headers.items() + }, + } + return resolved + + def resolve_shell_tool_environment( + self, + environment: ShellToolEnvironment | None, + ) -> ShellToolEnvironment: + """An absent environment comes back as the local one ``ShellTool`` normalizes it to.""" + if environment is None: + return {"type": "local"} + if environment.get("type") != "container_auto": + return environment + auto = cast(ShellToolContainerAutoEnvironment, environment) + network_policy = auto.get("network_policy") + if network_policy is None: + return environment + return { + **auto, + "network_policy": self._resolve_network_policy(network_policy), + } + + def resolve_code_interpreter_tool_config( + self, + tool_config: CodeInterpreter, + ) -> CodeInterpreter: + container = tool_config.get("container") + if not isinstance(container, Mapping): + return tool_config + network_policy = container.get("network_policy") + if network_policy is None: + return tool_config + return { + **tool_config, + "container": { + **container, + "network_policy": self._resolve_network_policy(network_policy), + }, + } diff --git a/temporalio/contrib/openai_agents/testing.py b/temporalio/contrib/openai_agents/testing.py index 110ca20b7..05d7864d5 100644 --- a/temporalio/contrib/openai_agents/testing.py +++ b/temporalio/contrib/openai_agents/testing.py @@ -1,6 +1,6 @@ """Testing utilities for OpenAI agents.""" -from collections.abc import AsyncIterator, Callable, Sequence +from collections.abc import AsyncIterator, Callable, Collection, Sequence from typing import Any from agents import ( @@ -181,6 +181,7 @@ def __init__( register_activities: bool = True, add_temporal_spans: bool = True, use_otel_instrumentation: bool = False, + resolvable_worker_env_vars: Collection[str] = (), ) -> None: """Initialize the AgentEnvironment. @@ -201,6 +202,10 @@ def __init__( use_otel_instrumentation: If set to true, enable open telemetry instrumentation. Warning: use_otel_instrumentation is experimental and behavior may change in future versions. Use with caution in production environments. + resolvable_worker_env_vars: Names of the environment variables that + ``temporal_worker_env_ref()`` may read on this environment's workers. + Warning: resolvable_worker_env_vars is experimental and behavior may change in future versions. + Use with caution in production environments. """ self._model_params = model_params self._model_provider = None @@ -213,6 +218,7 @@ def __init__( self._plugin: OpenAIAgentsPlugin | None = None self._add_temporal_spans = add_temporal_spans self._use_otel_instrumentation = use_otel_instrumentation + self._resolvable_worker_env_vars = resolvable_worker_env_vars async def __aenter__(self) -> "AgentEnvironment": """Enter the async context manager.""" @@ -224,6 +230,7 @@ async def __aenter__(self) -> "AgentEnvironment": register_activities=self._register_activities, add_temporal_spans=self._add_temporal_spans, use_otel_instrumentation=self._use_otel_instrumentation, + resolvable_worker_env_vars=self._resolvable_worker_env_vars, ) return self diff --git a/tests/contrib/openai_agents/test_openai.py b/tests/contrib/openai_agents/test_openai.py index 25597ee55..79f7f9dcd 100644 --- a/tests/contrib/openai_agents/test_openai.py +++ b/tests/contrib/openai_agents/test_openai.py @@ -96,6 +96,9 @@ from temporalio.contrib.openai_agents._temporal_model_stub import ( _TemporalModelStub, ) +from temporalio.contrib.openai_agents._temporal_worker_env_ref import ( + _WorkerEnvRefResolver, +) from temporalio.contrib.openai_agents.testing import ( AgentEnvironment, ResponseBuilders, @@ -2713,7 +2716,7 @@ class FakeSandboxSession: tool_inputs = activity_input.get("tools") or [] assert len(tool_inputs) == 1 - rebuilt = _build_tool(tool_inputs[0]) + rebuilt = _build_tool(tool_inputs[0], _WorkerEnvRefResolver(())) assert isinstance(rebuilt, CustomTool) assert rebuilt.name == tool.name assert rebuilt.description == tool.description @@ -2753,7 +2756,7 @@ async def stub(_ctx: Any, _payload: str) -> str: tool_inputs = activity_input.get("tools") or [] assert len(tool_inputs) == 1 - rebuilt = _build_tool(tool_inputs[0]) + rebuilt = _build_tool(tool_inputs[0], _WorkerEnvRefResolver(())) assert isinstance(rebuilt, CustomTool) assert rebuilt.tool_config == tool.tool_config assert rebuilt.defer_loading is True diff --git a/tests/contrib/openai_agents/test_openai_tool_secrets.py b/tests/contrib/openai_agents/test_openai_tool_secrets.py new file mode 100644 index 000000000..de4a3e0c4 --- /dev/null +++ b/tests/contrib/openai_agents/test_openai_tool_secrets.py @@ -0,0 +1,867 @@ +"""Tests for worker environment references in hosted tool secrets.""" + +import time +import uuid +from collections.abc import AsyncIterator, Collection +from typing import Any, cast + +import pytest +from agents import ( + Agent, + AgentOutputSchemaBase, + CodeInterpreterTool, + Handoff, + HostedMCPTool, + Model, + ModelResponse, + ModelSettings, + ModelTracing, + Runner, + Tool, + TResponseInputItem, + Usage, +) +from agents.items import TResponseStreamEvent +from agents.tool import ShellTool, ShellToolEnvironment + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.openai_agents import ( + ModelActivityParameters, + OpenAIAgentsPlugin, + OpenAIPayloadConverter, + temporal_worker_env_ref, +) +from temporalio.contrib.openai_agents._invoke_model_activity import ( + ActivityModelInput, + ModelActivity, + StreamingActivityModelInput, + _build_tool, +) +from temporalio.contrib.openai_agents._temporal_model_stub import _TemporalModelStub +from temporalio.contrib.openai_agents._temporal_worker_env_ref import ( + _WorkerEnvRefResolver, +) +from temporalio.contrib.openai_agents.testing import ( + AgentEnvironment, + TestModelProvider, +) +from temporalio.testing import ActivityEnvironment +from tests.helpers import new_worker + +SENTINEL = "sk-test-sentinel-4f1a9c7e2b" +ENV_NAME = "TEMPORAL_TEST_TOOL_SECRET" +OTHER_SENTINEL = "sk-test-other-8c3d5e0a1f" +OTHER_ENV_NAME = "TEMPORAL_TEST_OTHER_TOOL_SECRET" + +_RESOLVER_ALLOWING_TEST_NAMES = _WorkerEnvRefResolver([ENV_NAME, OTHER_ENV_NAME]) + + +def _round_trip_activity_input(tool: Tool) -> tuple[bytes, ActivityModelInput]: + stub = _TemporalModelStub( + model_name="gpt-5", + model_params=ModelActivityParameters(), + agent=None, + ) + activity_input, _summary = stub._build_activity_input( + system_instructions=None, + input="hi", + model_settings=ModelSettings(), + tools=[tool], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) + converter = OpenAIPayloadConverter() + payload = converter.to_payload(activity_input) + return payload.data, converter.from_payload(payload, ActivityModelInput) + + +def _activity_input_payload_and_tool(tool: Tool) -> tuple[bytes, Any]: + payload, received = _round_trip_activity_input(tool) + tools = received.get("tools") or [] + assert len(tools) == 1 + return payload, tools[0] + + +def _hosted_mcp_tool(authorization: str, header_value: str) -> HostedMCPTool: + return HostedMCPTool( + tool_config={ + "type": "mcp", + "server_label": "test_server", + "server_url": "https://example.com/mcp", + "authorization": authorization, + "headers": {"X-Token": header_value, "X-Plain": "not-a-secret"}, + } + ) + + +def _domain_secret(name: str, value: str) -> dict[str, str]: + return {"domain": "example.com", "name": name, "value": value} + + +def _network_policy(domain_secrets: tuple[Any, ...]) -> Any: + return { + "type": "allowlist", + "allowed_domains": ["example.com"], + "domain_secrets": list(domain_secrets), + } + + +def _shell_tool(*domain_secrets: Any) -> ShellTool: + environment: Any = { + "type": "container_auto", + "network_policy": _network_policy(domain_secrets), + } + return ShellTool(environment=environment) + + +def _code_interpreter_tool(*domain_secrets: Any) -> CodeInterpreterTool: + tool_config: Any = { + "type": "code_interpreter", + "container": { + "type": "auto", + "network_policy": _network_policy(domain_secrets), + }, + } + return CodeInterpreterTool(tool_config=tool_config) + + +def _as_dict(value: Any) -> dict[str, Any]: + return cast(dict[str, Any], value) + + +def _secrets_in(network_policy: Any) -> list[Any]: + return list(_as_dict(network_policy)["domain_secrets"]) + + +def _shell_secrets(built: Tool) -> list[Any]: + assert isinstance(built, ShellTool) + assert built.environment is not None + return _secrets_in(_as_dict(built.environment)["network_policy"]) + + +def _code_interpreter_secrets(built: Tool) -> list[Any]: + assert isinstance(built, CodeInterpreterTool) + container = _as_dict(built.tool_config)["container"] + return _secrets_in(_as_dict(container)["network_policy"]) + + +def test_hosted_mcp_secrets_stay_out_of_activity_arguments( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + + payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + assert SENTINEL.encode() not in payload + assert payload.count(ref.encode()) == 2 + assert received.tool_config["authorization"] == ref + assert received.tool_config["headers"]["X-Token"] == ref + + +def test_hosted_shell_domain_secret_stays_out_of_activity_arguments( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + + payload, received = _activity_input_payload_and_tool( + _shell_tool(_domain_secret("TOKEN", ref)) + ) + + assert SENTINEL.encode() not in payload + assert ref.encode() in payload + secrets = _secrets_in(received.environment["network_policy"]) + assert secrets[0]["value"] == ref + + +def test_code_interpreter_domain_secret_stays_out_of_activity_arguments( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + + payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool(_domain_secret("TOKEN", ref)) + ) + + assert SENTINEL.encode() not in payload + assert ref.encode() in payload + secrets = _secrets_in(received.tool_config["container"]["network_policy"]) + assert secrets[0]["value"] == ref + + +def test_hosted_mcp_secrets_resolve_for_the_model_call( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == SENTINEL + assert _as_dict(built.tool_config)["headers"] == { + "X-Token": SENTINEL, + "X-Plain": "not-a-secret", + } + assert received.tool_config["authorization"] == ref + assert received.tool_config["headers"]["X-Token"] == ref + + +@pytest.mark.parametrize( + ("resolvable", "resolves"), + [([ENV_NAME], True), (["*"], True), ([OTHER_ENV_NAME], False)], + ids=["the_name", "star", "another_name"], +) +def test_a_reference_resolves_only_from_a_variable_the_worker_allows( + monkeypatch: pytest.MonkeyPatch, resolvable: Collection[str], resolves: bool +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + built = _build_tool(received, _WorkerEnvRefResolver(resolvable)) + + assert isinstance(built, HostedMCPTool) + expected = SENTINEL if resolves else ref + assert _as_dict(built.tool_config)["authorization"] == expected + assert _as_dict(built.tool_config)["headers"]["X-Token"] == expected + + +def test_star_anywhere_in_the_resolvable_names_resolves_every_name( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + built = _build_tool(received, _WorkerEnvRefResolver([OTHER_ENV_NAME, "*"])) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == SENTINEL + assert _as_dict(built.tool_config)["headers"]["X-Token"] == SENTINEL + + +def test_a_glob_in_the_resolvable_names_matches_no_name( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + built = _build_tool(received, _WorkerEnvRefResolver(["TEMPORAL_TEST_*"])) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == ref + assert _as_dict(built.tool_config)["headers"]["X-Token"] == ref + + +def test_hosted_shell_domain_secret_resolves_for_the_model_call( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _shell_tool(_domain_secret("TOKEN", ref)) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _shell_secrets(built) == [_domain_secret("TOKEN", SENTINEL)] + + +def test_code_interpreter_domain_secret_resolves_for_the_model_call( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool(_domain_secret("TOKEN", ref)) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _code_interpreter_secrets(built) == [_domain_secret("TOKEN", SENTINEL)] + + +def test_code_interpreter_domain_secret_survives_a_second_build( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool(_domain_secret("TOKEN", ref)) + ) + + first = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + second = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _code_interpreter_secrets(first) == [_domain_secret("TOKEN", SENTINEL)] + assert _code_interpreter_secrets(second) == [_domain_secret("TOKEN", SENTINEL)] + assert _secrets_in(received.tool_config["container"]["network_policy"]) == [ + _domain_secret("TOKEN", ref) + ] + + +def test_only_domain_secrets_holding_a_worker_env_ref_are_resolved( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + literal = _domain_secret("PLAIN", "plain-token-value") + _payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool(literal, _domain_secret("TOKEN", ref)) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _code_interpreter_secrets(built) == [ + literal, + _domain_secret("TOKEN", SENTINEL), + ] + + +def test_two_domain_secrets_resolve_to_their_own_secrets( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + _payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool( + _domain_secret("TOKEN", temporal_worker_env_ref(ENV_NAME)), + _domain_secret("OTHER", temporal_worker_env_ref(OTHER_ENV_NAME)), + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _code_interpreter_secrets(built) == [ + _domain_secret("TOKEN", SENTINEL), + _domain_secret("OTHER", OTHER_SENTINEL), + ] + + +def test_shell_domain_secrets_resolve_to_their_own_secrets_in_order( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + literal = _domain_secret("PLAIN", "plain-token-value") + _payload, received = _activity_input_payload_and_tool( + _shell_tool( + literal, + _domain_secret("TOKEN", temporal_worker_env_ref(ENV_NAME)), + _domain_secret("OTHER", temporal_worker_env_ref(OTHER_ENV_NAME)), + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert _shell_secrets(built) == [ + literal, + _domain_secret("TOKEN", SENTINEL), + _domain_secret("OTHER", OTHER_SENTINEL), + ] + + +def test_a_shell_domain_secret_naming_a_denied_variable_is_passed_through( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + denied = temporal_worker_env_ref(OTHER_ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _shell_tool(_domain_secret("TOKEN", denied)) + ) + + built = _build_tool(received, _WorkerEnvRefResolver([ENV_NAME])) + + assert _shell_secrets(built) == [_domain_secret("TOKEN", denied)] + + +def test_a_code_interpreter_domain_secret_naming_a_denied_variable_is_passed_through( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + denied = temporal_worker_env_ref(OTHER_ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _code_interpreter_tool(_domain_secret("TOKEN", denied)) + ) + + built = _build_tool(received, _WorkerEnvRefResolver([ENV_NAME])) + + assert _code_interpreter_secrets(built) == [_domain_secret("TOKEN", denied)] + + +def test_two_worker_env_refs_in_one_mcp_config_resolve_to_their_own_secrets( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool( + temporal_worker_env_ref(ENV_NAME), temporal_worker_env_ref(OTHER_ENV_NAME) + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == SENTINEL + assert _as_dict(built.tool_config)["headers"] == { + "X-Token": OTHER_SENTINEL, + "X-Plain": "not-a-secret", + } + + +def test_a_worker_env_ref_in_a_header_name_is_passed_through( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + ref = temporal_worker_env_ref(ENV_NAME) + tool_config: Any = { + "type": "mcp", + "server_label": "test_server", + "server_url": "https://example.com/mcp", + "headers": {ref: "not-a-secret"}, + } + + payload, received = _activity_input_payload_and_tool( + HostedMCPTool(tool_config=tool_config) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["headers"] == {ref: "not-a-secret"} + assert SENTINEL.encode() not in payload + + +def test_an_mcp_config_without_an_authorization_resolves_its_headers( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + tool_config: Any = { + "type": "mcp", + "server_label": "test_server", + "server_url": "https://example.com/mcp", + "headers": {"X-Token": temporal_worker_env_ref(ENV_NAME)}, + } + + _payload, received = _activity_input_payload_and_tool( + HostedMCPTool(tool_config=tool_config) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert "authorization" not in _as_dict(built.tool_config) + assert _as_dict(built.tool_config)["headers"] == {"X-Token": SENTINEL} + + +def test_an_mcp_config_without_headers_resolves_its_authorization( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + tool_config: Any = { + "type": "mcp", + "server_label": "test_server", + "server_url": "https://example.com/mcp", + "authorization": temporal_worker_env_ref(ENV_NAME), + } + + _payload, received = _activity_input_payload_and_tool( + HostedMCPTool(tool_config=tool_config) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == SENTINEL + assert "headers" not in _as_dict(built.tool_config) + + +def test_local_shell_environment_keeps_its_executor(): + _payload, received = _activity_input_payload_and_tool( + ShellTool(environment={"type": "local"}, executor=lambda _request: "") + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, ShellTool) + assert built.executor is not None + assert _as_dict(built.environment) == {"type": "local"} + + +@pytest.mark.parametrize( + "environment", + [ + {"type": "container_auto"}, + { + "type": "container_auto", + "network_policy": {"type": "disabled"}, + }, + {"type": "container_reference", "container_id": "cntr_abc"}, + ], + ids=["container_auto", "container_auto_disabled_policy", "container_reference"], +) +def test_hosted_shell_environment_gets_no_executor( + environment: ShellToolEnvironment, +): + _payload, received = _activity_input_payload_and_tool( + ShellTool(environment=environment) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, ShellTool) + assert built.executor is None + assert _as_dict(built.environment) == environment + + +def test_code_interpreter_container_id_is_passed_through(): + _payload, received = _activity_input_payload_and_tool( + CodeInterpreterTool( + tool_config={"type": "code_interpreter", "container": "cntr_abc"} + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, CodeInterpreterTool) + assert _as_dict(built.tool_config)["container"] == "cntr_abc" + + +@pytest.mark.parametrize( + "container", + [ + {"type": "auto"}, + {"type": "auto", "network_policy": {"type": "disabled"}}, + ], + ids=["no_policy", "disabled_policy"], +) +def test_code_interpreter_container_without_domain_secrets_is_passed_through( + container: Any, +): + _payload, received = _activity_input_payload_and_tool( + CodeInterpreterTool( + tool_config={"type": "code_interpreter", "container": container} + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, CodeInterpreterTool) + assert _as_dict(built.tool_config)["container"] == container + + +def test_an_unset_environment_variable_resolves_to_an_empty_value( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv(ENV_NAME, raising=False) + ref = temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool(_hosted_mcp_tool(ref, ref)) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == "" + assert _as_dict(built.tool_config)["headers"]["X-Token"] == "" + + +def test_a_policy_with_no_domain_secrets_is_passed_through(): + policy: Any = {"type": "allowlist", "allowed_domains": ["example.com"]} + _payload, received = _activity_input_payload_and_tool( + CodeInterpreterTool( + tool_config={ + "type": "code_interpreter", + "container": {"type": "auto", "network_policy": policy}, + } + ) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, CodeInterpreterTool) + assert _as_dict(built.tool_config)["container"] == { + "type": "auto", + "network_policy": policy, + } + + +def test_plain_values_are_passed_through_unchanged( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + plain = "temporal.worker_env_ref-but-not-quite" + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool(plain, plain) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == plain + assert _as_dict(built.tool_config)["headers"] == { + "X-Token": plain, + "X-Plain": "not-a-secret", + } + + +def test_a_worker_env_ref_inside_a_larger_value_is_substituted_in_place( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + composed = "Bearer " + temporal_worker_env_ref(ENV_NAME) + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool(composed, composed) + ) + + built = _build_tool(received, _RESOLVER_ALLOWING_TEST_NAMES) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == f"Bearer {SENTINEL}" + assert _as_dict(built.tool_config)["headers"]["X-Token"] == f"Bearer {SENTINEL}" + + +def test_one_value_holding_two_refs_resolves_only_the_allowed_name( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + denied = temporal_worker_env_ref(OTHER_ENV_NAME) + composed = f"{temporal_worker_env_ref(ENV_NAME)} {denied}" + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool(composed, composed) + ) + + built = _build_tool(received, _WorkerEnvRefResolver([ENV_NAME])) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == f"{SENTINEL} {denied}" + assert _as_dict(built.tool_config)["headers"]["X-Token"] == f"{SENTINEL} {denied}" + + +def test_a_worker_env_ref_with_no_closing_brace_is_passed_through( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + truncated = "temporal.worker_env_ref:{" + ENV_NAME + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool(truncated, truncated) + ) + + built = _build_tool(received, _WorkerEnvRefResolver(["*"])) + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == truncated + assert _as_dict(built.tool_config)["headers"]["X-Token"] == truncated + + +def test_a_value_packed_with_unclosed_references_does_not_stall_the_resolver(): + opener = "temporal.worker_env_ref:{" + packed = opener * (1024 * 1024 // len(opener)) + _payload, received = _activity_input_payload_and_tool( + _hosted_mcp_tool(packed, "not-a-secret") + ) + + start = time.monotonic() + built = _build_tool(received, _WorkerEnvRefResolver([])) + elapsed = time.monotonic() - start + + assert isinstance(built, HostedMCPTool) + assert _as_dict(built.tool_config)["authorization"] == packed + assert elapsed < 5.0 + + +def test_a_bare_string_is_rejected_as_the_resolvable_variable_names(): + with pytest.raises(TypeError, match="resolvable_worker_env_vars"): + _WorkerEnvRefResolver(ENV_NAME) + + +async def _no_stream_events() -> AsyncIterator[TResponseStreamEvent]: + """Publishing an event here makes the flusher retry for ten minutes against a workflow that does not exist.""" + events: list[TResponseStreamEvent] = [] + for event in events: + yield event + + +class _ToolRecordingModel(Model): + def __init__(self) -> None: + self.tools: list[Tool] = [] + + async def get_response( + self, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + tracing: ModelTracing, + **kwargs: Any, + ) -> ModelResponse: + self.tools = tools + return ModelResponse(output=[], usage=Usage(), response_id=None) + + def stream_response( + self, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + tracing: ModelTracing, + **kwargs: Any, + ) -> AsyncIterator[TResponseStreamEvent]: + self.tools = tools + return _no_stream_events() + + +def _hosted_mcp_config_the_model_received(model: _ToolRecordingModel) -> dict[str, Any]: + assert len(model.tools) == 1 + tool = model.tools[0] + assert isinstance(tool, HostedMCPTool) + return _as_dict(tool.tool_config) + + +async def test_invoke_model_activity_resolves_tool_secrets( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + sent = _hosted_mcp_tool(temporal_worker_env_ref(ENV_NAME), "not-a-secret") + _payload, activity_input = _round_trip_activity_input(sent) + model = _ToolRecordingModel() + + await ActivityEnvironment().run( + ModelActivity( + TestModelProvider(model), resolvable_worker_env_vars=[ENV_NAME] + ).invoke_model_activity, + activity_input, + ) + + assert _hosted_mcp_config_the_model_received(model) == { + **_as_dict(sent.tool_config), + "authorization": SENTINEL, + } + + +async def test_invoke_model_activity_resolves_nothing_by_default( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + sent = _hosted_mcp_tool(temporal_worker_env_ref(ENV_NAME), "not-a-secret") + _payload, activity_input = _round_trip_activity_input(sent) + model = _ToolRecordingModel() + + await ActivityEnvironment().run( + ModelActivity(TestModelProvider(model)).invoke_model_activity, + activity_input, + ) + + assert _hosted_mcp_config_the_model_received(model) == _as_dict(sent.tool_config) + + +async def test_invoke_model_activity_streaming_resolves_tool_secrets( + monkeypatch: pytest.MonkeyPatch, client: Client +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + sent = _hosted_mcp_tool(temporal_worker_env_ref(ENV_NAME), "not-a-secret") + _payload, activity_input = _round_trip_activity_input(sent) + streaming_input: StreamingActivityModelInput = { + **activity_input, + "streaming_topic": "events", + } + model = _ToolRecordingModel() + + await ActivityEnvironment(client).run( + ModelActivity( + TestModelProvider(model), resolvable_worker_env_vars=[ENV_NAME] + ).invoke_model_activity_streaming, + streaming_input, + ) + + assert _hosted_mcp_config_the_model_received(model) == { + **_as_dict(sent.tool_config), + "authorization": SENTINEL, + } + + +@workflow.defn +class WorkerEnvRefAgentWorkflow: + @workflow.run + async def run(self) -> None: + tool_config: Any = { + "type": "mcp", + "server_label": "test_server", + "server_url": "https://example.com/mcp", + "authorization": temporal_worker_env_ref(ENV_NAME), + "headers": {"X-Token": temporal_worker_env_ref(OTHER_ENV_NAME)}, + } + agent = Agent[None]( + name="Worker env ref agent", + tools=[HostedMCPTool(tool_config=tool_config)], + ) + await Runner.run(starting_agent=agent, input="hi") + + +async def test_a_worker_resolves_only_the_variables_its_plugin_names( + monkeypatch: pytest.MonkeyPatch, client: Client +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + model = _ToolRecordingModel() + config = client.config() + config["plugins"] = [ + *config.get("plugins", []), + OpenAIAgentsPlugin( + model_provider=TestModelProvider(model), + resolvable_worker_env_vars=[ENV_NAME], + ), + ] + client = Client(**config) + + async with new_worker(client, WorkerEnvRefAgentWorkflow) as worker: + await client.execute_workflow( + WorkerEnvRefAgentWorkflow.run, + id=f"worker-env-ref-allowlist-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + tool_config = _hosted_mcp_config_the_model_received(model) + assert tool_config["authorization"] == SENTINEL + assert tool_config["headers"] == { + "X-Token": temporal_worker_env_ref(OTHER_ENV_NAME) + } + + +async def test_an_agent_environment_forwards_the_variables_it_names_to_the_worker( + monkeypatch: pytest.MonkeyPatch, client: Client +): + monkeypatch.setenv(ENV_NAME, SENTINEL) + monkeypatch.setenv(OTHER_ENV_NAME, OTHER_SENTINEL) + model = _ToolRecordingModel() + + async with AgentEnvironment( + model=model, resolvable_worker_env_vars=[ENV_NAME] + ) as env: + client = env.applied_on_client(client) + async with new_worker(client, WorkerEnvRefAgentWorkflow) as worker: + await client.execute_workflow( + WorkerEnvRefAgentWorkflow.run, + id=f"agent-environment-env-ref-allowlist-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + tool_config = _hosted_mcp_config_the_model_received(model) + assert tool_config["authorization"] == SENTINEL + assert tool_config["headers"] == { + "X-Token": temporal_worker_env_ref(OTHER_ENV_NAME) + }