fix(weave): resolve ch_client off the loop and map cost columns like sync

acalls_query_stats passed self.ch_client as an argument to asyncio.to_thread,
so it was evaluated on the event loop: minting blocks there on first use, and
concurrent calls then handed one thread's client to every other thread.

acalls_query ignored include_costs. A cost query's SELECT carries sort-only
columns ahead of summary_dump, so zipping against select_fields alone shifts
every later value onto the wrong key -- silently, since the zip was strict=False.
Uses get_cost_result_columns like the sync path, and strict=True so a future
mismatch fails loudly.

Author: nikumar1206

tests/trace_server/test_async_calls_queries.py

@@ -0,0 +1,112 @@
+"""Guards for the two ways the async calls reads can silently diverge from sync.
+
+Both are failure modes that produce wrong data or wrong threading rather than an
+exception, so neither would show up in a smoke test.
+"""
+
+import asyncio
+import threading
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from weave.trace_server import trace_server_interface as tsi
+from weave.trace_server.async_clickhouse_trace_server import (
+    AsyncClickHouseTraceServer,
+)
+from weave.trace_server.calls_query_builder.calls_query_builder import (
+    CallsMergedField,
+    OrderField,
+)
+from weave.trace_server.token_costs import get_cost_result_columns
+
+PROJECT = "UHJvamVjdEludGVybmFsSWQ6MQ=="
+
+
+def test_acalls_query_stats_never_touches_ch_client_on_the_event_loop():
+    """`ch_client` is thread-local and minting one blocks.
+
+    Passed as an argument to `to_thread` it would be evaluated on the loop:
+    blocking it on first use, and then handing one thread's client to every
+    other thread running this concurrently.
+    """
+    server = AsyncClickHouseTraceServer(host="test_host")
+    loop_thread = threading.get_ident()
+    touched_on_loop: list[bool] = []
+
+    class Guard:
+        def __get__(self, obj, owner=None):
+            touched_on_loop.append(threading.get_ident() == loop_thread)
+            return MagicMock()
+
+    async def run():
+        with (
+            patch.object(type(server), "ch_client", Guard()),
+            patch.object(type(server), "table_routing_resolver", MagicMock()),
+            patch.object(server, "_aquery", return_value=MagicMock(result_rows=[[0]])),
+            patch(
+                "weave.trace_server.async_clickhouse_trace_server.build_calls_stats_query",
+                return_value=("SELECT 1", ["count"], None),
+            ),
+            patch(
+                "weave.trace_server.async_clickhouse_trace_server.calls_stats_res",
+                return_value=tsi.CallsQueryStatsRes(count=0),
+            ),
+        ):
+            await server.acalls_query_stats(tsi.CallsQueryStatsReq(project_id=PROJECT))
+
+    asyncio.run(run())
+    assert touched_on_loop, "ch_client was never resolved; the test proves nothing"
+    assert not any(touched_on_loop), (
+        "ch_client was resolved on the event loop thread; it must be resolved "
+        "inside the executor callback"
+    )
+
+
+class _Captured(Exception):
+    """Carries the zip keys out before any downstream validation runs."""
+
+    def __init__(self, columns):
+        self.columns = columns
+
+
+@pytest.mark.parametrize("include_costs", [True, False])
+def test_acalls_query_maps_columns_the_way_sync_does(include_costs: bool):
+    """A cost query's SELECT carries sort-only columns ahead of `summary_dump`.
+
+    Zipping row values against `select_fields` alone shifts every value after
+    that point onto the wrong key, which is silent rather than an error.
+    """
+    server = AsyncClickHouseTraceServer(host="test_host")
+    select = ["id", "project_id", "summary_dump"]
+    order = [OrderField(field=CallsMergedField(field="started_at"), direction="ASC")]
+    cq = MagicMock(
+        select_fields=[MagicMock(field=f) for f in select], order_fields=order
+    )
+    cq.as_sql.return_value = "SELECT 1"
+
+    expected = get_cost_result_columns(select, order) if include_costs else select
+
+    def capture(d):
+        raise _Captured(list(d))
+
+    async def run():
+        with (
+            patch.object(server, "_build_calls_query", return_value=(cq, None)),
+            patch.object(
+                server,
+                "_aquery",
+                return_value=MagicMock(result_rows=[list(range(len(expected)))]),
+            ),
+            patch(
+                "weave.trace_server.async_clickhouse_trace_server.ch_call_dict_to_call_schema_dict",
+                side_effect=capture,
+            ),
+        ):
+            await server.acalls_query(
+                tsi.CallsQueryReq(project_id=PROJECT, include_costs=include_costs)
+            )
+
+    with pytest.raises(_Captured) as exc:
+        asyncio.run(run())
+    assert exc.value.columns == expected

weave/trace_server/async_clickhouse_trace_server.py

@@ -56,6 +56,7 @@
     make_spans_count_query,
     make_spans_list_query,
 )
+from weave.trace_server.token_costs import get_cost_result_columns
 from weave.trace_server.tracing import traced
 
 _T = TypeVar("_T")
@@ -284,11 +285,17 @@ async def acalls_query(self, req: tsi.CallsQueryReq) -> tsi.CallsQueryRes:
         cq, settings = await asyncio.to_thread(self._build_calls_query, req)
         pb = ParamBuilder()
         raw_res = await self._aquery(cq.as_sql(pb), pb.get_params(), settings=settings)
-        select_columns = [c.field for c in cq.select_fields]
+        if req.include_costs:
+            # Cost query SELECT adds ORDER BY fields; result columns must match.
+            select_columns = get_cost_result_columns(
+                [c.field for c in cq.select_fields], cq.order_fields
+            )
+        else:
+            select_columns = [c.field for c in cq.select_fields]
         calls = [
             tsi.CallSchema.model_validate(
                 ch_call_dict_to_call_schema_dict(
-                    dict(zip(select_columns, row, strict=False))
+                    dict(zip(select_columns, row, strict=True))
                 )
             )
             for row in raw_res.result_rows
@@ -300,10 +307,14 @@ async def acalls_query_stats(
         self, req: tsi.CallsQueryStatsReq
     ) -> tsi.CallsQueryStatsRes:
         """Native-async twin of `calls_query_stats`."""
-        read_table = await asyncio.to_thread(
-            self.table_routing_resolver.resolve_read_table,
-            req.project_id,
-            self.ch_client,
+        # `ch_client` is thread-local and minting one blocks. Resolve it inside
+        # the pool thread: as an argument it would be evaluated on the event
+        # loop, blocking it on first use and then handing one thread's client to
+        # every other thread that runs this concurrently.
+        read_table = await self._run_on_ch_executor(
+            lambda: self.table_routing_resolver.resolve_read_table(
+                req.project_id, self.ch_client
+            )
         )
         pb = ParamBuilder()
         query, columns, settings = build_calls_stats_query(req, pb, read_table)