feat(run): propagate the run session as an x-session-id header (ENG-3473)

aixplain/v2/model.py

@@ -493,9 +493,16 @@ class ModelRunParams(BaseRunParams):
     Attributes:
         stream: If True, returns a ModelResponseStreamer for streaming responses.
             The model must support streaming (check supports_streaming attribute).
+        session_id: Conversation this run belongs to, emitted as the
+            ``x-session-id`` header so downstream services can correlate the call
+            with the session that triggered it. Header-only — stripped from the
+            model/action input payload and the run URL, exactly like
+            ``identifier`` (→ ``x-user-id``). Omit it (or pass ``None``) to send
+            no header.
     """
 
     stream: NotRequired[bool]
+    session_id: NotRequired[Optional[str]]
 
 
 @dataclass_json
@@ -664,22 +671,37 @@ def __setattr__(self, name: str, value):
             super().__setattr__(name, value)
 
     _SDK_ONLY_PARAMS = frozenset(
-        {"timeout", "wait_time", "show_progress", "stream", "run_retries", "run_retry_wait", "identifier"}
+        {
+            "timeout",
+            "wait_time",
+            "show_progress",
+            "stream",
+            "run_retries",
+            "run_retry_wait",
+            "identifier",
+            "session_id",
+        }
     )
 
     # ``identifier`` is a per-run caller identity emitted as the ``x-user-id``
     # header (RunnableResourceMixin._headers_for_run), never a model/action
     # input. Exclude it from the v2 payload builder here (and from the v1/URL
     # builders via _SDK_ONLY_PARAMS above) so it cannot leak into model inputs
     # or supplier-facing logs. Agent keeps it in the body by NOT overriding this.
+    #
+    # ``session_id`` is the same kind of per-run metadata (emitted as
+    # ``x-session-id``) and is header-only for every runnable, so the base
+    # _RUN_CONTROL_KEYS already excludes it; it is repeated in
+    # _SDK_ONLY_PARAMS above to also cover the v1/URL builder paths.
     _RUN_CONTROL_KEYS = RunnableResourceMixin._RUN_CONTROL_KEYS | {"identifier"}
 
     def build_run_payload(self, **kwargs: Unpack[ModelRunParams]) -> dict:
         """Build the JSON payload for a model execution request.
 
         Strips SDK-only orchestration params (``timeout``, ``wait_time``,
-        ``show_progress``, ``stream``, ``run_retries``, ``run_retry_wait``)
-        so they are never forwarded to the backend API.
+        ``show_progress``, ``stream``, ``run_retries``, ``run_retry_wait``) and
+        the header-only run metadata (``identifier``, ``session_id``) so they are
+        never forwarded to the backend API.
         """
         filtered = {k: v for k, v in kwargs.items() if k not in self._SDK_ONLY_PARAMS}
         return super().build_run_payload(**filtered)

aixplain/v2/resource.py

@@ -578,6 +578,14 @@ class BaseRunParams(BaseParams):
         wait_time: Initial interval in seconds between poll attempts.
         run_retries: Extra attempts after the first failure (total attempts = 1 + run_retries).
         run_retry_wait: Seconds to wait between retry attempts (default 1.0).
+
+    Note:
+        ``session_id`` is a header-only run key handled by every runnable
+        (``_headers_for_run`` → ``x-session-id``, stripped from the body by
+        ``_RUN_CONTROL_KEYS``). It is declared on :class:`ModelRunParams` rather
+        than here: :class:`~aixplain.v2.agent.AgentRunParams` deliberately has no
+        ``session_id`` — agent runs join a conversation via ``session=``, and a
+        look-alike key that only set a header would be a footgun.
     """
 
     timeout: NotRequired[int]
@@ -1223,13 +1231,20 @@ class RunnableResourceMixin(BaseMixin, Generic[RunParamsT, ResultT]):
     # is a legitimate execution-config body field (see Agent). Model/Tool, where
     # it is header-only, add it to their own override so it never leaks into the
     # model/action input payload.
+    #
+    # ``session_id`` IS excluded here (unlike ``identifier``): it is header-only
+    # for every runnable. No run params type declares it as a body field — the
+    # Agent session path takes ``session=`` (a Session or id) and routes through
+    # ``POST /v1/sessions/{id}/messages``, so a top-level ``session_id`` kwarg is
+    # purely the per-run correlation channel emitted as ``x-session-id``.
     _RUN_CONTROL_KEYS: frozenset[str] = frozenset(
         {
             "run_retries",
             "run_retry_wait",
             "timeout",
             "wait_time",
             "show_progress",
+            "session_id",
         }
     )
 
@@ -1262,11 +1277,21 @@ def _payload_kwargs_for_run(self, kwargs: dict) -> dict:
         return {k: v for k, v in kwargs.items() if k not in self._RUN_CONTROL_KEYS}
 
     def _headers_for_run(self, kwargs: dict) -> Optional[dict]:
-        """Build per-run headers from optional runtime metadata."""
+        """Build per-run headers from optional runtime metadata.
+
+        ``identifier`` → ``x-user-id`` (caller identity) and ``session_id`` →
+        ``x-session-id`` (conversation the run belongs to). Both are optional and
+        independent: a key that is absent or ``None`` simply omits its header, and
+        with neither present this returns ``None`` so no headers are attached.
+        """
+        headers = {}
         identifier = kwargs.get("identifier")
-        if identifier is None:
-            return None
-        return {"x-user-id": str(identifier)}
+        if identifier is not None:
+            headers["x-user-id"] = str(identifier)
+        session_id = kwargs.get("session_id")
+        if session_id is not None:
+            headers["x-session-id"] = str(session_id)
+        return headers or None
 
     def _post_and_handle_run(self, **kwargs: Unpack[RunParamsT]) -> ResultT:
         """Single POST + handle_run_response (no retries, no before_run)."""

tests/unit/v2/test_model.py

@@ -384,6 +384,47 @@ def test_identifier_excluded_from_build_run_payload(self):
         assert "identifier" not in payload
 
 
+class TestModelSessionHeader:
+    """session_id must ride as the x-session-id header, never as a model input."""
+
+    def _create_model(self):
+        model = Model.__new__(Model)
+        model.id = "test-model-id"
+        model.name = "Test Model"
+        model.connection_type = ["synchronous"]
+        model.params = None
+        model.__post_init__()
+        model.context = Mock()
+        return model
+
+    def test_session_id_excluded_from_payload(self):
+        """The run's session is per-run metadata, not a model input — it must not
+        reach the model (or supplier-facing logs) as an input field."""
+        model = self._create_model()
+        payload_kwargs = model._payload_kwargs_for_run({"text": "hi", "session_id": "sess-1"})
+        assert "session_id" not in payload_kwargs
+        assert payload_kwargs["text"] == "hi"
+
+    def test_session_id_emitted_as_header(self):
+        model = self._create_model()
+        assert model._headers_for_run({"text": "hi", "session_id": "sess-1"}) == {"x-session-id": "sess-1"}
+
+    def test_session_id_excluded_from_build_run_payload(self):
+        """Also excluded from the v1/URL payload builder path (_SDK_ONLY_PARAMS)."""
+        model = self._create_model()
+        payload = model.build_run_payload(text="hi", session_id="sess-1")
+        assert "session_id" not in payload
+
+    def test_identifier_and_session_id_ride_together(self):
+        """Both headers on one run, and neither in the payload."""
+        model = self._create_model()
+        kwargs = {"text": "hi", "identifier": "alice", "session_id": "sess-1"}
+        assert model._headers_for_run(kwargs) == {"x-user-id": "alice", "x-session-id": "sess-1"}
+        payload_kwargs = model._payload_kwargs_for_run(kwargs)
+        assert "identifier" not in payload_kwargs
+        assert "session_id" not in payload_kwargs
+
+
 class TestModelV1Fallback:
     """Tests for _run_async_v1() V1 endpoint integration."""
 

tests/unit/v2/test_resource.py

@@ -989,6 +989,36 @@ def test_base_keeps_identifier_in_payload(self):
         assert payload_kwargs["identifier"] == "alice"
         assert payload_kwargs["text"] == "hi"
 
+    def test_session_id_emitted_as_header(self):
+        """session_id is emitted as the x-session-id header for any runnable resource."""
+        resource = self._create_runnable_resource()
+
+        headers = resource._headers_for_run({"text": "hi", "session_id": "sess-1"})
+        assert headers == {"x-session-id": "sess-1"}
+
+    def test_base_excludes_session_id_from_payload(self):
+        """Unlike identifier, session_id is header-only for every runnable — no run
+        params type declares it as a body field, so the base mixin strips it."""
+        resource = self._create_runnable_resource()
+
+        payload_kwargs = resource._payload_kwargs_for_run({"text": "hi", "session_id": "sess-1"})
+        assert "session_id" not in payload_kwargs
+        assert payload_kwargs["text"] == "hi"
+
+    def test_both_run_metadata_headers_emitted_together(self):
+        resource = self._create_runnable_resource()
+
+        headers = resource._headers_for_run({"identifier": "alice", "session_id": "sess-1"})
+        assert headers == {"x-user-id": "alice", "x-session-id": "sess-1"}
+
+    def test_no_headers_when_no_run_metadata(self):
+        """Neither key present → no headers dict at all, so _post_and_handle_run
+        attaches nothing (documented default, never a crash)."""
+        resource = self._create_runnable_resource()
+
+        assert resource._headers_for_run({"text": "hi"}) is None
+        assert resource._headers_for_run({"identifier": None, "session_id": None}) is None
+
 
 # =============================================================================
 # Result Class Tests

tests/unit/v2/test_tool.py

@@ -668,3 +668,18 @@ def test_list_data_preserves_extra_kwargs(self):
 
         assert result["identifier"] == "alice"
         assert result["data"] == ["x"]
+
+    def test_run_metadata_survives_merge_and_stays_header_only(self):
+        """``identifier``/``session_id`` must survive the merge (so
+        ``_headers_for_run`` can see them) while staying out of the action payload."""
+        tool = self._make_tool()
+
+        merged = tool._merge_with_dynamic_attrs(
+            action="search", data={"q": "hi"}, identifier="alice", session_id="sess-1"
+        )
+
+        assert tool._headers_for_run(merged) == {"x-user-id": "alice", "x-session-id": "sess-1"}
+        payload_kwargs = tool._payload_kwargs_for_run(merged)
+        assert "identifier" not in payload_kwargs
+        assert "session_id" not in payload_kwargs
+        assert payload_kwargs["data"] == {"q": "hi"}