feat(weave): rebuild RemoteHTTPTraceServer on weave-server-sdk

pyproject.toml

@@ -58,6 +58,11 @@ dependencies = [
   # (gen_ai.tool.call.arguments, cache token names,
   # gen_ai.usage.reasoning.output_tokens) added in 0.63b0.
   "opentelemetry-semantic-conventions>=0.63b0",
+
+  # Generated trace-server client (source of truth for the API types).
+  # TEMPORARY: resolved from test PyPI via [tool.uv.sources] below until the
+  # package is published to real PyPI.
+  "weave-server-sdk==0.0.1",
 ]
 
 [project.optional-dependencies]
@@ -233,6 +238,7 @@ select = [
   "FIX003",  # https://docs.astral.sh/ruff/rules/line-contains-xxx/
   "I",       # https://docs.astral.sh/ruff/rules/#isort-i
   "W",       # https://docs.astral.sh/ruff/rules/#warning-w
+  "TID251",  # banned-api: enforces the client/server import boundary (see flake8-tidy-imports below)
   "TID252",  # https://docs.astral.sh/ruff/rules/relative-imports/#relative-imports-tid252
   "UP",      # https://docs.astral.sh/ruff/rules/#pyupgrade-up
   "TRY",     # https://docs.astral.sh/ruff/rules/#tryceratops-try
@@ -318,12 +324,25 @@ line-length = 88
 show-fixes = true
 exclude = ["rules"]
 
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+# Client code must not depend on the trace server's interface modules; API
+# types come from weave_server_sdk (generated from the OpenAPI spec). The ban
+# list grows as the migration proceeds. The server package itself, tests, and
+# scripts are exempt via per-file-ignores below.
+"weave.trace_server.http_service_interface".msg = "Client code must not use the server's HTTP body models. Use weave_server_sdk.models (generated from the OpenAPI spec) instead."
+
 [tool.ruff.lint.per-file-ignores]
 "!/weave/trace/**/*.py" = ["T201"]
 "!/tests/**/*.py" = ["RUF059"]
-# Tests intentionally use async functions without await, compare known float values,
-# and use pytest.raises match strings with regex-like characters.
-"tests/**/*.py" = ["RUF029", "RUF043", "RUF069"]
+# The trace server may import its own modules; tests exercise the server
+# directly; scripts and the mock server are operational server tooling.
+"weave/trace_server/**/*.py" = ["TID251"]
+"scripts/**/*.py" = ["TID251"]
+"trace_server_mock/**/*.py" = ["TID251"]
+# Tests intentionally use async functions without await, compare known float
+# values, use pytest.raises match strings with regex-like characters, and may
+# exercise the trace server directly (TID251).
+"tests/**/*.py" = ["RUF029", "RUF043", "RUF069", "TID251"]
 "weave/trace/serialization/op_type.py" = ["RUF100", "N802"]
 "weave/trace_server/costs/update_costs.py" = ["PLW1514"]
 "weave/type_handlers/Video/video.py" = ["PLW0603"]
@@ -526,6 +545,17 @@ conflicts = [
   ],
 ]
 
+# TEMPORARY: weave-server-sdk is only published to test PyPI so far. Resolve
+# just that package from the test index; everything else stays on real PyPI.
+# Remove this (and the index below) once it is published to real PyPI.
+[tool.uv.sources]
+weave-server-sdk = { index = "testpypi" }
+
+[[tool.uv.index]]
+name = "testpypi"
+url = "https://test.pypi.org/simple/"
+explicit = true
+
 # uv workspace members. Each listed path is a sibling Python package that uv
 # manages alongside the root `weave` project. Members can depend on each
 # other and on the root by name; uv resolves those to the local workspace

tests/conftest.py

@@ -7,6 +7,7 @@
 from typing import Any
 from unittest.mock import MagicMock, patch
 
+import httpx
 import pytest
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
@@ -597,11 +598,12 @@ def client(
 
 @pytest.fixture
 def network_proxy_client(client, monkeypatch):
-    """This fixture is used to test the `RemoteHTTPTraceServer` class. There is
-    almost no logic in this class, other than a little batching, so we typically
-    skip it for simplicity. However, we can use this fixture to test such logic.
-    It initializes a mini FastAPI app that proxies requests from the
-    `RemoteHTTPTraceServer` to the underlying `client.server` object.
+    """This fixture is used to test the `RemoteHTTPTraceServer` class.
+    There is almost no logic in this class, other than a little batching, so we
+    typically skip it for simplicity. However, we can use this fixture to test
+    such logic. It initializes a mini FastAPI app and routes the server's HTTP
+    transport into it, proxying requests to the underlying `client.server`
+    object.
 
     We probably will want to flesh this out more in the future, but this is a
     starting point.
@@ -707,12 +709,25 @@ def obj_read(req: tsi.ObjReadReq) -> tsi.ObjReadRes:
 
     with TestClient(app) as c:
 
-        def post(url, data=None, json=None, **kwargs):
-            kwargs.pop("stream", None)
-            return c.post(url, data=data, json=json, **kwargs)
-
-        orig_post = weave.utils.http_requests.post
-        weave.utils.http_requests.post = post
+        class TestClientTransport(httpx.BaseTransport):
+            """Routes the server's httpx requests into the FastAPI TestClient."""
+
+            def handle_request(self, request: httpx.Request) -> httpx.Response:
+                request.read()
+                resp = c.request(
+                    request.method,
+                    request.url.path,
+                    params=request.url.params,
+                    content=request.content,
+                    headers={
+                        k: v
+                        for k, v in request.headers.items()
+                        if k.lower() not in {"host", "content-length"}
+                    },
+                )
+                return httpx.Response(
+                    resp.status_code, headers=resp.headers, content=resp.content
+                )
 
         def make_fast_async_batch_processor(*args, **kwargs):
             kwargs.setdefault("min_batch_interval", 0)
@@ -733,14 +748,15 @@ def make_fast_call_batch_processor(*args, **kwargs):
             make_fast_call_batch_processor,
         )
 
+        # Absolute base URL required for httpx cookie handling; the transport
+        # routes by path, so the host is never dialed.
         remote_client = RemoteHTTPTraceServer(
-            trace_server_url="",
+            trace_server_url="http://testserver",
             should_batch=True,
+            transport=TestClientTransport(),
         )
         yield (client, remote_client, records)
 
-        weave.utils.http_requests.post = orig_post
-
 
 @pytest.fixture(autouse=True)
 def caching_client_isolation(monkeypatch, tmp_path):

tests/trace_server_bindings/conftest.py

@@ -2,6 +2,7 @@
 from types import MethodType
 from unittest.mock import MagicMock
 
+import httpx
 import pytest
 import tenacity
 
@@ -59,24 +60,60 @@ def generate_call_start_end_pair(
     return tsi.CallStartReq(start=start), tsi.CallEndReq(end=end)
 
 
+# =============================================================================
+# HTTP transport spy
+# =============================================================================
+
+
+class SpyTransport(httpx.BaseTransport):
+    """httpx transport that records requests and replays queued responses.
+
+    Queue items may be ``httpx.Response`` objects or exceptions to raise.
+    When the queue is empty, returns ``default_response`` (200 ``{}`` unless
+    overridden).
+    """
+
+    def __init__(
+        self,
+        *items: httpx.Response | Exception,
+        default_response: httpx.Response | None = None,
+    ) -> None:
+        self.queue: list[httpx.Response | Exception] = list(items)
+        self.requests: list[httpx.Request] = []
+        self.default_response = default_response
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        request.read()
+        self.requests.append(request)
+        if self.queue:
+            item = self.queue.pop(0)
+            if isinstance(item, Exception):
+                raise item
+            return item
+        if self.default_response is not None:
+            return self.default_response
+        return httpx.Response(200, json={})
+
+    @property
+    def urls(self) -> list[str]:
+        return [str(r.url) for r in self.requests]
+
+
 # =============================================================================
 # Fixtures
 # =============================================================================
 
 
 @pytest.fixture
-def success_response():
-    """Common fixture for mocking a successful HTTP response."""
-    response = MagicMock()
-    response.status_code = 200
-    response.json.return_value = {"id": "test_id", "trace_id": "test_trace_id"}
-    return response
+def server_class():
+    """The remote trace server implementation under test."""
+    return RemoteHTTPTraceServer
 
 
 @pytest.fixture
-def server(request):
-    """Common server fixture configured by the indirect parameter."""
-    server_ = RemoteHTTPTraceServer("http://example.com", should_batch=True)
+def server(request, server_class):
+    """Common server fixture parametrized by batching/retry behavior."""
+    server_ = server_class("http://example.com", should_batch=True)
 
     if request.param == "normal":
         server_._send_batch_to_server = MagicMock()