feat(weave): propagate Custom Runtime conversation context

tests/trace/test_chat_completions.py

@@ -37,6 +37,22 @@ def _chunk_payload(content: str = "Hello") -> dict:
     }
 
 
+def _completion_payload(content: str = "Hello") -> dict:
+    return {
+        "id": "completion-1",
+        "choices": [
+            {
+                "finish_reason": "stop",
+                "index": 0,
+                "message": {"role": "assistant", "content": content},
+            }
+        ],
+        "created": 1,
+        "model": "runtime-model",
+        "object": "chat.completion",
+    }
+
+
 class _TrackingByteStream(httpx.SyncByteStream):
     def __init__(self, content: bytes, close_error: Exception | None = None) -> None:
         self.content = content
@@ -208,3 +224,102 @@ def handler(_request: httpx.Request) -> httpx.Response:
     assert stream.response.is_closed
     assert body.closed
     assert clients[0].is_closed
+
+
+def test_playground_completion_reuses_conversation_across_provider_switches(
+    monkeypatch,
+) -> None:
+    request_bodies = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        body = json.loads(request.content)
+        request_bodies.append(body)
+        conversation_id = body.get("conversation_id", "server-conversation")
+        return httpx.Response(
+            200,
+            json={
+                "response": _completion_payload(),
+                "conversation_id": conversation_id,
+            },
+        )
+
+    _install_mock_transport(monkeypatch, handler)
+    completions = Completions(_client())
+
+    first = completions.create(
+        endpoint="playground",
+        model="custom::runtime-a::model-a",
+        messages=[{"role": "user", "content": "Hello"}],
+        track_llm_call=False,
+    )
+    second = completions.create(
+        endpoint="playground",
+        model="openai/gpt-4o",
+        messages=[{"role": "user", "content": "Continue"}],
+        conversation_id=first.conversation_id,
+        track_llm_call=False,
+    )
+    third = completions.create(
+        endpoint="playground",
+        model="custom::runtime-b::model-b",
+        messages=[{"role": "user", "content": "Finish"}],
+        conversation_id=second.conversation_id,
+        track_llm_call=False,
+    )
+
+    assert second.conversation_id == first.conversation_id
+    assert third.conversation_id == first.conversation_id
+    assert [body["inputs"]["model"] for body in request_bodies] == [
+        "custom::runtime-a::model-a",
+        "openai/gpt-4o",
+        "custom::runtime-b::model-b",
+    ]
+    assert "conversation_id" not in request_bodies[0]
+    assert request_bodies[1]["conversation_id"] == first.conversation_id
+    assert request_bodies[2]["conversation_id"] == first.conversation_id
+
+
+def test_custom_runtime_stream_consumes_server_conversation_context(
+    monkeypatch,
+) -> None:
+    def handler(_request: httpx.Request) -> httpx.Response:
+        content = (
+            json.dumps({"_meta": {"conversation_id": "server-conversation"}}).encode()
+            + b"\n"
+            + json.dumps(_chunk_payload()).encode()
+            + b"\n"
+        )
+        return httpx.Response(200, content=content)
+
+    _install_mock_transport(monkeypatch, handler)
+    completions = Completions(_client())
+
+    stream = completions.create(
+        endpoint="playground",
+        model="custom::runtime::model",
+        messages=[{"role": "user", "content": "Hello"}],
+        stream=True,
+        track_llm_call=False,
+    )
+
+    assert [chunk.choices[0].delta.content for chunk in stream] == ["Hello"]
+    assert stream.conversation_id == "server-conversation"
+
+
+def test_inference_completion_has_no_playground_conversation_context(
+    monkeypatch,
+) -> None:
+    def handler(_request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, json=_completion_payload())
+
+    _install_mock_transport(monkeypatch, handler)
+    completions = Completions(_client())
+
+    result = completions.create(
+        endpoint="inference",
+        model="coreweave/runtime-model",
+        messages=[{"role": "user", "content": "Hello"}],
+        track_llm_call=False,
+    )
+
+    assert "conversation_id" not in result.model_dump()

tests/trace_server/test_async_clickhouse_trace_server.py

@@ -24,6 +24,7 @@
 from weave.trace_server.external_to_internal_trace_server_adapter import (
     ExternalTraceServer,
 )
+from weave.trace_server.llm_completion import CustomProviderInfo
 
 LITELLM_ACOMPLETION_PATCH = (
     "weave.trace_server.async_clickhouse_trace_server.lite_llm_acompletion"
@@ -114,6 +115,42 @@ async def test_tracking_routes_through_log_completion_call(
     assert forwarded_res is llm_res
 
 
+@pytest.mark.asyncio
+async def test_async_custom_runtime_preserves_tracking_and_forwards_context(
+    server: AsyncClickHouseTraceServer,
+) -> None:
+    llm_res = tsi.CompletionsCreateRes(response={"choices": [{"x": 1}]})
+    completion = AsyncMock(return_value=llm_res)
+    req = _make_req(track_llm_call=True, model="custom::runtime::model")
+    with (
+        patch(
+            "weave.trace_server.clickhouse_trace_server_batched.get_custom_provider_info",
+            return_value=CustomProviderInfo(
+                base_url="https://runtime.example.com/v1",
+                api_key="runtime-key",
+                extra_headers={"X-Tenant": "customer"},
+                return_type="openai",
+                actual_model_name="model",
+            ),
+        ),
+        patch(LITELLM_ACOMPLETION_PATCH, new=completion),
+        patch.object(
+            server,
+            "_log_completion_call",
+            return_value=tsi.CompletionsCreateRes(response=llm_res.response),
+        ) as log_completion,
+    ):
+        await server.acompletions_create(req)
+
+    assert req.conversation_id
+    assert req.track_llm_call is True
+    log_completion.assert_called_once()
+    assert completion.await_args.kwargs["extra_headers"] == {
+        "X-Tenant": "customer",
+        "X-Weave-Conversation-Id": req.conversation_id,
+    }
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("_mock_secret_fetcher")
 async def test_deferred_returns_real_span_without_inserting() -> None:

tests/trace_server/test_clickhouse_trace_server_batched.py

@@ -566,10 +566,11 @@ def mock_obj_read_func(req):
         stream = server.completions_create_stream(req)
         chunks = list(stream)
 
-        assert len(chunks) == 2
-        assert chunks[0]["choices"][0]["delta"]["content"] == "Streamed"
-        assert chunks[1]["choices"][0]["finish_reason"] == "stop"
-        assert "usage" in chunks[1]
+        assert req.conversation_id
+        assert chunks[0] == {"_meta": {"conversation_id": req.conversation_id}}
+        assert chunks[1]["choices"][0]["delta"]["content"] == "Streamed"
+        assert chunks[2]["choices"][0]["finish_reason"] == "stop"
+        assert "usage" in chunks[2]
 
         # Verify litellm was called with correct parameters
         mock_litellm.assert_called_once()
@@ -578,7 +579,10 @@ def mock_obj_read_func(req):
             call_args.get("api_base")
             or call_args.get("base_url") == "https://api.custom.com"
         )
-        assert call_args["extra_headers"] == {"X-Custom": "value"}
+        assert call_args["extra_headers"] == {
+            "X-Custom": "value",
+            "X-Weave-Conversation-Id": req.conversation_id,
+        }
 
 
 def test_completions_create_stream_custom_provider_with_tracking():
@@ -709,6 +713,8 @@ def mock_obj_read_func(req):
         assert len(chunks) == 3  # Meta chunk + 2 content chunks
         assert "_meta" in chunks[0]
         assert "weave_call_id" in chunks[0]["_meta"]
+        assert req.track_llm_call is True
+        assert chunks[0]["_meta"]["conversation_id"] == req.conversation_id
         assert chunks[1]["choices"][0]["delta"]["content"] == "Streamed"
         assert chunks[2]["choices"][0]["finish_reason"] == "stop"
         assert "usage" in chunks[2]
@@ -724,7 +730,10 @@ def mock_obj_read_func(req):
             call_args.get("api_base")
             or call_args.get("base_url") == "https://api.custom.com"
         )
-        assert call_args["extra_headers"] == {"X-Custom": "value"}
+        assert call_args["extra_headers"] == {
+            "X-Custom": "value",
+            "X-Weave-Conversation-Id": req.conversation_id,
+        }
 
 
 def test_completions_create_stream_multiple_choices():

tests/trace_server/test_custom_provider.py

@@ -363,10 +363,11 @@ def test_custom_provider_completions_create(client):
                 f"API base URL mismatch. Expected 'https://api.example.com', "
                 f"got '{call_args['api_base']}'"
             )
-            assert call_args["extra_headers"] == {"X-Custom-Header": "value"}, (
-                f"Extra headers mismatch. Expected {{'X-Custom-Header': 'value'}}, "
-                f"got {call_args['extra_headers']}"
-            )
+            assert res.conversation_id
+            assert call_args["extra_headers"] == {
+                "X-Custom-Header": "value",
+                "X-Weave-Conversation-Id": res.conversation_id,
+            }
 
             # Completions now write to the spans table, not calls.
             # Verify the span was created with correct identifiers.

tests/trace_server/test_llm_completion.py

@@ -672,17 +672,21 @@ def mock_obj_read_func(req):
             chunks = list(stream)
 
             # Verify the chunks
-            assert len(chunks) == 2
-            assert chunks[0]["choices"][0]["delta"]["content"] == "Custom"
-            assert chunks[1]["choices"][0]["finish_reason"] == "stop"
+            assert req.conversation_id
+            assert chunks[0] == {"_meta": {"conversation_id": req.conversation_id}}
+            assert chunks[1]["choices"][0]["delta"]["content"] == "Custom"
+            assert chunks[2]["choices"][0]["finish_reason"] == "stop"
 
             # Verify litellm was called with correct parameters
             mock_litellm.assert_called_once()
             call_args = mock_litellm.call_args[1]
             assert (
                 call_args.get("api_base") or call_args.get("base_url")
             ) == "https://api.custom.com"
-            assert call_args["extra_headers"] == {"X-Custom": "value"}
+            assert call_args["extra_headers"] == {
+                "X-Custom": "value",
+                "X-Weave-Conversation-Id": req.conversation_id,
+            }
 
     def test_missing_api_key(self):
         """Test handling of missing API key in streaming completion."""
@@ -2281,6 +2285,64 @@ def test_custom_provider_name_matching_selector_prefix_is_preserved(monkeypatch)
     assert model_info.model_name == "custom/gpt-4"
 
 
+def test_custom_runtime_context_overrides_configured_header(monkeypatch):
+    req = tsi.CompletionsCreateReq(
+        project_id="entity/project",
+        inputs=_completion_inputs(model="custom::runtime::model"),
+        track_llm_call=False,
+        conversation_id="conv-1",
+    )
+    monkeypatch.setattr(
+        chts,
+        "get_custom_provider_info",
+        MagicMock(
+            return_value=llm_mod.CustomProviderInfo(
+                base_url="https://runtime.example.com/v1",
+                api_key=None,
+                extra_headers={
+                    "X-Tenant": "customer",
+                    "x-weave-conversation-id": "configured-value",
+                },
+                return_type="openai",
+                actual_model_name="model",
+            )
+        ),
+    )
+
+    info = chts._setup_completion_model_info(None, req, MagicMock())
+
+    assert req.track_llm_call is False
+    assert info.extra_headers == {
+        "X-Tenant": "customer",
+        "X-Weave-Conversation-Id": "conv-1",
+    }
+
+
+def test_custom_runtime_context_is_header_safe(monkeypatch):
+    req = tsi.CompletionsCreateReq(
+        project_id="entity/project",
+        inputs=_completion_inputs(model="custom::runtime::model"),
+        conversation_id="support/café\n",
+    )
+    monkeypatch.setattr(
+        chts,
+        "get_custom_provider_info",
+        MagicMock(
+            return_value=llm_mod.CustomProviderInfo(
+                base_url="https://runtime.example.com/v1",
+                api_key=None,
+                extra_headers={},
+                return_type="openai",
+                actual_model_name="model",
+            )
+        ),
+    )
+
+    info = chts._setup_completion_model_info(None, req, MagicMock())
+
+    assert info.extra_headers == {"X-Weave-Conversation-Id": "support%2Fcaf%C3%A9%0A"}
+
+
 def test_coreweave_with_api_key_keeps_litellm_configuration():
     req = tsi.CompletionsCreateReq(
         project_id="entity/project",