Fix in-process Execution API secrets routing in client contexts

airflow-core/docs/core-concepts/overview.rst

@@ -212,9 +212,10 @@ The two processes talk over a socket, and the Supervisor is the only side that e
 task JWT or talks to the *Execution API* — the user's code never sees the token and never touches the
 database.
 
-The same runtime can also run *in-process* (a single Python process, no fork, no sockets, no HTTP) for
-``dag.test()`` and local runs. The diagram below contrasts the two paths and marks where each Python process
-lives:
+The same runtime can also run *in-process* (a single Python process, no fork, no HTTP) for
+``dag.test()`` and local runs. A supervisor socket is still set up, because operators such as
+``PythonVirtualenvOperator`` spawn their own child process that has to reconnect to ask for Connections
+and Variables. The diagram below contrasts the two paths and marks where each Python process lives:
 
 .. image:: ../img/diagram_task_sdk_execution_architecture.png
 

airflow-core/src/airflow/api_fastapi/execution_api/app.py

@@ -44,9 +44,11 @@
     get_sig_validation_args,
     get_signing_args,
 )
+from airflow.process_context import override_process_context
 
 if TYPE_CHECKING:
     import httpx
+    from starlette.types import Receive, Scope, Send
 
 import structlog
 from structlog.contextvars import bind_contextvars
@@ -372,13 +374,28 @@ def _shutdown_loop(
     thread.join(timeout=5)
 
 
+class _RequestScopedServerContextApp:
+    """Wrap an ASGI app so in-process requests behave like server-side API handling."""
+
+    def __init__(self, app: FastAPI) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        with override_process_context("server"):
+            await self.app(scope, receive, send)
+
+
 @attrs.define()
 class InProcessExecutionAPI:
     """
     A helper class to make it possible to run the ExecutionAPI "in-process".
 
     The sync version of this makes use of a2wsgi which runs the async loop in a separate thread. This is
     needed so that we can use the sync httpx client
+
+    Requests are always dispatched in a server process context, with no way to opt out: this app *is*
+    the server side of the Execution API, and in a single process its route handlers would otherwise
+    read the caller's ``SUPERVISOR_COMMS`` and re-enter the API through the Task SDK.
     """
 
     _app: FastAPI | None = None
@@ -433,7 +450,8 @@ def transport(self) -> httpx.WSGITransport:
         thread = threading.Thread(target=loop.run_forever, name="InProcessExecutionAPI-loop", daemon=True)
         thread.start()
 
-        middleware = ASGIMiddleware(self.app, loop=loop)
+        app = self.app
+        middleware = ASGIMiddleware(cast("Any", _RequestScopedServerContextApp(app)), loop=loop)
 
         # https://github.com/abersheeran/a2wsgi/discussions/64
         async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
@@ -443,7 +461,7 @@ async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
 
         # Wait for lifespan startup to complete so callers see a ready app and so the finalizer can
         # safely aclose() a context whose __aenter__ has actually run.
-        asyncio.run_coroutine_threadsafe(start_lifespan(cm, self.app), loop).result()
+        asyncio.run_coroutine_threadsafe(start_lifespan(cm, app), loop).result()
 
         transport = httpx.WSGITransport(app=middleware)  # type: ignore[arg-type]
 
@@ -460,4 +478,4 @@ async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
     def atransport(self) -> httpx.ASGITransport:
         import httpx
 
-        return httpx.ASGITransport(app=self.app)
+        return httpx.ASGITransport(app=_RequestScopedServerContextApp(self.app))

airflow-core/src/airflow/models/connection.py

@@ -20,7 +20,6 @@
 import json
 import logging
 import re
-import sys
 import warnings
 from contextlib import suppress
 from json import JSONDecodeError
@@ -50,6 +49,7 @@ class AirflowSecretsBackendAccessDenied(PermissionError):  # type: ignore[no-red
         """Compat stub — never raised by task-sdk <1.2.2."""
 
 
+from airflow.process_context import should_use_task_sdk_api_path
 from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -475,7 +475,7 @@ def get_connection_from_secrets(cls, conn_id: str, team_name: str | None = None)
 
         # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             from airflow.sdk import Connection as TaskSDKConnection
             from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
 
@@ -566,7 +566,7 @@ def to_dict(self, *, prune_empty: bool = False, validate: bool = True) -> dict[s
 
     @classmethod
     def from_json(cls, value, conn_id=None) -> Connection:
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             from airflow.sdk import Connection as TaskSDKConnection
 
             warnings.warn(

airflow-core/src/airflow/models/variable.py

@@ -20,7 +20,6 @@
 import contextlib
 import json
 import logging
-import sys
 import warnings
 from typing import TYPE_CHECKING, Any
 
@@ -47,6 +46,7 @@ class AirflowSecretsBackendAccessDenied(PermissionError):  # type: ignore[no-red
         """Compat stub — never raised by task-sdk <1.2.2."""
 
 
+from airflow.process_context import should_use_task_sdk_api_path
 from airflow.secrets.metastore import MetastoreBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
@@ -166,7 +166,7 @@ def get(
 
         # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             warnings.warn(
                 "Using Variable.get from `airflow.models` is deprecated."
                 "Please use `get` on Variable from sdk(`airflow.sdk.Variable`) instead",
@@ -226,7 +226,7 @@ def set(
 
         # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             warnings.warn(
                 "Using Variable.set from `airflow.models` is deprecated."
                 "Please use `set` on Variable from sdk(`airflow.sdk.Variable`) instead",
@@ -314,7 +314,7 @@ def update(
 
         # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             warnings.warn(
                 "Using Variable.update from `airflow.models` is deprecated."
                 "Please use `set` on Variable from sdk(`airflow.sdk.Variable`) instead as it is an upsert.",
@@ -380,7 +380,7 @@ def delete(key: str, team_name: str | None = None, session: Session | None = Non
 
         # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
         # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+        if should_use_task_sdk_api_path():
             warnings.warn(
                 "Using Variable.delete from `airflow.models` is deprecated."
                 "Please use `delete` on Variable from sdk(`airflow.sdk.Variable`) instead",

airflow-core/src/airflow/process_context.py

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import sys
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Literal
+
+__all__ = [
+    "override_process_context",
+    "should_use_task_sdk_api_path",
+]
+
+_PROCESS_CONTEXT_OVERRIDE: ContextVar[str | None] = ContextVar(
+    "_AIRFLOW_PROCESS_CONTEXT_OVERRIDE",
+    default=None,
+)
+
+
+@contextmanager
+def override_process_context(context: Literal["server", "client"]) -> Generator[None, None, None]:
+    """Temporarily override the current process context for the active execution flow."""
+    token = _PROCESS_CONTEXT_OVERRIDE.set(context)
+    try:
+        yield
+    finally:
+        _PROCESS_CONTEXT_OVERRIDE.reset(token)
+
+
+def should_use_task_sdk_api_path() -> bool:
+    """Return True when execution-context helpers should route through Task SDK APIs."""
+    # Only the ContextVar, never the ``_AIRFLOW_PROCESS_CONTEXT`` env var: that env var is
+    # process-wide and inherited by children (``action_cli`` sets it around the whole
+    # ``airflow dags test`` body, and PythonVirtualenvOperator passes it to the venv child), so
+    # letting it win here would send worker-side code straight to the metastore. ``SUPERVISOR_COMMS``
+    # keeps precedence over it, matching ``ensure_secrets_backend_loaded()`` in the Task SDK.
+    if _PROCESS_CONTEXT_OVERRIDE.get() == "server":
+        return False
+
+    task_runner_module = sys.modules.get("airflow.sdk.execution_time.task_runner")
+    return bool(getattr(task_runner_module, "SUPERVISOR_COMMS", None))

airflow-core/tests/unit/api_fastapi/execution_api/test_app.py

@@ -163,6 +163,39 @@ def test_in_process_execution_api_runs_without_jwt_secret():
     assert response.status_code == 200
 
 
+def test_in_process_execution_api_does_not_reenter_task_sdk(session):
+    """A request served in-process must not route back out through the Task SDK.
+
+    Under ``dag.test()`` the API server and the task share one process, so ``SUPERVISOR_COMMS`` is
+    visible to the route handler too. Without a request-scoped server context the handler reads it,
+    decides it is client-side, and issues another Execution API request -- looping until the caller
+    is killed. No env var is set here on purpose: ``dag.test()`` from a script sets none, so the
+    ContextVar is the only thing preventing re-entry.
+    """
+    from airflow.models.variable import Variable
+
+    # Patched below, not by decorator: this setup must run before ``SUPERVISOR_COMMS`` is visible.
+    Variable.set(key="inproc_key", value="VALUE", session=session)
+    session.commit()
+
+    api = InProcessExecutionAPI()
+    with (
+        mock.patch.dict(
+            "sys.modules",
+            {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])},
+        ),
+        mock.patch(
+            "airflow.sdk.Variable.get",
+            side_effect=AssertionError("in-process Execution API re-entered the Task SDK path"),
+        ),
+        httpx.Client(transport=api.transport) as client,
+    ):
+        response = client.get("http://localhost/variables/inproc_key")
+
+    assert response.status_code == 200
+    assert response.json() == {"key": "inproc_key", "value": "VALUE"}
+
+
 def test_in_process_execution_api_transport_lifecycle():
     """The background loop + thread lifecycle is tied to the ``.transport``, not the factory instance.
 

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connections.py

@@ -17,12 +17,14 @@
 
 from __future__ import annotations
 
+import sys
 from unittest import mock
 
 import pytest
 from fastapi import FastAPI, HTTPException, status
 
 from airflow.models.connection import Connection
+from airflow.process_context import override_process_context
 
 pytestmark = pytest.mark.db_test
 
@@ -105,6 +107,38 @@ def test_connection_get_from_env_var(self, client, session):
             "extra": '{"headers": "header"}',
         }
 
+    @mock.patch.dict(
+        "os.environ",
+        {
+            "AIRFLOW_CONN_TEST_CONN_SERVER": '{"uri": "http://root:admin@localhost:8080/https?headers=header"}',
+        },
+    )
+    @mock.patch.dict(
+        sys.modules,
+        {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])},
+    )
+    @mock.patch(
+        "airflow.sdk.Connection.get",
+        side_effect=AssertionError(
+            "Execution API should not route through Task SDK Connection.get in server context"
+        ),
+    )
+    def test_connection_get_uses_server_path_when_supervisor_comms_exists(self, mock_sdk_get, client):
+        with override_process_context("server"):
+            response = client.get("/execution/connections/test_conn_server")
+
+        assert response.status_code == 200
+        assert response.json() == {
+            "conn_id": "test_conn_server",
+            "conn_type": "http",
+            "host": "localhost",
+            "login": "root",
+            "password": "admin",
+            "schema": "https",
+            "port": 8080,
+            "extra": '{"headers": "header"}',
+        }
+
     def test_connection_get_not_found(self, client):
         response = client.get("/execution/connections/non_existent_test_conn")