From 00d9d2e4b62f00959bfa72972d64f7600bc8daac Mon Sep 17 00:00:00 2001 From: Simon Schampijer <3714785+erikos@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:34:04 +0200 Subject: [PATCH 01/11] feat: expose request headers to pipeline wrapper run methods (opt-in) OpenAI-compatible run_chat_completion / run_response wrappers can now receive the incoming request headers by declaring an optional 'headers' parameter. The OpenAI router forwards headers to the wrapper only when the method declares the parameter (or **kwargs), so existing wrappers are unaffected. Removes the need for custom middleware + a ContextVar to bridge request headers to a wrapper for use cases like per-request identity/token forwarding. Assisted-by: Claude:claude-opus-4-8 --- src/hayhooks/server/routers/openai.py | 39 ++++++++++++++++--- .../server/utils/base_pipeline_wrapper.py | 4 ++ 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/hayhooks/server/routers/openai.py b/src/hayhooks/server/routers/openai.py index c1f8292f..86362244 100644 --- a/src/hayhooks/server/routers/openai.py +++ b/src/hayhooks/server/routers/openai.py @@ -1,3 +1,4 @@ +import inspect import time from collections.abc import AsyncGenerator, Generator from dataclasses import dataclass @@ -104,6 +105,25 @@ def _select_execution_mode(wrapper: BasePipelineWrapper, dispatch: _OpenAIDispat raise HTTPException(status_code=501, detail=dispatch.not_implemented_detail) +def _method_accepts_kwarg(method: Any, name: str) -> bool: + """True if `method` declares keyword argument `name` (explicitly or via **kwargs).""" + try: + params = inspect.signature(method).parameters + except (TypeError, ValueError): + return False + return name in params or any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + + +def _build_call_kwargs( + wrapper: BasePipelineWrapper, method_name: str, base_kwargs: dict[str, Any], body: dict[str, Any], headers: dict[str, str] | None +) -> dict[str, Any]: + """Assemble the wrapper-call kwargs, forwarding `headers` only if the wrapper method opts in.""" + call_kwargs = {**base_kwargs, "body": body} + if headers is not None and _method_accepts_kwarg(getattr(wrapper, method_name), "headers"): + call_kwargs["headers"] = headers + return call_kwargs + + async def _invoke_pipeline_method( wrapper: BasePipelineWrapper, *, mode: str, method_name: str, model: str, call_kwargs: dict[str, Any] ) -> Any: @@ -140,6 +160,7 @@ async def _run_pipeline_method( model: str, kwargs: dict[str, Any], body: dict[str, Any], + headers: dict[str, str] | None = None, ) -> str | Generator | AsyncGenerator: """Shared dispatch logic for chat completions and responses endpoints.""" stream_requested = bool(body.get("stream", False)) @@ -155,8 +176,9 @@ async def _run_pipeline_method( try: wrapper = _resolve_pipeline_wrapper(model) mode, method_name = _select_execution_mode(wrapper, dispatch) + call_kwargs = _build_call_kwargs(wrapper, method_name, kwargs, body, headers) result = await _invoke_pipeline_method( - wrapper, mode=mode, method_name=method_name, model=model, call_kwargs={**kwargs, "body": body} + wrapper, mode=mode, method_name=method_name, model=model, call_kwargs=call_kwargs ) normalized_result = await _normalize_result(result, stream_requested=stream_requested) except BaseException: @@ -175,22 +197,27 @@ async def _run_pipeline_method( wrapper = _resolve_pipeline_wrapper(model) mode, method_name = _select_execution_mode(wrapper, dispatch) span.set_tag("hayhooks.openai.execution_mode", mode) + call_kwargs = _build_call_kwargs(wrapper, method_name, kwargs, body, headers) result = await _invoke_pipeline_method( - wrapper, mode=mode, method_name=method_name, model=model, call_kwargs={**kwargs, "body": body} + wrapper, mode=mode, method_name=method_name, model=model, call_kwargs=call_kwargs ) return await _normalize_result(result, stream_requested=stream_requested) async def _run_completion( - model: str, messages: list[dict[str, Any]], body: dict[str, Any] + model: str, messages: list[dict[str, Any]], body: dict[str, Any], headers: dict[str, str] | None = None ) -> str | Generator | AsyncGenerator: - return await _run_pipeline_method(_CHAT_COMPLETION_DISPATCH, model=model, kwargs={"messages": messages}, body=body) + return await _run_pipeline_method( + _CHAT_COMPLETION_DISPATCH, model=model, kwargs={"messages": messages}, body=body, headers=headers + ) async def _run_response( - model: str, input_items: list[dict[str, Any]], body: dict[str, Any] + model: str, input_items: list[dict[str, Any]], body: dict[str, Any], headers: dict[str, str] | None = None ) -> str | Generator | AsyncGenerator: - return await _run_pipeline_method(_RESPONSE_DISPATCH, model=model, kwargs={"input_items": input_items}, body=body) + return await _run_pipeline_method( + _RESPONSE_DISPATCH, model=model, kwargs={"input_items": input_items}, body=body, headers=headers + ) def _find_file_upload_wrapper() -> BasePipelineWrapper | None: diff --git a/src/hayhooks/server/utils/base_pipeline_wrapper.py b/src/hayhooks/server/utils/base_pipeline_wrapper.py index 2e0ae372..24738a84 100644 --- a/src/hayhooks/server/utils/base_pipeline_wrapper.py +++ b/src/hayhooks/server/utils/base_pipeline_wrapper.py @@ -77,6 +77,10 @@ def run_chat_completion(self, model: str, messages: list[dict], body: dict) -> s model: The `name` of the deployed Haystack pipeline to run messages: The history of messages as OpenAI-compatible list of dicts body: Additional parameters and configuration options + + Optionally, override this method with an extra ``headers: dict[str, str]`` parameter to + receive the incoming request headers (e.g. ``Authorization``). Headers are only passed when + the overriding method declares the parameter, so existing implementations are unaffected. """ msg = "run_chat_completion not implemented" raise NotImplementedError(msg) From f121e74e7f8d0a3ed9c04032e867c74ec8a01c08 Mon Sep 17 00:00:00 2001 From: Simon Schampijer <3714785+erikos@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:47:14 +0200 Subject: [PATCH 02/11] test: cover the opt-in header forwarding to wrappers Adds tests for the behaviour introduced in the previous commit. Two wrapper fixtures under tests/test_files/files/ deploy through the normal path: one declares `headers` and echoes the Authorization header back, the other keeps the existing (model, messages, body) signature and must be called unchanged. The end-to-end checks are skipped when the installed fastapi-openai-compat predates header forwarding, since the wrapper cannot receive headers the router never passes on. They activate on their own once a version that forwards them is installed. `_method_accepts_kwarg` and `_build_call_kwargs` get unit tests, covering explicit parameters, **kwargs, wrappers that do not opt in, callables without an introspectable signature, and headers=None from non-HTTP entry points. Also wraps the `_build_call_kwargs` signature to stay inside the 120-column limit, and documents the opt-in on run_response and on both async variants, which the previous commit only did for run_chat_completion. Assisted-by: Claude:claude-opus-5 --- src/hayhooks/server/routers/openai.py | 6 +- .../server/utils/base_pipeline_wrapper.py | 10 ++ .../chat_with_headers/pipeline_wrapper.py | 14 ++ .../chat_without_headers/pipeline_wrapper.py | 14 ++ tests/test_it_openai_headers.py | 146 ++++++++++++++++++ 5 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 tests/test_files/files/chat_with_headers/pipeline_wrapper.py create mode 100644 tests/test_files/files/chat_without_headers/pipeline_wrapper.py create mode 100644 tests/test_it_openai_headers.py diff --git a/src/hayhooks/server/routers/openai.py b/src/hayhooks/server/routers/openai.py index 86362244..882ab807 100644 --- a/src/hayhooks/server/routers/openai.py +++ b/src/hayhooks/server/routers/openai.py @@ -115,7 +115,11 @@ def _method_accepts_kwarg(method: Any, name: str) -> bool: def _build_call_kwargs( - wrapper: BasePipelineWrapper, method_name: str, base_kwargs: dict[str, Any], body: dict[str, Any], headers: dict[str, str] | None + wrapper: BasePipelineWrapper, + method_name: str, + base_kwargs: dict[str, Any], + body: dict[str, Any], + headers: dict[str, str] | None, ) -> dict[str, Any]: """Assemble the wrapper-call kwargs, forwarding `headers` only if the wrapper method opts in.""" call_kwargs = {**base_kwargs, "body": body} diff --git a/src/hayhooks/server/utils/base_pipeline_wrapper.py b/src/hayhooks/server/utils/base_pipeline_wrapper.py index 24738a84..90ee02d4 100644 --- a/src/hayhooks/server/utils/base_pipeline_wrapper.py +++ b/src/hayhooks/server/utils/base_pipeline_wrapper.py @@ -88,6 +88,9 @@ def run_chat_completion(self, model: str, messages: list[dict], body: dict) -> s async def run_chat_completion_async(self, model: str, messages: list[dict], body: dict) -> str | AsyncGenerator: """ Asynchronous version of run_chat_completion. + + Accepts an optional ``headers: dict[str, str]`` parameter on the same opt-in basis as + run_chat_completion. """ msg = "run_chat_completion_async not implemented" raise NotImplementedError(msg) @@ -103,6 +106,10 @@ def run_response(self, model: str, input_items: list[dict], body: dict) -> str | model: The `name` of the deployed Haystack pipeline to run input_items: Normalized input items in OpenAI Responses API format body: Additional parameters and configuration options (e.g. temperature, tools, instructions) + + Optionally, override this method with an extra ``headers: dict[str, str]`` parameter to + receive the incoming request headers (e.g. ``Authorization``). Headers are only passed when + the overriding method declares the parameter, so existing implementations are unaffected. """ msg = "run_response not implemented" raise NotImplementedError(msg) @@ -110,6 +117,9 @@ def run_response(self, model: str, input_items: list[dict], body: dict) -> str | async def run_response_async(self, model: str, input_items: list[dict], body: dict) -> str | AsyncGenerator: """ Asynchronous version of run_response. + + Accepts an optional ``headers: dict[str, str]`` parameter on the same opt-in basis as + run_response. """ msg = "run_response_async not implemented" raise NotImplementedError(msg) diff --git a/tests/test_files/files/chat_with_headers/pipeline_wrapper.py b/tests/test_files/files/chat_with_headers/pipeline_wrapper.py new file mode 100644 index 00000000..60a3f527 --- /dev/null +++ b/tests/test_files/files/chat_with_headers/pipeline_wrapper.py @@ -0,0 +1,14 @@ +from haystack import Pipeline + +from hayhooks import BasePipelineWrapper + + +class PipelineWrapper(BasePipelineWrapper): + """Declares the optional `headers` parameter, so Hayhooks forwards the request headers.""" + + def setup(self) -> None: + self.pipeline = Pipeline() + + def run_chat_completion(self, model: str, messages: list[dict], body: dict, headers: dict[str, str]) -> str: + # NOTE: This is used in tests, please don't change it + return f"authorization={headers.get('authorization', 'missing')}" diff --git a/tests/test_files/files/chat_without_headers/pipeline_wrapper.py b/tests/test_files/files/chat_without_headers/pipeline_wrapper.py new file mode 100644 index 00000000..dd9da8b5 --- /dev/null +++ b/tests/test_files/files/chat_without_headers/pipeline_wrapper.py @@ -0,0 +1,14 @@ +from haystack import Pipeline + +from hayhooks import BasePipelineWrapper + + +class PipelineWrapper(BasePipelineWrapper): + """Keeps the pre-existing signature, so Hayhooks must not forward request headers.""" + + def setup(self) -> None: + self.pipeline = Pipeline() + + def run_chat_completion(self, model: str, messages: list[dict], body: dict) -> str: + # NOTE: This is used in tests, please don't change it + return "no headers parameter declared" diff --git a/tests/test_it_openai_headers.py b/tests/test_it_openai_headers.py new file mode 100644 index 00000000..e61d5360 --- /dev/null +++ b/tests/test_it_openai_headers.py @@ -0,0 +1,146 @@ +"""Request headers reach pipeline wrappers that declare a `headers` parameter, and only those.""" + +import shutil +from pathlib import Path + +import pytest +from fastapi_openai_compat import ChatRequest + +from hayhooks.server.pipelines import registry +from hayhooks.server.routers.openai import _build_call_kwargs, _method_accepts_kwarg +from hayhooks.server.utils.base_pipeline_wrapper import BasePipelineWrapper +from hayhooks.settings import settings + +# Headers can only reach a wrapper if the installed fastapi-openai-compat forwards them to the +# run_completion callback. Skip the end-to-end checks on versions that predate that. +try: + from fastapi_openai_compat._shared import callable_accepts_kwarg as _compat_opt_in + + _COMPAT_FORWARDS_HEADERS = callable(_compat_opt_in) +except ImportError: # pragma: no cover - depends on the installed version + _COMPAT_FORWARDS_HEADERS = False + +requires_header_forwarding = pytest.mark.skipif( + not _COMPAT_FORWARDS_HEADERS, + reason="installed fastapi-openai-compat does not forward request headers to run_completion", +) + +TEST_FILES_DIR_WITH_HEADERS = Path(__file__).parent / "test_files/files/chat_with_headers" +PIPELINE_FILES_WITH_HEADERS = { + "pipeline_wrapper.py": (TEST_FILES_DIR_WITH_HEADERS / "pipeline_wrapper.py").read_text(), +} + +TEST_FILES_DIR_WITHOUT_HEADERS = Path(__file__).parent / "test_files/files/chat_without_headers" +PIPELINE_FILES_WITHOUT_HEADERS = { + "pipeline_wrapper.py": (TEST_FILES_DIR_WITHOUT_HEADERS / "pipeline_wrapper.py").read_text(), +} + + +@pytest.fixture(autouse=True) +def clear_registry(): + registry.clear() + if Path(settings.pipelines_dir).exists(): + shutil.rmtree(settings.pipelines_dir) + yield + + +def _chat(client, model: str, headers: dict[str, str] | None = None): + request = ChatRequest(stream=False, model=model, messages=[{"role": "user", "content": "who am I?"}]) + return client.post("/chat/completions", json=request.model_dump(), headers=headers) + + +@requires_header_forwarding +def test_headers_forwarded_when_wrapper_declares_them(client, deploy_files): + assert deploy_files(client, "with_headers", PIPELINE_FILES_WITH_HEADERS).status_code == 200 + + response = _chat(client, "with_headers", {"Authorization": "Bearer alice-token"}) + + assert response.status_code == 200 + content = response.json()["choices"][0]["message"]["content"] + assert content == "authorization=Bearer alice-token" + + +def test_wrapper_without_headers_parameter_is_unaffected(client, deploy_files): + """The pre-existing (model, messages, body) signature must keep working untouched.""" + assert deploy_files(client, "without_headers", PIPELINE_FILES_WITHOUT_HEADERS).status_code == 200 + + response = _chat(client, "without_headers", {"Authorization": "Bearer bob-token"}) + + assert response.status_code == 200 + content = response.json()["choices"][0]["message"]["content"] + assert content == "no headers parameter declared" + + +@requires_header_forwarding +def test_wrapper_declaring_headers_without_request_headers(client, deploy_files): + """A request always carries some headers, so the wrapper still gets a dict, just without ours.""" + assert deploy_files(client, "with_headers", PIPELINE_FILES_WITH_HEADERS).status_code == 200 + + response = _chat(client, "with_headers") + + assert response.status_code == 200 + assert response.json()["choices"][0]["message"]["content"] == "authorization=missing" + + +# --- the opt-in helpers ---------------------------------------------------------------------- + + +class _WithHeaders(BasePipelineWrapper): + def setup(self) -> None: ... + + def run_chat_completion(self, model: str, messages: list[dict], body: dict, headers: dict[str, str]) -> str: + return "ok" + + +class _WithKwargs(BasePipelineWrapper): + def setup(self) -> None: ... + + def run_chat_completion(self, model: str, messages: list[dict], body: dict, **kwargs) -> str: + return "ok" + + +class _WithoutHeaders(BasePipelineWrapper): + def setup(self) -> None: ... + + def run_chat_completion(self, model: str, messages: list[dict], body: dict) -> str: + return "ok" + + +def test_method_accepts_kwarg_detects_explicit_parameter(): + assert _method_accepts_kwarg(_WithHeaders().run_chat_completion, "headers") is True + + +def test_method_accepts_kwarg_detects_var_keyword(): + assert _method_accepts_kwarg(_WithKwargs().run_chat_completion, "headers") is True + + +def test_method_accepts_kwarg_rejects_missing_parameter(): + assert _method_accepts_kwarg(_WithoutHeaders().run_chat_completion, "headers") is False + + +def test_method_accepts_kwarg_on_unintrospectable_callable(): + """Callables without a signature must not raise; they simply do not opt in.""" + assert _method_accepts_kwarg(print, "headers") is False + + +def test_build_call_kwargs_includes_headers_for_opted_in_wrapper(): + call_kwargs = _build_call_kwargs( + _WithHeaders(), "run_chat_completion", {"messages": []}, {"stream": False}, {"authorization": "Bearer x"} + ) + + assert call_kwargs == {"messages": [], "body": {"stream": False}, "headers": {"authorization": "Bearer x"}} + + +def test_build_call_kwargs_omits_headers_for_other_wrappers(): + call_kwargs = _build_call_kwargs( + _WithoutHeaders(), "run_chat_completion", {"messages": []}, {"stream": False}, {"authorization": "Bearer x"} + ) + + assert call_kwargs == {"messages": [], "body": {"stream": False}} + + +def test_build_call_kwargs_omits_headers_when_none_available(): + """Non-HTTP entry points pass headers=None, which must not reach the wrapper.""" + call_kwargs = _build_call_kwargs(_WithHeaders(), "run_chat_completion", {"messages": []}, {}, None) + + assert call_kwargs == {"messages": [], "body": {}} From 1221b2623e27d832c9538659ac2bb1a683344cef Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Tue, 18 Aug 2026 13:00:49 +0200 Subject: [PATCH 03/11] fix(streaming): make pipeline task cancellation configurable (#262) * fix(streaming): make pipeline task cancellation configurable * refactor(streaming): simplify shielded task cleanup * docs(streaming): explain pipeline task shielding --- docs/concepts/pipeline-wrapper.md | 19 +++++ src/hayhooks/server/pipelines/streaming.py | 57 +++++++++---- tests/test_it_streaming_disconnect.py | 97 ++++++++++++++++++++++ tests/test_streaming.py | 79 ++++++++++++++++++ 4 files changed, 237 insertions(+), 15 deletions(-) create mode 100644 tests/test_it_streaming_disconnect.py diff --git a/docs/concepts/pipeline-wrapper.md b/docs/concepts/pipeline-wrapper.md index c21cd4ce..cf2d4b88 100644 --- a/docs/concepts/pipeline-wrapper.md +++ b/docs/concepts/pipeline-wrapper.md @@ -185,6 +185,25 @@ async def run_api_async(self, query: str) -> AsyncGenerator: ) ``` +##### Client disconnects + +By default, closing the response stream or cancelling its consumer also cancels the pipeline task. Cancellation still +propagates to the request handler so HTTP cleanup can finish normally. + +Set `shield_pipeline_task=True` when the request should stop but the pipeline must continue running: + +```python +return async_streaming_generator( + pipeline=self.pipeline, + pipeline_run_args={"prompt": {"query": query}}, + shield_pipeline_task=True, +) +``` + +Shielding only changes cleanup of the pipeline task: it does not suppress cancellation in the request handler, and no +more chunks can be delivered to a disconnected client. Hayhooks keeps the detached task alive until it finishes and +logs any eventual exception. + When a generator is detected, Hayhooks automatically wraps it in a FastAPI `StreamingResponse` using the `text/plain` media type. The behaviour is identical for both `run_api()` and `run_api_async()`—the only difference is whether the underlying generator is sync or async. | Method | What you return | Response media type | Notes | diff --git a/src/hayhooks/server/pipelines/streaming.py b/src/hayhooks/server/pipelines/streaming.py index fdb6ccb5..42eb1476 100644 --- a/src/hayhooks/server/pipelines/streaming.py +++ b/src/hayhooks/server/pipelines/streaming.py @@ -34,6 +34,7 @@ _ASYNC_STREAMING_LOOP: contextvars.ContextVar[asyncio.AbstractEventLoop] = contextvars.ContextVar( "_hayhooks_async_streaming_loop" ) +_SHIELDED_PIPELINE_TASKS: set[asyncio.Task[Any]] = set() # Streaming callbacks are module-level so Haystack can serialize snapshot inputs. # The active queues live in ContextVars and are set around pipeline execution, so @@ -789,6 +790,25 @@ def _check_pipeline_task_exception(pipeline_task: asyncio.Task) -> None: raise exception +def _handle_shielded_pipeline_task_done(pipeline_task: asyncio.Task) -> None: + """Keep shielded tasks alive until completion and retrieve detached exceptions.""" + _SHIELDED_PIPELINE_TASKS.discard(pipeline_task) + with contextlib.suppress(asyncio.CancelledError): + if error := pipeline_task.exception(): + log.opt(exception=error).error("Error in shielded pipeline task") + + +def _detach_pipeline_task(pipeline_task: asyncio.Task) -> None: + """Keep an unfinished pipeline task alive after its stream closes.""" + if pipeline_task.done(): + with contextlib.suppress(asyncio.CancelledError): + pipeline_task.exception() + return + + _SHIELDED_PIPELINE_TASKS.add(pipeline_task) + pipeline_task.add_done_callback(_handle_shielded_pipeline_task_done) + + async def _stream_chunks_from_queue( queue: asyncio.Queue[StreamingChunk], pipeline_task: asyncio.Task, @@ -826,9 +846,6 @@ async def _stream_chunks_from_queue( if queue.empty(): _check_pipeline_task_exception(pipeline_task) continue - except asyncio.CancelledError: - log.warning("Async streaming generator was cancelled") - break except Exception as e: log.opt(exception=True).error("Unexpected error in async streaming generator: {}", e) raise @@ -837,22 +854,29 @@ async def _stream_chunks_from_queue( _check_pipeline_task_exception(pipeline_task) -async def _cleanup_pipeline_async(pipeline_task: asyncio.Task) -> None: +async def _cleanup_pipeline_async(pipeline_task: asyncio.Task, *, shield_pipeline_task: bool = False) -> None: """ Cleans up the pipeline task if it's still running. Args: pipeline_task: The task to clean up + shield_pipeline_task: Keep the task running after the stream closes instead of cancelling it """ - if not pipeline_task.done(): - pipeline_task.cancel() - try: - await asyncio.wait_for(pipeline_task, timeout=1.0) - except (asyncio.TimeoutError, asyncio.CancelledError): - pass - except Exception as e: - # Don't re-raise - this runs in finally block, so we don't want to mask original errors - log.opt(exception=True).warning("Error during pipeline task cleanup: {}", e) + if shield_pipeline_task: + _detach_pipeline_task(pipeline_task) + return + + if pipeline_task.done(): + return + + pipeline_task.cancel() + try: + await asyncio.wait_for(pipeline_task, timeout=1.0) + except (asyncio.TimeoutError, asyncio.CancelledError): + pass + except Exception as e: + # Don't re-raise - this runs in finally block, so we don't want to mask original errors + log.opt(exception=True).warning("Error during pipeline task cleanup: {}", e) def async_streaming_generator( # noqa: PLR0913, C901 @@ -867,6 +891,7 @@ def async_streaming_generator( # noqa: PLR0913, C901 include_outputs_from: set[str] | None = None, allow_sync_streaming_callbacks: bool = False, external_event_queue: asyncio.Queue[StreamingChunk | PipelineEvent | str | dict[str, Any]] | None = None, + shield_pipeline_task: bool = False, ) -> AsyncGenerator[StreamingChunk | PipelineEvent | str | dict[str, Any], None]: """ Creates an async generator that yields streaming chunks from a pipeline or agent execution. @@ -897,6 +922,8 @@ def async_streaming_generator( # noqa: PLR0913, C901 external_event_queue: Optional external asyncio queue to merge with internal events. Events from this queue will be yielded alongside streaming chunks from the pipeline. Supports StreamingChunk, PipelineEvent, str, or custom dict events. + shield_pipeline_task: Keep the pipeline task running if the stream closes or its consumer is cancelled. + Defaults to False, preserving the existing behavior of cancelling the task. Yields: StreamingChunk: Individual chunks from the streaming execution @@ -966,7 +993,7 @@ async def generator() -> AsyncGenerator[StreamingChunk | PipelineEvent | str | d yield result yield chunk - await pipeline_task + await (asyncio.shield(pipeline_task) if shield_pipeline_task else pipeline_task) final_chunk = _process_pipeline_end(pipeline_task.result(), on_pipeline_end) if final_chunk: yield final_chunk @@ -975,6 +1002,6 @@ async def generator() -> AsyncGenerator[StreamingChunk | PipelineEvent | str | d log.opt(exception=True).error("Unexpected error in async streaming generator: {}", e) raise finally: - await _cleanup_pipeline_async(pipeline_task) + await _cleanup_pipeline_async(pipeline_task, shield_pipeline_task=shield_pipeline_task) return generator() diff --git a/tests/test_it_streaming_disconnect.py b/tests/test_it_streaming_disconnect.py new file mode 100644 index 00000000..647c671d --- /dev/null +++ b/tests/test_it_streaming_disconnect.py @@ -0,0 +1,97 @@ +import asyncio +import contextlib +from typing import Any + +import pytest +from fastapi import FastAPI +from haystack import component +from haystack.dataclasses import StreamingChunk + +from hayhooks.server.pipelines.streaming import _SHIELDED_PIPELINE_TASKS, async_streaming_generator +from hayhooks.server.utils.haystack_compat import AsyncPipeline +from hayhooks.server.utils.streaming_response_utils import _streaming_response_from_async_gen + + +@component +class _BlockingStreamingComponent: + def __init__(self) -> None: + self.release = asyncio.Event() + self.completed = asyncio.Event() + + @component.output_types(result=str) + def run(self, streaming_callback: Any | None = None) -> dict[str, str]: + raise AssertionError("The async pipeline must call run_async") + + @component.output_types(result=str) + async def run_async(self, streaming_callback: Any | None = None) -> dict[str, str]: + await streaming_callback(StreamingChunk(content="first", index=0)) + await self.release.wait() + self.completed.set() + return {"result": "done"} + + +async def _disconnect_after_first_chunk(app: FastAPI) -> None: + incoming = asyncio.Queue() + first_chunk = asyncio.Event() + await incoming.put({"type": "http.request", "body": b"", "more_body": False}) + + async def receive(): + return await incoming.get() + + async def send(message): + if message["type"] == "http.response.body" and message.get("body"): + first_chunk.set() + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": "/stream", + "raw_path": b"/stream", + "query_string": b"", + "headers": [], + "client": ("test", 1), + "server": ("test", 80), + "root_path": "", + } + request_task = asyncio.create_task(app(scope, receive, send)) + + try: + await asyncio.wait_for(first_chunk.wait(), timeout=1.0) + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request_task, timeout=1.0) + finally: + if not request_task.done(): + request_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await request_task + + +@pytest.mark.integration +@pytest.mark.parametrize("shield_pipeline_task", [False, True], ids=["cancel", "shield"]) +async def test_http_disconnect_pipeline_task(shield_pipeline_task): + component = _BlockingStreamingComponent() + pipeline = AsyncPipeline() + pipeline.add_component("blocking", component) + app = FastAPI() + + @app.get("/stream") + async def stream(): + generator = async_streaming_generator(pipeline, shield_pipeline_task=shield_pipeline_task) + return _streaming_response_from_async_gen(generator) + + tasks_before = set(_SHIELDED_PIPELINE_TASKS) + await _disconnect_after_first_chunk(app) + detached_tasks = _SHIELDED_PIPELINE_TASKS - tasks_before + + assert bool(detached_tasks) is shield_pipeline_task + assert all(not task.done() for task in detached_tasks) + + component.release.set() + await asyncio.wait_for(component.completed.wait(), timeout=1.0) + if detached_tasks: + await asyncio.wait_for(asyncio.gather(*detached_tasks), timeout=1.0) + await asyncio.sleep(0) + assert detached_tasks.isdisjoint(_SHIELDED_PIPELINE_TASKS) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 0e419f3e..d91cdac1 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,3 +1,4 @@ +import asyncio import contextvars import os import time @@ -178,6 +179,28 @@ async def mock_run_async(messages=None, streaming_callback=None, **kwargs): return _factory +@pytest.fixture +def blocking_mock_agent(mocker): + release = asyncio.Event() + cancelled = asyncio.Event() + completed = asyncio.Event() + mock_agent = mocker.Mock(spec=Agent) + assert isinstance(mock_agent, Agent) + + async def mock_run_async(messages=None, streaming_callback=None, **kwargs): + await streaming_callback(StreamingChunk(content="First chunk", index=0)) + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + completed.set() + return {"messages": ["Done"]} + + mock_agent.run_async = mocker.AsyncMock(side_effect=mock_run_async) + return mock_agent, release, cancelled, completed + + def test_streaming_generator_with_sync_only_generator(pipeline_with_sync_only_generator): generator = streaming_generator(pipeline_with_sync_only_generator, pipeline_run_args={}) @@ -415,6 +438,62 @@ async def consume_with_cancel(): assert len(chunks) >= 1 +@pytest.mark.parametrize("shield_pipeline_task", [False, True], ids=["cancel", "shield"]) +async def test_async_streaming_generator_close_pipeline_task(blocking_mock_agent, shield_pipeline_task): + mock_agent, release, cancelled, completed = blocking_mock_agent + gen = async_streaming_generator( + mock_agent, + pipeline_run_args={"messages": []}, + shield_pipeline_task=shield_pipeline_task, + ) + + assert (await anext(gen)).content == "First chunk" + await gen.aclose() + + if shield_pipeline_task: + assert not cancelled.is_set() + release.set() + await asyncio.wait_for(completed.wait(), timeout=1.0) + else: + assert cancelled.is_set() + assert not completed.is_set() + + mock_agent.run_async.assert_awaited_once() + + +@pytest.mark.parametrize("shield_pipeline_task", [False, True], ids=["cancel", "shield"]) +async def test_async_streaming_consumer_cancellation_pipeline_task(blocking_mock_agent, shield_pipeline_task): + mock_agent, release, cancelled, completed = blocking_mock_agent + gen = async_streaming_generator( + mock_agent, + pipeline_run_args={"messages": []}, + shield_pipeline_task=shield_pipeline_task, + ) + first_chunk_received = asyncio.Event() + + async def consume(): + assert (await anext(gen)).content == "First chunk" + first_chunk_received.set() + await anext(gen) + + consumer_task = asyncio.create_task(consume()) + await first_chunk_received.wait() + consumer_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await consumer_task + + if shield_pipeline_task: + assert not cancelled.is_set() + release.set() + await asyncio.wait_for(completed.wait(), timeout=1.0) + else: + assert cancelled.is_set() + assert not completed.is_set() + + mock_agent.run_async.assert_awaited_once() + + # ContextVar used to verify context propagation into the streaming thread _test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar("_test_context_var") From 279c1e790e319789beddcbc4e1373be898621ac9 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Tue, 18 Aug 2026 15:33:31 +0200 Subject: [PATCH 04/11] fix(ci): support core metadata 2.5 publishing (#263) --- .github/workflows/pypi.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index ad055038..182584e9 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -23,4 +23,4 @@ jobs: run: hatch build - name: Publish on PyPi - uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 From 8f3e705077a278e0771fc9f3bbeee2ecdd498bd3 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 10:44:39 +0200 Subject: [PATCH 05/11] feat: define typed request header opt-in for wrappers --- src/hayhooks/server/pipelines/models.py | 4 ++ .../server/utils/base_pipeline_wrapper.py | 10 +++++ src/hayhooks/server/utils/module_loader.py | 5 ++- src/hayhooks/server/utils/request_headers.py | 42 +++++++++++++++++++ 4 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 src/hayhooks/server/utils/request_headers.py diff --git a/src/hayhooks/server/pipelines/models.py b/src/hayhooks/server/pipelines/models.py index aec68137..1b2c03cc 100644 --- a/src/hayhooks/server/pipelines/models.py +++ b/src/hayhooks/server/pipelines/models.py @@ -7,6 +7,7 @@ from pydantic import BaseModel, Field, create_model from hayhooks.server.exceptions import PipelineWrapperError +from hayhooks.server.utils.request_headers import accepts_request_headers from hayhooks.server.utils.yaml_utils import InputResolution, OutputResolution @@ -70,10 +71,13 @@ def create_request_model_from_callable(func: Callable, model_name: str, docstrin """ params = inspect.signature(func).parameters + inject_headers = accepts_request_headers(func) param_docs = {p.arg_name: p.description for p in docstring.params} fields: dict[str, Any] = {} for name, param in params.items(): + if name == "headers" and inject_headers: + continue default_value = ... if param.default == param.empty else param.default description = param_docs.get(name) or f"Parameter '{name}'" field_info = Field(default=default_value, description=description) diff --git a/src/hayhooks/server/utils/base_pipeline_wrapper.py b/src/hayhooks/server/utils/base_pipeline_wrapper.py index 2e0ae372..8134b362 100644 --- a/src/hayhooks/server/utils/base_pipeline_wrapper.py +++ b/src/hayhooks/server/utils/base_pipeline_wrapper.py @@ -4,6 +4,16 @@ class BasePipelineWrapper(ABC): + """ + Base class for deployed pipelines. + + The run_api, run_chat_completion and run_response methods, including their async + variants, can declare ``headers: Headers | None = None`` to receive HTTP request + headers. Import ``Headers`` from ``starlette.datastructures`` at runtime. Header + lookup is case-insensitive; A2A, MCP and direct calls use the default of None. + Ordinary dict parameters and **kwargs do not opt in to header injection. + """ + # Class attribute to skip MCP listing of the pipeline # If True, the pipeline will not be listed as an MCP tool # Even if it has a description and a request model diff --git a/src/hayhooks/server/utils/module_loader.py b/src/hayhooks/server/utils/module_loader.py index 9883ca18..140d3efc 100644 --- a/src/hayhooks/server/utils/module_loader.py +++ b/src/hayhooks/server/utils/module_loader.py @@ -15,6 +15,7 @@ from hayhooks.server.exceptions import PipelineModuleLoadError, PipelineWrapperError from hayhooks.server.logger import log from hayhooks.server.utils.base_pipeline_wrapper import BasePipelineWrapper +from hayhooks.server.utils.request_headers import accepts_request_headers from hayhooks.settings import settings @@ -241,7 +242,7 @@ def _raise_load_error(self, original_error: Exception) -> NoReturn: def _set_method_implementation_flags(pipeline_wrapper: BasePipelineWrapper) -> None: """ - Set implementation flags for all supported run methods. + Set implementation flags and validate request-header declarations on supported run methods. Args: pipeline_wrapper: The wrapper instance to annotate with flags. @@ -258,6 +259,8 @@ def _set_method_implementation_flags(pipeline_wrapper: BasePipelineWrapper) -> N for attr_name, method_name in methods_to_check: is_implemented = _is_method_overridden(pipeline_wrapper, method_name) + if is_implemented and method_name != "run_file_upload": + accepts_request_headers(getattr(pipeline_wrapper, method_name)) setattr(pipeline_wrapper, attr_name, is_implemented) log.debug("pipeline_wrapper.{}: {}", attr_name, is_implemented) diff --git a/src/hayhooks/server/utils/request_headers.py b/src/hayhooks/server/utils/request_headers.py new file mode 100644 index 00000000..2588b3a0 --- /dev/null +++ b/src/hayhooks/server/utils/request_headers.py @@ -0,0 +1,42 @@ +import inspect +from collections.abc import Callable +from types import SimpleNamespace, UnionType +from typing import Any, Union, get_args, get_origin, get_type_hints + +from starlette.datastructures import Headers + +from hayhooks.server.exceptions import PipelineWrapperError + + +def accepts_request_headers(method: Callable[..., Any]) -> bool: + """Recognize and validate the explicit ``headers: Headers | None = None`` opt-in.""" + try: + parameter = inspect.signature(method).parameters.get("headers") + except (TypeError, ValueError): + return False + if parameter is None or parameter.annotation is inspect.Parameter.empty: + return False + + # Resolve only this annotation; unrelated forward references need not be importable at runtime. + try: + annotation = get_type_hints( + SimpleNamespace(__annotations__={"headers": parameter.annotation}), + globalns=getattr(inspect.unwrap(method), "__globals__", {}), + )["headers"] + except NameError: + # An unresolved application type does not opt in to HTTP metadata. + return False + if annotation is not Headers and not ( + get_origin(annotation) in (Union, UnionType) and Headers in get_args(annotation) + ): + return False + + if ( + annotation != Headers | None + or parameter.default is not None + or parameter.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) + ): + name = getattr(method, "__name__", type(method).__name__) + msg = f"{name}: request headers must be declared as a keyword parameter 'headers: Headers | None = None'" + raise PipelineWrapperError(msg) + return True From a6d6f82cbe8bf1a9f610d7d1720a1efeaa99c4c4 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 10:44:39 +0200 Subject: [PATCH 06/11] feat: forward request headers across pipeline endpoints --- src/hayhooks/server/routers/openai.py | 35 ++- src/hayhooks/server/utils/deploy_utils.py | 36 ++- src/hayhooks/server/utils/mcp_utils.py | 7 +- tests/test_request_headers.py | 276 ++++++++++++++++++++++ 4 files changed, 330 insertions(+), 24 deletions(-) create mode 100644 tests/test_request_headers.py diff --git a/src/hayhooks/server/routers/openai.py b/src/hayhooks/server/routers/openai.py index 8577f3fd..fc4ca229 100644 --- a/src/hayhooks/server/routers/openai.py +++ b/src/hayhooks/server/routers/openai.py @@ -14,6 +14,7 @@ create_responses_router, ) from haystack.dataclasses import StreamingChunk +from starlette.datastructures import Headers from hayhooks.server.logger import log from hayhooks.server.pipelines.registry import registry @@ -27,6 +28,7 @@ trace_sync_stream, ) from hayhooks.server.utils.base_pipeline_wrapper import BasePipelineWrapper +from hayhooks.server.utils.request_headers import accepts_request_headers @dataclass(frozen=True) @@ -105,14 +107,21 @@ def _select_execution_mode(wrapper: BasePipelineWrapper, dispatch: _OpenAIDispat async def _invoke_pipeline_method( - wrapper: BasePipelineWrapper, *, mode: str, method_name: str, model: str, call_kwargs: dict[str, Any] + wrapper: BasePipelineWrapper, + *, + mode: str, + method_name: str, + call_kwargs: dict[str, Any], + headers: dict[str, str] | None, ) -> Any: """Invoke the resolved pipeline method in either async or threadpool-sync mode.""" method = getattr(wrapper, method_name) - log.debug("Using {} ({}) for model: {}", method_name, mode, model) + if accepts_request_headers(method): + call_kwargs = {**call_kwargs, "headers": Headers(headers) if headers is not None else None} + log.debug("Using {} ({}) for model: {}", method_name, mode, call_kwargs["model"]) if mode == "async": - return await method(model=model, **call_kwargs) - return await run_in_threadpool(method, model=model, **call_kwargs) + return await method(**call_kwargs) + return await run_in_threadpool(method, **call_kwargs) def _wrap_string_as_streaming(text: str) -> Generator[StreamingChunk, None, None]: @@ -140,9 +149,11 @@ async def _run_pipeline_method( model: str, kwargs: dict[str, Any], body: dict[str, Any], + headers: dict[str, str] | None, ) -> str | Generator | AsyncGenerator: """Shared dispatch logic for chat completions and responses endpoints.""" stream_requested = bool(body.get("stream", False)) + call_kwargs = {"model": model, **kwargs, "body": body} trace_tags = build_trace_tags( { "hayhooks.transport": "openai", @@ -156,7 +167,7 @@ async def _run_pipeline_method( wrapper = _resolve_pipeline_wrapper(model) mode, method_name = _select_execution_mode(wrapper, dispatch) result = await _invoke_pipeline_method( - wrapper, mode=mode, method_name=method_name, model=model, call_kwargs={**kwargs, "body": body} + wrapper, mode=mode, method_name=method_name, call_kwargs=call_kwargs, headers=headers ) normalized_result = await _normalize_result(result, stream_requested=stream_requested) except BaseException: @@ -184,21 +195,25 @@ async def _run_pipeline_method( mode, method_name = _select_execution_mode(wrapper, dispatch) span.set_tag("hayhooks.openai.execution_mode", mode) result = await _invoke_pipeline_method( - wrapper, mode=mode, method_name=method_name, model=model, call_kwargs={**kwargs, "body": body} + wrapper, mode=mode, method_name=method_name, call_kwargs=call_kwargs, headers=headers ) return await _normalize_result(result, stream_requested=stream_requested) async def _run_completion( - model: str, messages: list[dict[str, Any]], body: dict[str, Any] + model: str, messages: list[dict[str, Any]], body: dict[str, Any], headers: dict[str, str] | None = None ) -> str | Generator | AsyncGenerator: - return await _run_pipeline_method(_CHAT_COMPLETION_DISPATCH, model=model, kwargs={"messages": messages}, body=body) + return await _run_pipeline_method( + _CHAT_COMPLETION_DISPATCH, model=model, kwargs={"messages": messages}, body=body, headers=headers + ) async def _run_response( - model: str, input_items: list[dict[str, Any]], body: dict[str, Any] + model: str, input_items: list[dict[str, Any]], body: dict[str, Any], headers: dict[str, str] | None = None ) -> str | Generator | AsyncGenerator: - return await _run_pipeline_method(_RESPONSE_DISPATCH, model=model, kwargs={"input_items": input_items}, body=body) + return await _run_pipeline_method( + _RESPONSE_DISPATCH, model=model, kwargs={"input_items": input_items}, body=body, headers=headers + ) def _find_file_upload_wrapper() -> BasePipelineWrapper | None: diff --git a/src/hayhooks/server/utils/deploy_utils.py b/src/hayhooks/server/utils/deploy_utils.py index fc90918e..8115b918 100644 --- a/src/hayhooks/server/utils/deploy_utils.py +++ b/src/hayhooks/server/utils/deploy_utils.py @@ -12,11 +12,12 @@ from typing import Any, cast import docstring_parser -from fastapi import FastAPI, Form, HTTPException +from fastapi import FastAPI, Form, HTTPException, Request from fastapi.concurrency import run_in_threadpool from fastapi.responses import Response, StreamingResponse from fastapi.routing import APIRoute from pydantic import BaseModel +from starlette.datastructures import Headers from hayhooks.server.exceptions import PipelineAlreadyExistsError, PipelineFilesError from hayhooks.server.logger import log, log_elapsed @@ -46,6 +47,7 @@ load_pipeline_module, unload_pipeline_modules, ) +from hayhooks.server.utils.request_headers import accepts_request_headers from hayhooks.server.utils.streaming_response_utils import _streaming_response_from_result from hayhooks.server.utils.yaml_pipeline_wrapper import YAMLPipelineWrapper from hayhooks.settings import DeployConcurrencyPolicy, settings @@ -233,10 +235,20 @@ async def wrapper(*args, **kwargs): async def _execute_pipeline_run( pipeline_wrapper: BasePipelineWrapper, payload: dict[str, Any], + *, + headers: Headers | None = None, ) -> Any: + method = ( + pipeline_wrapper.run_api_async if pipeline_wrapper._is_run_api_async_implemented else pipeline_wrapper.run_api + ) + if accepts_request_headers(method): + if "headers" in payload: + msg = "Request headers cannot be supplied as pipeline arguments" + raise ValueError(msg) + payload = {**payload, "headers": headers} if pipeline_wrapper._is_run_api_async_implemented: - return await pipeline_wrapper.run_api_async(**payload) - return await run_in_threadpool(pipeline_wrapper.run_api, **payload) + return await method(**payload) + return await run_in_threadpool(method, **payload) _SENSITIVE_KEY_PATTERNS = { @@ -330,18 +342,19 @@ async def _execute_pipeline_run_with_tracing( pipeline_wrapper: BasePipelineWrapper, payload: dict[str, Any], *, + headers: Headers, trace_tags: dict[str, Any], is_streaming_response: bool, ) -> Any: if is_streaming_response: try: - return await _execute_pipeline_run(pipeline_wrapper, payload) + return await _execute_pipeline_run(pipeline_wrapper, payload, headers=headers) except BaseException: with trace_operation(SPAN_PIPELINE_RUN, tags=trace_tags): raise with trace_operation(SPAN_PIPELINE_RUN, tags=trace_tags): - return await _execute_pipeline_run(pipeline_wrapper, payload) + return await _execute_pipeline_run(pipeline_wrapper, payload, headers=headers) def _trace_streaming_run_result(result: Any, trace_tags: dict[str, Any]) -> Any: @@ -405,7 +418,7 @@ def create_run_endpoint_handler( ) is_streaming_response = get_response_class_from_callable(run_method) is StreamingResponse - async def _handle_request(run_req: BaseModel) -> Response | BaseModel: + async def _handle_request(run_req: BaseModel, request: Request) -> Response | BaseModel: payload = run_req.model_dump() trace_tags = _build_run_trace_tags(pipeline_name, payload) @@ -414,6 +427,7 @@ async def _handle_request(run_req: BaseModel) -> Response | BaseModel: result = await _execute_pipeline_run_with_tracing( pipeline_wrapper, payload, + headers=request.headers, trace_tags=trace_tags, is_streaming_response=is_streaming_response, ) @@ -449,13 +463,17 @@ async def _handle_request(run_req: BaseModel) -> Response | BaseModel: @handle_pipeline_exceptions() async def run_endpoint_with_files( + request: Request, run_req: request_model = Form(..., media_type="multipart/form-data"), # ty: ignore[invalid-type-form] # noqa: B008 ) -> response_model: # ty: ignore[invalid-type-form] - return await _handle_request(run_req) + return await _handle_request(run_req, request) @handle_pipeline_exceptions() - async def run_endpoint_without_files(run_req: request_model) -> response_model: # ty: ignore[invalid-type-form] - return await _handle_request(run_req) + async def run_endpoint_without_files( + request: Request, + run_req: request_model, # ty: ignore[invalid-type-form] + ) -> response_model: # ty: ignore[invalid-type-form] + return await _handle_request(run_req, request) return run_endpoint_with_files if requires_files else run_endpoint_without_files diff --git a/src/hayhooks/server/utils/mcp_utils.py b/src/hayhooks/server/utils/mcp_utils.py index 382e8441..6b505850 100644 --- a/src/hayhooks/server/utils/mcp_utils.py +++ b/src/hayhooks/server/utils/mcp_utils.py @@ -5,7 +5,6 @@ from enum import Enum from typing import Any -from fastapi.concurrency import run_in_threadpool from haystack.lazy_imports import LazyImport from starlette.applications import Starlette from starlette.responses import JSONResponse @@ -26,6 +25,7 @@ ) from hayhooks.server.utils.base_pipeline_wrapper import BasePipelineWrapper from hayhooks.server.utils.deploy_utils import ( + _execute_pipeline_run, deploy_pipeline_files_async, deploy_pipelines, # noqa: F401 (re-exported; historically lived in this module) undeploy_pipeline_async, @@ -140,10 +140,7 @@ async def run_pipeline_as_tool(name: str, arguments: dict[str, Any]) -> list["Te msg = f"Pipeline '{name}' not found" raise ValueError(msg) - if pipeline_wrapper._is_run_api_async_implemented: - result = await pipeline_wrapper.run_api_async(**arguments) - else: - result = await run_in_threadpool(pipeline_wrapper.run_api, **arguments) + result = await _execute_pipeline_run(pipeline_wrapper, arguments) log.trace("Pipeline '{}' returned result: {}", name, result) diff --git a/tests/test_request_headers.py b/tests/test_request_headers.py new file mode 100644 index 00000000..ac6c3ba1 --- /dev/null +++ b/tests/test_request_headers.py @@ -0,0 +1,276 @@ +import asyncio +import importlib.util +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from httpx import ASGITransport, AsyncClient + +from hayhooks.server.pipelines.registry import registry +from hayhooks.server.routers.deploy import router as deploy_router +from hayhooks.server.routers.openai import _run_completion, _run_response +from hayhooks.server.routers.openai import router as openai_router +from hayhooks.server.utils.a2a_utils import _run_chat_completion +from hayhooks.server.utils.mcp_utils import list_pipelines_as_tools, run_pipeline_as_tool +from hayhooks.settings import settings + +METHODS = ["run_api", "run_chat_completion", "run_response"] +METHODS += [method + "_async" for method in METHODS] +BASE_SOURCE = """\ +import asyncio +import time +from collections.abc import Generator, AsyncGenerator +from fastapi import UploadFile +from starlette.datastructures import Headers +from hayhooks import BasePipelineWrapper + +class PipelineWrapper(BasePipelineWrapper): + def setup(self): + pass +""" + + +@pytest.fixture +def headers_client(): + registry.clear() + app = FastAPI() + app.include_router(deploy_router) + app.include_router(openai_router) + with TestClient(app) as client: + yield client + registry.clear() + + +def deploy(client, source): + response = client.post( + "/deploy_files", + json={ + "name": "headers_test", + "files": {"pipeline_wrapper.py": source}, + "save_files": False, + }, + ) + assert response.status_code == 200, response.text + + +def wrapper_source(method, stream=False, files=False): + is_async = method.endswith("_async") + arguments = ( + "query: str" + if method.startswith("run_api") + else ( + "model: str, messages: list[dict], body: dict" + if "chat" in method + else "model: str, input_items: list[dict], body: dict" + ) + ) + if files: + arguments += ", files: list[UploadFile]" + prefix = "async " if is_async else "" + delay = "await asyncio.sleep(0.01)" if is_async else "time.sleep(0.01)" + return_type = ("AsyncGenerator" if is_async else "Generator") if stream else "str" + result = ( + f" {prefix}def chunks():\n {delay}\n yield token\n return chunks()\n" + if stream + else " return token\n" + ) + return BASE_SOURCE + ( + f" {prefix}def {method}(self, {arguments}, *, headers: Headers | None = None) -> {return_type}:\n" + f" {delay}\n" + " token = headers.get('Authorization', 'missing') if headers is not None else 'no-context'\n" + result + ) + + +def request_for(method, stream=False, prefix="/v1"): + if method.startswith("run_api"): + return "/headers_test/run", {"query": "hi", "headers": {"authorization": "body-spoof"}} + body = {"model": "headers_test", "stream": stream, "headers": {"authorization": "body-spoof"}} + if "chat" in method: + return prefix + "/chat/completions", {**body, "messages": [{"role": "user", "content": "hi"}]} + return prefix + "/responses", {**body, "input": "hi"} + + +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("stream", [False, True]) +async def test_headers_are_request_local_through_execution_and_streaming(headers_client, method, stream): + deploy(headers_client, wrapper_source(method, stream)) + url, body = request_for(method, stream) + tokens = ["Bearer alice-test", "Bearer bob-test", None] + async with AsyncClient(transport=ASGITransport(app=headers_client.app), base_url="http://test") as client: + responses = await asyncio.gather( + *[ + client.post(url, json=body, headers={"Authorization": token} if token is not None else {}) + for token in tokens + ] + ) + for token, response in zip(tokens, responses, strict=True): + assert response.status_code == 200, response.text + assert (token or "missing") in response.text + assert "body-spoof" not in response.text + assert all((other or "missing") not in response.text for other in tokens if other != token) + + +@pytest.mark.parametrize("method", [m for m in METHODS if not m.startswith("run_api")]) +def test_openai_aliases_forward_headers(headers_client, method): + deploy(headers_client, wrapper_source(method)) + url, body = request_for(method, prefix="") + response = headers_client.post(url, json=body, headers={"Authorization": "alias-token"}) + assert response.status_code == 200, response.text + assert "alias-token" in response.text + + +@pytest.mark.parametrize("method", ["run_api", "run_api_async"]) +def test_multipart_headers_are_injected_outside_the_form(headers_client, method): + deploy(headers_client, wrapper_source(method, files=True)) + response = headers_client.post( + "/headers_test/run", + data={"query": "hi", "headers": "body-spoof"}, + files={"files": ("test.txt", b"test", "text/plain")}, + headers={"Authorization": "upload-token"}, + ) + assert response.status_code == 200, response.text + assert response.json() == {"result": "upload-token"} + + +@pytest.mark.parametrize("method", ["run_api", "run_api_async"]) +@pytest.mark.skipif(importlib.util.find_spec("mcp") is None, reason="MCP is not installed") +@pytest.mark.mcp +async def test_injected_headers_stay_out_of_schemas_and_mcp_arguments(headers_client, method): + deploy(headers_client, wrapper_source(method)) + schema = headers_client.get("/openapi.json").json()["components"]["schemas"]["headers_testRunRequest"] + assert set(schema["properties"]) == {"query"} + tools = await list_pipelines_as_tools() + assert set(tools[0].inputSchema["properties"]) == {"query"} + result = await run_pipeline_as_tool("headers_test", {"query": "hi"}) + assert result[0].text == "no-context" + with pytest.raises(ValueError, match="cannot be supplied as pipeline arguments"): + await run_pipeline_as_tool("headers_test", {"query": "hi", "headers": {"authorization": "tool-spoof"}}) + + +@pytest.mark.parametrize("method", [m for m in METHODS if not m.startswith("run_api")]) +async def test_non_http_openai_calls_use_the_default(headers_client, method): + deploy(headers_client, wrapper_source(method)) + if "chat" in method: + assert await _run_completion("headers_test", [], {}) == "no-context" + else: + assert await _run_response("headers_test", [], {}) == "no-context" + + +@pytest.mark.parametrize("method", ["run_chat_completion", "run_chat_completion_async"]) +@pytest.mark.skipif(importlib.util.find_spec("a2a") is None, reason="A2A is not installed") +@pytest.mark.a2a +async def test_a2a_calls_use_the_default(headers_client, method): + deploy(headers_client, wrapper_source(method)) + assert await _run_chat_completion("headers_test", SimpleNamespace(message=None, current_task=None)) == "no-context" + + +@pytest.mark.parametrize( + "extra,assertion", + [ + ("", "True"), + (", **kwargs", "kwargs == {}"), + (", *headers", "headers == ()"), + (", headers: dict[str, str] | None = None", "headers is None"), + (", headers: 'NotImportedAtRuntime' = None", "headers is None"), + (", headers: dict[str, Headers] = None", "headers is None"), + ], +) +@pytest.mark.parametrize("method", [m for m in METHODS if not m.startswith("run_api")]) +def test_legacy_openai_signatures_do_not_receive_headers(headers_client, extra, assertion, method): + args = ( + "model: str, messages: list[dict], body: dict" + if "chat" in method + else "model: str, input_items: list[dict], body: dict" + ) + prefix = "async " if method.endswith("_async") else "" + deploy( + headers_client, + BASE_SOURCE + + ( + f" {prefix}def {method}(self, {args}{extra}) -> str:\n" + f" assert {assertion}\n return 'legacy'\n" + ), + ) + url, body = request_for(method) + response = headers_client.post(url, json=body, headers={"Authorization": "must-not-inject"}) + assert response.status_code == 200, response.text + assert "legacy" in response.text + + +@pytest.mark.parametrize("method", ["run_api", "run_api_async"]) +@pytest.mark.skipif(importlib.util.find_spec("mcp") is None, reason="MCP is not installed") +@pytest.mark.mcp +async def test_regular_headers_body_field_keeps_its_schema_and_value(headers_client, method): + prefix = "async " if method.endswith("_async") else "" + deploy( + headers_client, + BASE_SOURCE + + ( + f" {prefix}def {method}(self, headers: dict[str, str]) -> str:\n" + " return headers['authorization']\n" + ), + ) + response = headers_client.post( + "/headers_test/run", + json={"headers": {"authorization": "body-value"}}, + headers={"Authorization": "transport-value"}, + ) + assert response.json() == {"result": "body-value"} + tools = await list_pipelines_as_tools() + assert tools[0].inputSchema["required"] == ["headers"] + result = await run_pipeline_as_tool("headers_test", {"headers": {"authorization": "tool-value"}}) + assert result[0].text == "tool-value" + + +@pytest.mark.parametrize( + "declaration", + [ + "headers: Headers", + "headers: Headers | None", + "headers: Headers = None", + "headers: Headers | str | None = None", + "headers: Headers | None = None, /", + "*headers: Headers | None", + "**headers: Headers | None", + ], +) +def test_invalid_opt_in_is_rejected_at_deployment(headers_client, declaration): + source = BASE_SOURCE + f" def run_api(self, {declaration}) -> str:\n return 'unused'\n" + response = headers_client.post( + "/deploy_files", + json={ + "name": "invalid_headers", + "files": {"pipeline_wrapper.py": source}, + "save_files": False, + }, + ) + assert response.status_code == 422, response.text + assert "headers: Headers | None = None" in response.json()["detail"] + assert registry.get("invalid_headers") is None + + +def test_postponed_header_annotation_is_resolved_without_resolving_unrelated_types(headers_client): + source = "from __future__ import annotations\n" + wrapper_source("run_chat_completion") + source = source.replace("messages: list[dict]", "messages: NotImportedAtRuntime") + deploy(headers_client, source) + response = headers_client.post( + "/chat/completions", + json={"model": "headers_test", "messages": []}, + headers={"Authorization": "resolved-token"}, + ) + assert response.status_code == 200, response.text + assert "resolved-token" in response.text + + +def test_injected_headers_are_not_payload_logs_or_trace_tags(headers_client, recording_tracer, caplog, monkeypatch): + monkeypatch.setattr(settings, "dashboard_trace_include_payload_values", True) + deploy(headers_client, wrapper_source("run_api")) + response = headers_client.post( + "/headers_test/run", + json={"query": "hi"}, + headers={"Authorization": "secret-transport-token"}, + ) + assert response.status_code == 200 + assert "secret-transport-token" not in caplog.text + assert "secret-transport-token" not in repr([span.tags for span in recording_tracer.spans]) From 7ba73012bee3db718ec1387c9d38612efaf5e979 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 10:44:39 +0200 Subject: [PATCH 07/11] docs: describe request header access and require compat support --- docs/concepts/pipeline-wrapper.md | 38 +++++++++++++++++++++++++++ docs/features/openai-compatibility.md | 3 +++ pyproject.toml | 2 +- 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/docs/concepts/pipeline-wrapper.md b/docs/concepts/pipeline-wrapper.md index cf2d4b88..96a734f6 100644 --- a/docs/concepts/pipeline-wrapper.md +++ b/docs/concepts/pipeline-wrapper.md @@ -112,6 +112,44 @@ def setup(self) -> None: !!! tip "Consider YAML-only deployment" If your pipeline is simple and doesn't need custom logic, consider using [YAML Pipeline Deployment](yaml-pipeline-deployment.md) instead, which doesn't require a wrapper at all. +### Request headers + +To access HTTP request headers, explicitly declare `headers: Headers | None = None` +on any of `run_api`, `run_chat_completion`, `run_response`, or their async variants: + +```python +from fastapi import HTTPException +from starlette.datastructures import Headers + +def run_api(self, query: str, *, headers: Headers | None = None) -> str: + authorization = headers.get("authorization") if headers is not None else None + if not authorization: + raise HTTPException(status_code=401, detail="Authorization is required") + # Validate the credential or forward it only to the intended trusted service. + return self.answer(query, authorization=authorization) +``` + +Import `Headers` at runtime, including when using postponed annotations. The +parameter must accept keyword arguments and have the annotation `Headers | None` +and default `None`; invalid declarations fail during deployment. Headers are +request-local and available throughout streaming. Lookups are case-insensitive. +The OpenAI-compatible endpoints receive a header dictionary from +`fastapi-openai-compat`, so repeated values for a header are not preserved there. + +The injected parameter is excluded from the `/run` JSON or multipart body and from +OpenAPI and MCP input schemas. A body field cannot override it. Existing application +parameters named `headers` with other types (such as `dict[str, str]`) keep their +normal behavior. Declaring only `**kwargs` does not enable injection. + +A2A, MCP and direct calls do not automatically supply HTTP headers, so the default +is `None`. MCP tool arguments cannot populate the injected parameter. Wrappers +that require authentication must handle missing credentials explicitly. Receiving +headers does not itself authenticate the caller. + +Do not store caller credentials on the shared wrapper instance, echo them in +responses, or include them in logs or traces. Hayhooks keeps injected headers out +of its automatic `/run` payload logs and trace tags. + ### run_api() The `run_api()` method is called for each API request to the `{pipeline_name}/run` endpoint. diff --git a/docs/features/openai-compatibility.md b/docs/features/openai-compatibility.md index 99137c31..8e1b2bd9 100644 --- a/docs/features/openai-compatibility.md +++ b/docs/features/openai-compatibility.md @@ -16,6 +16,9 @@ Hayhooks supports two OpenAI API surfaces: Both APIs are available simultaneously. A pipeline wrapper can implement one or both. +To access caller credentials or other HTTP headers, declare the optional typed +[`headers` parameter](../concepts/pipeline-wrapper.md#request-headers) on your wrapper method. + ## Key Features - **Automatic Endpoint Generation**: OpenAI-compatible endpoints are created automatically diff --git a/pyproject.toml b/pyproject.toml index 7a8e0dad..377e3af6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "pydantic-settings", "python-dotenv", "docstring-parser", - "fastapi-openai-compat>=1.2.0", + "fastapi-openai-compat>=1.3.0", ] [project.optional-dependencies] From c4e848cda81d7c75dfd9f2f5f890615c3f629198 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 11:02:55 +0200 Subject: [PATCH 08/11] fix: resolve postponed wrapper annotations and strengthen header tests --- docs/concepts/pipeline-wrapper.md | 79 +++++++++++++------------ docs/features/a2a-support.md | 1 + docs/features/mcp-support.md | 6 +- docs/features/openai-compatibility.md | 1 + src/hayhooks/server/pipelines/models.py | 54 +---------------- tests/test_deploy_utils.py | 57 ++++-------------- tests/test_request_headers.py | 70 ++++++++++++++-------- 7 files changed, 108 insertions(+), 160 deletions(-) diff --git a/docs/concepts/pipeline-wrapper.md b/docs/concepts/pipeline-wrapper.md index 96a734f6..60c4375b 100644 --- a/docs/concepts/pipeline-wrapper.md +++ b/docs/concepts/pipeline-wrapper.md @@ -30,6 +30,44 @@ class PipelineWrapper(BasePipelineWrapper): return result["llm"]["replies"][0] ``` +## Request headers + +To access HTTP request headers, explicitly declare `headers: Headers | None = None` +on any of `run_api`, `run_chat_completion`, `run_response`, or their async variants: + +```python +from fastapi import HTTPException +from starlette.datastructures import Headers + +def run_api(self, query: str, *, headers: Headers | None = None) -> str: + authorization = headers.get("authorization") if headers is not None else None + if not authorization: + raise HTTPException(status_code=401, detail="Authorization is required") + # Application code: validate the credential or forward it only to your trusted service. + return self.answer(query, authorization=authorization) +``` + +Import `Headers` at runtime, including when using postponed annotations. The +parameter must accept keyword arguments and have the annotation `Headers | None` +and default `None`; invalid declarations fail during deployment. Headers are +request-local and available throughout streaming. Lookups are case-insensitive. +The OpenAI-compatible endpoints receive a header dictionary from +`fastapi-openai-compat`, so repeated values for a header are not preserved there. + +The injected parameter is excluded from the `/run` JSON or multipart body and from +OpenAPI and MCP input schemas. A body field cannot override it. Existing application +parameters named `headers` with other types (such as `dict[str, str]`) keep their +normal behavior. Declaring only `**kwargs` does not enable injection. + +A2A, MCP and direct calls do not automatically supply HTTP headers, so the default +is `None`. MCP tool arguments cannot populate the injected parameter. Wrappers +that require authentication must handle missing credentials explicitly. Receiving +headers does not itself authenticate the caller. + +Do not store caller credentials on the shared wrapper instance, echo them in +responses, or include them in logs or traces. Hayhooks keeps injected headers out +of its automatic `/run` payload logs and trace tags. + ## Required Methods ### setup() @@ -112,44 +150,6 @@ def setup(self) -> None: !!! tip "Consider YAML-only deployment" If your pipeline is simple and doesn't need custom logic, consider using [YAML Pipeline Deployment](yaml-pipeline-deployment.md) instead, which doesn't require a wrapper at all. -### Request headers - -To access HTTP request headers, explicitly declare `headers: Headers | None = None` -on any of `run_api`, `run_chat_completion`, `run_response`, or their async variants: - -```python -from fastapi import HTTPException -from starlette.datastructures import Headers - -def run_api(self, query: str, *, headers: Headers | None = None) -> str: - authorization = headers.get("authorization") if headers is not None else None - if not authorization: - raise HTTPException(status_code=401, detail="Authorization is required") - # Validate the credential or forward it only to the intended trusted service. - return self.answer(query, authorization=authorization) -``` - -Import `Headers` at runtime, including when using postponed annotations. The -parameter must accept keyword arguments and have the annotation `Headers | None` -and default `None`; invalid declarations fail during deployment. Headers are -request-local and available throughout streaming. Lookups are case-insensitive. -The OpenAI-compatible endpoints receive a header dictionary from -`fastapi-openai-compat`, so repeated values for a header are not preserved there. - -The injected parameter is excluded from the `/run` JSON or multipart body and from -OpenAPI and MCP input schemas. A body field cannot override it. Existing application -parameters named `headers` with other types (such as `dict[str, str]`) keep their -normal behavior. Declaring only `**kwargs` does not enable injection. - -A2A, MCP and direct calls do not automatically supply HTTP headers, so the default -is `None`. MCP tool arguments cannot populate the injected parameter. Wrappers -that require authentication must handle missing credentials explicitly. Receiving -headers does not itself authenticate the caller. - -Do not store caller credentials on the shared wrapper instance, echo them in -responses, or include them in logs or traces. Hayhooks keeps injected headers out -of its automatic `/run` payload logs and trace tags. - ### run_api() The `run_api()` method is called for each API request to the `{pipeline_name}/run` endpoint. @@ -169,7 +169,8 @@ def run_api(self, urls: list[str], question: str) -> str: **Input argument rules:** -- Arguments must be JSON-serializable +- Body arguments must be JSON-serializable, except [file uploads](#file-upload-support) +- The typed [`headers` parameter](#request-headers) is injected separately from the body - Use proper type hints (`list[str]`, `int | None`, etc.) - Default values are supported - Complex types like `dict[str, Any]` are allowed diff --git a/docs/features/a2a-support.md b/docs/features/a2a-support.md index e0da2f17..50442d3f 100644 --- a/docs/features/a2a-support.md +++ b/docs/features/a2a-support.md @@ -155,6 +155,7 @@ See [examples/a2a_multi_agent](https://github.com/deepset-ai/hayhooks/tree/main/ ## Current limitations +- **Request headers**: A2A calls leave the optional typed [`headers` parameter](../concepts/pipeline-wrapper.md#request-headers) at `None`; transport headers are not forwarded to the wrapper. - **Request-bound task execution**: Hayhooks currently treats A2A as a chat-shaped bridge. Each task runs inside the request handler by calling `run_chat_completion` / `run_chat_completion_async`, so non-streaming `SendMessage` returns after the task has completed or failed. This means detached task execution via [`returnImmediately`](https://a2a-protocol.org/latest/specification/#322-sendmessageconfiguration), [`input-required`](https://a2a-protocol.org/latest/specification/#63-multi-turn-interaction) pauses, and [push notification delivery](https://a2a-protocol.org/latest/specification/#353-push-notification-delivery) are not supported yet. - **Static agents list**: A2A routes are built from the registry at startup. Pipelines deployed or undeployed at runtime require restarting `hayhooks a2a run`. - **In-memory task store**: task state is kept in memory and lost on restart. diff --git a/docs/features/mcp-support.md b/docs/features/mcp-support.md index 69ad4689..e68046b5 100644 --- a/docs/features/mcp-support.md +++ b/docs/features/mcp-support.md @@ -129,7 +129,11 @@ For each deployed pipeline, Hayhooks will: - Parse **`run_api` method docstring**: - If you use Google-style or reStructuredText-style docstrings, use the first line as MCP Tool `description` and the rest as `parameters` (if present) - Each parameter description will be used as the `description` of the corresponding Pydantic model field (if present) -- Generate a Pydantic model from the `inputSchema` using the **`run_api` method arguments as fields** +- Generate the `inputSchema` from the **`run_api` method arguments**, excluding the optional typed [`headers` parameter](../concepts/pipeline-wrapper.md#request-headers) + +MCP calls leave the typed `headers` parameter at `None`, including over HTTP transports. +Supplying it as a tool argument is rejected. Ordinary application fields named `headers` +(for example, `headers: dict[str, str]`) remain normal tool arguments. **Example:** diff --git a/docs/features/openai-compatibility.md b/docs/features/openai-compatibility.md index 8e1b2bd9..c7155e17 100644 --- a/docs/features/openai-compatibility.md +++ b/docs/features/openai-compatibility.md @@ -18,6 +18,7 @@ Both APIs are available simultaneously. A pipeline wrapper can implement one or To access caller credentials or other HTTP headers, declare the optional typed [`headers` parameter](../concepts/pipeline-wrapper.md#request-headers) on your wrapper method. +This requires `fastapi-openai-compat>=1.3.0`, which Hayhooks declares as a dependency. ## Key Features diff --git a/src/hayhooks/server/pipelines/models.py b/src/hayhooks/server/pipelines/models.py index 1b2c03cc..14d4a812 100644 --- a/src/hayhooks/server/pipelines/models.py +++ b/src/hayhooks/server/pipelines/models.py @@ -8,54 +8,6 @@ from hayhooks.server.exceptions import PipelineWrapperError from hayhooks.server.utils.request_headers import accepts_request_headers -from hayhooks.server.utils.yaml_utils import InputResolution, OutputResolution - - -def get_request_model_from_resolved_io( - pipeline_name: str, declared_inputs: dict[str, InputResolution] -) -> type[BaseModel]: - """ - Create a flat Pydantic request model from declared inputs resolved by yaml_utils. - - Args: - pipeline_name: Name of the pipeline used for model naming. - declared_inputs: Mapping of declared input name to InputResolution. - - Returns: - A Pydantic model with top-level fields matching declared input names. - """ - fields: dict[str, Any] = {} - - for input_name, resolution in declared_inputs.items(): - input_type = resolution.type - default_value = ... if resolution.required else None - fields[input_name] = (input_type, default_value) - - return create_model(f"{pipeline_name.capitalize()}RunRequest", **fields) - - -def get_response_model_from_resolved_io( - pipeline_name: str, declared_outputs: dict[str, OutputResolution] -) -> type[BaseModel]: - """ - Create a flat Pydantic response model from declared outputs resolved by yaml_utils. - - Args: - pipeline_name: Name of the pipeline used for model naming. - declared_outputs: Mapping of declared output name to OutputResolution. - - Returns: - A Pydantic model with top-level fields matching declared output names. - """ - fields: dict[str, Any] = {} - - for output_name, resolution in declared_outputs.items(): - output_type = resolution.type - fields[output_name] = (output_type, ...) - - return create_model( - f"{pipeline_name.capitalize()}RunResponse", result=(dict, Field(..., description="Pipeline result")) - ) def create_request_model_from_callable(func: Callable, model_name: str, docstring: Docstring) -> type[BaseModel]: @@ -70,7 +22,7 @@ def create_request_model_from_callable(func: Callable, model_name: str, docstrin Pydantic model class for request """ - params = inspect.signature(func).parameters + params = inspect.signature(func, eval_str=True).parameters inject_headers = accepts_request_headers(func) param_docs = {p.arg_name: p.description for p in docstring.params} @@ -117,7 +69,7 @@ def create_response_model_from_callable( Pydantic model class for response, or None for streaming/file responses. """ - return_type = inspect.signature(func).return_annotation + return_type = inspect.signature(func, eval_str=True).return_annotation if return_type is inspect.Signature.empty: msg = f"Pipeline wrapper is missing a return type for '{func.__name__}' method" # ty: ignore[unresolved-attribute] @@ -158,7 +110,7 @@ def get_response_class_from_callable(func: Callable) -> type[Response] | None: * ``None`` for normal JSON endpoints (the caller should omit the ``response_class`` kwarg so FastAPI uses its default ``JSONResponse``). """ - return_type = inspect.signature(func).return_annotation + return_type = inspect.signature(func, eval_str=True).return_annotation if return_type is inspect.Signature.empty: return None diff --git a/tests/test_deploy_utils.py b/tests/test_deploy_utils.py index 1869b946..942f58ce 100644 --- a/tests/test_deploy_utils.py +++ b/tests/test_deploy_utils.py @@ -448,37 +448,6 @@ def sample_func_no_doc() -> int: assert "result" in schema["required"] -@pytest.mark.parametrize( - "return_type", - [ - Response, - FileResponse, - StreamingResponse, - Generator, - AsyncGenerator, - Generator[str, None, None], - AsyncGenerator[str, None], - ], - ids=[ - "Response", - "FileResponse", - "StreamingResponse", - "Generator", - "AsyncGenerator", - "Generator[str, None, None]", - "AsyncGenerator[str, None]", - ], -) -def test_create_response_model_returns_none_for_non_json_types(return_type): - func = lambda: None # noqa: E731 - func.__annotations__["return"] = return_type - - docstring = docstring_parser.parse("") - result = create_response_model_from_callable(func, "Test", docstring) - - assert result is None - - @pytest.mark.parametrize( ("return_type", "expected_class"), [ @@ -489,24 +458,21 @@ def test_create_response_model_returns_none_for_non_json_types(return_type): (AsyncGenerator, StreamingResponse), (Generator[str, None, None], StreamingResponse), (AsyncGenerator[str, None], StreamingResponse), - ], - ids=[ - "Response", - "FileResponse", - "StreamingResponse", - "Generator", - "AsyncGenerator", - "Generator[str, None, None]", - "AsyncGenerator[str, None]", + ("Response", Response), + ("FileResponse", FileResponse), + ("StreamingResponse", StreamingResponse), + ("Generator", StreamingResponse), + ("AsyncGenerator", StreamingResponse), + ("Generator[str, None, None]", StreamingResponse), + ("AsyncGenerator[str, None]", StreamingResponse), ], ) -def test_get_response_class_for_non_json_types(return_type, expected_class): +def test_non_json_return_types_skip_response_models_and_select_response_class(return_type, expected_class): func = lambda: None # noqa: E731 func.__annotations__["return"] = return_type - result = get_response_class_from_callable(func) - - assert result is expected_class + assert create_response_model_from_callable(func, "Test", docstring_parser.parse("")) is None + assert get_response_class_from_callable(func) is expected_class def test_get_response_class_returns_none_for_json_types(): @@ -586,7 +552,8 @@ def setup(self): with pytest.raises( PipelineWrapperError, match=re.escape( - "At least one of run_api, run_api_async, run_chat_completion, run_chat_completion_async, run_response, or run_response_async must be implemented" + "At least one of run_api, run_api_async, run_chat_completion, run_chat_completion_async, " + "run_response, or run_response_async must be implemented" ), ): create_pipeline_wrapper_instance(module) diff --git a/tests/test_request_headers.py b/tests/test_request_headers.py index ac6c3ba1..de3b34b5 100644 --- a/tests/test_request_headers.py +++ b/tests/test_request_headers.py @@ -70,15 +70,16 @@ def wrapper_source(method, stream=False, files=False): prefix = "async " if is_async else "" delay = "await asyncio.sleep(0.01)" if is_async else "time.sleep(0.01)" return_type = ("AsyncGenerator" if is_async else "Generator") if stream else "str" + header_value = "headers.get('Authorization', 'missing') if headers is not None else 'no-context'" result = ( - f" {prefix}def chunks():\n {delay}\n yield token\n return chunks()\n" + f" {prefix}def chunks():\n {delay}\n" + f" yield {header_value}\n return chunks()\n" if stream - else " return token\n" + else f" return {header_value}\n" ) return BASE_SOURCE + ( f" {prefix}def {method}(self, {arguments}, *, headers: Headers | None = None) -> {return_type}:\n" - f" {delay}\n" - " token = headers.get('Authorization', 'missing') if headers is not None else 'no-context'\n" + result + f" {delay}\n" + result ) @@ -93,8 +94,13 @@ def request_for(method, stream=False, prefix="/v1"): @pytest.mark.parametrize("method", METHODS) @pytest.mark.parametrize("stream", [False, True]) -async def test_headers_are_request_local_through_execution_and_streaming(headers_client, method, stream): - deploy(headers_client, wrapper_source(method, stream)) +@pytest.mark.parametrize("postponed", [False, True]) +async def test_headers_are_request_local_through_execution_and_streaming(headers_client, method, stream, postponed): + source = wrapper_source(method, stream) + deploy(headers_client, "from __future__ import annotations\n" + source if postponed else source) + if method.startswith("run_api"): + schema = headers_client.get("/openapi.json").json()["components"]["schemas"]["headers_testRunRequest"] + assert set(schema["properties"]) == {"query"} url, body = request_for(method, stream) tokens = ["Bearer alice-test", "Bearer bob-test", None] async with AsyncClient(transport=ASGITransport(app=headers_client.app), base_url="http://test") as client: @@ -121,8 +127,10 @@ def test_openai_aliases_forward_headers(headers_client, method): @pytest.mark.parametrize("method", ["run_api", "run_api_async"]) -def test_multipart_headers_are_injected_outside_the_form(headers_client, method): - deploy(headers_client, wrapper_source(method, files=True)) +@pytest.mark.parametrize("postponed", [False, True]) +def test_multipart_headers_are_injected_outside_the_form(headers_client, method, postponed): + source = wrapper_source(method, files=True) + deploy(headers_client, "from __future__ import annotations\n" + source if postponed else source) response = headers_client.post( "/headers_test/run", data={"query": "hi", "headers": "body-spoof"}, @@ -136,10 +144,8 @@ def test_multipart_headers_are_injected_outside_the_form(headers_client, method) @pytest.mark.parametrize("method", ["run_api", "run_api_async"]) @pytest.mark.skipif(importlib.util.find_spec("mcp") is None, reason="MCP is not installed") @pytest.mark.mcp -async def test_injected_headers_stay_out_of_schemas_and_mcp_arguments(headers_client, method): +async def test_injected_headers_stay_out_of_mcp_schema_and_arguments(headers_client, method): deploy(headers_client, wrapper_source(method)) - schema = headers_client.get("/openapi.json").json()["components"]["schemas"]["headers_testRunRequest"] - assert set(schema["properties"]) == {"query"} tools = await list_pipelines_as_tools() assert set(tools[0].inputSchema["properties"]) == {"query"} result = await run_pipeline_as_tool("headers_test", {"query": "hi"}) @@ -199,9 +205,20 @@ def test_legacy_openai_signatures_do_not_receive_headers(headers_client, extra, @pytest.mark.parametrize("method", ["run_api", "run_api_async"]) -@pytest.mark.skipif(importlib.util.find_spec("mcp") is None, reason="MCP is not installed") -@pytest.mark.mcp -async def test_regular_headers_body_field_keeps_its_schema_and_value(headers_client, method): +@pytest.mark.parametrize( + "transport", + [ + "http", + pytest.param( + "mcp", + marks=[ + pytest.mark.mcp, + pytest.mark.skipif(importlib.util.find_spec("mcp") is None, reason="MCP is not installed"), + ], + ), + ], +) +async def test_regular_headers_body_field_keeps_its_schema_and_value(headers_client, method, transport): prefix = "async " if method.endswith("_async") else "" deploy( headers_client, @@ -211,16 +228,21 @@ async def test_regular_headers_body_field_keeps_its_schema_and_value(headers_cli " return headers['authorization']\n" ), ) - response = headers_client.post( - "/headers_test/run", - json={"headers": {"authorization": "body-value"}}, - headers={"Authorization": "transport-value"}, - ) - assert response.json() == {"result": "body-value"} - tools = await list_pipelines_as_tools() - assert tools[0].inputSchema["required"] == ["headers"] - result = await run_pipeline_as_tool("headers_test", {"headers": {"authorization": "tool-value"}}) - assert result[0].text == "tool-value" + if transport == "http": + schema = headers_client.get("/openapi.json").json()["components"]["schemas"]["headers_testRunRequest"] + assert schema["required"] == ["headers"] + response = headers_client.post( + "/headers_test/run", + json={"headers": {"authorization": "body-value"}}, + headers={"Authorization": "transport-value"}, + ) + assert response.status_code == 200, response.text + assert response.json() == {"result": "body-value"} + else: + tools = await list_pipelines_as_tools() + assert tools[0].inputSchema["required"] == ["headers"] + result = await run_pipeline_as_tool("headers_test", {"headers": {"authorization": "tool-value"}}) + assert result[0].text == "tool-value" @pytest.mark.parametrize( From 63f36c79c00d421d11592cf3eaae44a31d2bdac4 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 11:21:38 +0200 Subject: [PATCH 09/11] fix: require published compat 1.4.0 and remove stale build include --- docs/features/openai-compatibility.md | 2 +- pyproject.toml | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/features/openai-compatibility.md b/docs/features/openai-compatibility.md index c7155e17..65df1c91 100644 --- a/docs/features/openai-compatibility.md +++ b/docs/features/openai-compatibility.md @@ -18,7 +18,7 @@ Both APIs are available simultaneously. A pipeline wrapper can implement one or To access caller credentials or other HTTP headers, declare the optional typed [`headers` parameter](../concepts/pipeline-wrapper.md#request-headers) on your wrapper method. -This requires `fastapi-openai-compat>=1.3.0`, which Hayhooks declares as a dependency. +This requires `fastapi-openai-compat>=1.4.0`, which Hayhooks declares as a dependency. ## Key Features diff --git a/pyproject.toml b/pyproject.toml index 377e3af6..67f9392a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "pydantic-settings", "python-dotenv", "docstring-parser", - "fastapi-openai-compat>=1.3.0", + "fastapi-openai-compat>=1.4.0", ] [project.optional-dependencies] @@ -76,8 +76,6 @@ source = "vcs" packages = ["src/hayhooks"] [tool.hatch.build.targets.wheel.force-include] -# .chainlit is a hidden dir excluded by default, so force-include is needed -"src/hayhooks/server/chainlit_app/.chainlit" = "hayhooks/server/chainlit_app/.chainlit" # Bundle dashboard source files so --with-tracing-dashboard can build locally at runtime. "dashboard/index.html" = "hayhooks/dashboard/index.html" "dashboard/components.json" = "hayhooks/dashboard/components.json" From b52a0b07f9350625fd09ed8ae64eedf26ee3b4e2 Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 11:21:38 +0200 Subject: [PATCH 10/11] test: handle component cancellation during disconnect cleanup --- tests/test_it_streaming_disconnect.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/test_it_streaming_disconnect.py b/tests/test_it_streaming_disconnect.py index 647c671d..b6ef8e94 100644 --- a/tests/test_it_streaming_disconnect.py +++ b/tests/test_it_streaming_disconnect.py @@ -16,18 +16,21 @@ class _BlockingStreamingComponent: def __init__(self) -> None: self.release = asyncio.Event() - self.completed = asyncio.Event() + self.finished = asyncio.Event() @component.output_types(result=str) def run(self, streaming_callback: Any | None = None) -> dict[str, str]: - raise AssertionError("The async pipeline must call run_async") + msg = "The async pipeline must call run_async" + raise AssertionError(msg) @component.output_types(result=str) async def run_async(self, streaming_callback: Any | None = None) -> dict[str, str]: - await streaming_callback(StreamingChunk(content="first", index=0)) - await self.release.wait() - self.completed.set() - return {"result": "done"} + try: + await streaming_callback(StreamingChunk(content="first", index=0)) + await self.release.wait() + return {"result": "done"} + finally: + self.finished.set() async def _disconnect_after_first_chunk(app: FastAPI) -> None: @@ -88,9 +91,12 @@ async def stream(): assert bool(detached_tasks) is shield_pipeline_task assert all(not task.done() for task in detached_tasks) + if shield_pipeline_task: + assert not component.finished.is_set() component.release.set() - await asyncio.wait_for(component.completed.wait(), timeout=1.0) + # Haystack may either complete the component or cancel it with the pipeline task. + await asyncio.wait_for(component.finished.wait(), timeout=1.0) if detached_tasks: await asyncio.wait_for(asyncio.gather(*detached_tasks), timeout=1.0) await asyncio.sleep(0) From 52317ef420a8fb71887b919ccd9e5900074f18ee Mon Sep 17 00:00:00 2001 From: Michele Pangrazzi Date: Mon, 28 Sep 2026 11:54:14 +0200 Subject: [PATCH 11/11] fix: keep tracing type checks valid with optional extras --- pyproject.toml | 7 +++++++ src/hayhooks/server/tracing.py | 3 ++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 67f9392a..bad1ba28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -143,6 +143,13 @@ python-version = "3.10" [tool.ty.rules] unused-ignore-comment = "warn" +[[tool.ty.overrides]] +include = ["src/hayhooks/server/tracing.py"] + +[tool.ty.overrides.rules] +# Optional tracing imports resolve only when their extras are installed. +unused-ignore-comment = "ignore" + [tool.ty.src] exclude = ["tests/**/*"] diff --git a/src/hayhooks/server/tracing.py b/src/hayhooks/server/tracing.py index 9a883ab5..e4ce4972 100644 --- a/src/hayhooks/server/tracing.py +++ b/src/hayhooks/server/tracing.py @@ -640,10 +640,11 @@ def finish(self, exc: BaseException | None = None) -> None: if self._finished: return + live_tags: dict[str, Any] = {} try: if exc is None: _mark_success(span) - live_tags: dict[str, Any] = {_TAG_SUCCESS: True} + live_tags = {_TAG_SUCCESS: True} elif isinstance(exc, HTTPException): _mark_http_exception(span, exc) live_tags = {