test: verify read_gbq_colab label preservation in anywidget display mode (#17887)

This PR ensures that `read_gbq_colab` telemetry labels are reliably
preserved across BigQuery job configurations when evaluating DataFrames
in `anywidget` display mode.

Fixes #<538249690> 🦕

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: shuoweil

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -68,6 +68,7 @@ def create_job_configs_labels(
         job_configs_labels = dict(job_configs_labels)
 
     if api_methods and "bigframes-api" not in job_configs_labels:
+        api_methods = list(api_methods)
         job_configs_labels["bigframes-api"] = api_methods[0]
         del api_methods[0]
 

packages/bigframes/tests/unit/session/test_io_bigquery.py

@@ -172,7 +172,7 @@ def test_add_and_trim_labels_length_limit_met():
     for _ in range(52):
         df.head()
 
-    io_bq.add_and_trim_labels(job_config=job_config)
+    io_bq.add_and_trim_labels(job_config=job_config, session=df._session)
     assert job_config.labels is not None
     assert len(job_config.labels) == 56
     assert "dataframe-max" not in job_config.labels.values()
@@ -218,6 +218,7 @@ def test_start_query_with_job_labels_length_limit_met(
         timeout=timeout,
         metrics=None,
         publisher=bigframes.core.events.Publisher(),
+        session=df._session,
     )
 
     assert job_config.labels is not None

packages/bigframes/tests/unit/session/test_read_gbq_colab.py

@@ -47,6 +47,36 @@ def test_read_gbq_colab_includes_label():
     assert "session-read_gbq_colab" in label_values
 
 
+def test_read_gbq_colab_includes_label_in_anywidget_mode():
+    """Make sure read_gbq_colab label is preserved in recent-bigframes-api labels in anywidget mode."""
+    pytest.importorskip("anywidget")
+    pytest.importorskip("traitlets")
+
+    import bigframes
+    import bigframes.display.html as bf_html
+
+    bqclient = mock.create_autospec(bigquery.Client, instance=True)
+    bqclient.project = "proj"
+    session = mocks.create_bigquery_session(bqclient=bqclient)
+    df = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")
+
+    with bigframes.option_context("display.render_mode", "anywidget"):
+        _ = bf_html.get_anywidget_bundle(df)
+
+    label_values = []
+    for kall in itertools.chain(
+        bqclient.query_and_wait.call_args_list,
+        bqclient._query_and_wait_bigframes.call_args_list,
+        bqclient.query.call_args_list,
+    ):
+        job_config = kall.kwargs.get("job_config")
+        if job_config is None:
+            continue
+        label_values.extend(job_config.labels.values())
+
+    assert "session-read_gbq_colab" in label_values
+
+
 @pytest.mark.parametrize("dry_run", [True, False])
 def test_read_gbq_colab_includes_formatted_values_in_dry_run(monkeypatch, dry_run):
     bqclient = mock.create_autospec(bigquery.Client, instance=True)