From 5cd8b59afc6bd50cbea8ebfb5aa57e79e2946481 Mon Sep 17 00:00:00 2001
From: andrewtruong <itsandrewtruong@gmail.com>
Date: Tue, 27 Jan 2026 20:08:15 -0500
Subject: [PATCH 1/2] test

---
 tests/trace/test_call_object.py  |  5 +++++
 tests/trace/test_weave_client.py | 31 ++++++++++++++++++++++++++++---
 weave/trace/call.py              | 10 ++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/tests/trace/test_call_object.py b/tests/trace/test_call_object.py
index f975feef68b1..7bae82b30429 100644
--- a/tests/trace/test_call_object.py
+++ b/tests/trace/test_call_object.py
@@ -24,4 +24,9 @@ def greet(name: str, age: int) -> str:
         "thread_id": call.thread_id,
         "turn_id": call.turn_id,
         "project_id": call.project_id,
+        "wb_run_id": call.wb_run_id,
+        "wb_run_step": call.wb_run_step,
+        "wb_run_step_end": call.wb_run_step_end,
+        "storage_size_bytes": call.storage_size_bytes,
+        "total_storage_size_bytes": call.total_storage_size_bytes,
     }
diff --git a/tests/trace/test_weave_client.py b/tests/trace/test_weave_client.py
index 45cc7030df35..ee36fa510fed 100644
--- a/tests/trace/test_weave_client.py
+++ b/tests/trace/test_weave_client.py
@@ -32,6 +32,7 @@
 from weave.trace.context.call_context import tracing_disabled
 from weave.trace.isinstance import weave_isinstance
 from weave.trace.op import is_op
+from weave.trace.wandb_run_context import WandbRunContext
 from weave.trace.refs import (
     DICT_KEY_EDGE_NAME,
     LIST_INDEX_EDGE_NAME,
@@ -4054,9 +4055,6 @@ def test_table_create_from_digests(network_proxy_client):
 def test_calls_query_with_wb_run_id_not_null(client, monkeypatch):
     """Test optimized stats query for wb_run_id not null."""
     # Mock wandb to simulate a run
-    from weave.trace import weave_client
-    from weave.trace.wandb_run_context import WandbRunContext
-
     mock_run_id = f"{client._project_id()}/test_run_123"
     monkeypatch.setattr(
         weave_client,
@@ -4079,6 +4077,33 @@ def test_op(x: int) -> int:
     assert len(calls) == 1
     assert calls[0].wb_run_id == mock_run_id
 
+
+def test_get_calls_columns_wb_run_id(client, monkeypatch):
+    mock_run_id = f"{client._project_id()}/test_run_456"
+    monkeypatch.setattr(
+        weave_client,
+        "get_global_wb_run_context",
+        lambda: WandbRunContext(run_id="test_run_456", step=7),
+    )
+
+    @weave.op
+    def test_op(x: int) -> int:
+        return x * 3
+
+    _, call = test_op.call(2)
+    client.flush()
+
+    calls = list(
+        client.get_calls(
+            columns=["wb_run_id"],
+            filter=tsi.CallsFilter(call_ids=[call.id]),
+        )
+    )
+
+    assert len(calls) == 1
+    assert hasattr(calls[0], "wb_run_id")
+    assert calls[0].wb_run_id == mock_run_id
+
     # Now query for calls with wb_run_id not null using limit=1 to trigger optimization
     query = tsi.Query(
         **{
diff --git a/weave/trace/call.py b/weave/trace/call.py
index 32c6bf7e58ef..224a427ced76 100644
--- a/weave/trace/call.py
+++ b/weave/trace/call.py
@@ -286,6 +286,11 @@ def to_dict(self) -> CallDict:
             deleted_at=self.deleted_at,
             thread_id=self.thread_id,
             turn_id=self.turn_id,
+            wb_run_id=self.wb_run_id,
+            wb_run_step=self.wb_run_step,
+            wb_run_step_end=self.wb_run_step_end,
+            storage_size_bytes=self.storage_size_bytes,
+            total_storage_size_bytes=self.total_storage_size_bytes,
         )
 
 
@@ -313,6 +318,11 @@ class CallDict(TypedDict):
     deleted_at: datetime.datetime | None
     thread_id: str | None
     turn_id: str | None
+    wb_run_id: str | None
+    wb_run_step: int | None
+    wb_run_step_end: int | None
+    storage_size_bytes: int | None
+    total_storage_size_bytes: int | None
 
 
 CallsIter = PaginatedIterator[CallSchema, WeaveObject]

From 655d5ee8177a617f6876f4533fb3dd593f1034ba Mon Sep 17 00:00:00 2001
From: andrewtruong <itsandrewtruong@gmail.com>
Date: Tue, 27 Jan 2026 23:55:51 -0500
Subject: [PATCH 2/2] test

---
 tests/trace/test_weave_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/trace/test_weave_client.py b/tests/trace/test_weave_client.py
index ee36fa510fed..2de35c1c3b03 100644
--- a/tests/trace/test_weave_client.py
+++ b/tests/trace/test_weave_client.py
@@ -32,7 +32,6 @@
 from weave.trace.context.call_context import tracing_disabled
 from weave.trace.isinstance import weave_isinstance
 from weave.trace.op import is_op
-from weave.trace.wandb_run_context import WandbRunContext
 from weave.trace.refs import (
     DICT_KEY_EDGE_NAME,
     LIST_INDEX_EDGE_NAME,
@@ -45,6 +44,7 @@
     get_serializer_for_obj,
     register_serializer,
 )
+from weave.trace.wandb_run_context import WandbRunContext
 from weave.trace_server.clickhouse_trace_server_batched import NotFoundError
 from weave.trace_server.common_interface import SortBy
 from weave.trace_server.constants import MAX_DISPLAY_NAME_LENGTH