feat: cap rows returned to the agent context

app/agent/tools/bigquery.py

@@ -1,4 +1,5 @@
 import inspect
+import itertools
 import json
 import uuid
 from functools import cache
@@ -15,6 +16,11 @@
 
 MAX_BYTES_BILLED = 10 * 10**9
 
+# Cap on how many rows are serialized into the agent's context. The full result set
+# is still materialized in BigQuery's destination table, so downloads (via query_ref)
+# get every row regardless — this only keeps a large result from blowing up context.
+MAX_CONTEXT_ROWS = 1000
+
 
 @cache
 def _bq_client() -> bq.Client:  # pragma: no cover
@@ -83,7 +89,9 @@ def execute_bigquery_sql(
                 labels=labels,
             ),
         )
-        rows = [dict(row) for row in job.result()]
+        result = job.result()
+        total_rows = result.total_rows
+        rows = [dict(row) for row in itertools.islice(result, MAX_CONTEXT_ROWS)]
     except GoogleAPICallError as e:
         reason = e.errors[0].get("reason") if getattr(e, "errors", None) else None
         if reason == "bytesBilledLimitExceeded":
@@ -102,9 +110,18 @@ def execute_bigquery_sql(
     # (~24h TTL), so a later export hands back exactly these rows without re-running.
     query_ref = f"qr_{uuid.uuid4().hex}"
 
-    content = json.dumps(
-        {"row_count": len(rows), "rows": rows}, ensure_ascii=False, default=str
-    )
+    payload = {"row_count": total_rows, "rows": rows}
+
+    # Surface truncation only when it actually happened, so the agent knows the rows
+    # it sees are a subset — and that the full set is still available for download.
+    if total_rows > len(rows):
+        payload["truncated"] = True
+        payload["truncation_note"] = (
+            f"Only the first {len(rows)} of {total_rows} rows are shown here to keep "
+            "the context small. The full result can still be downloaded from the interface."
+        )
+
+    content = json.dumps(payload, ensure_ascii=False, default=str)
 
     artifact = {
         "type": "query_result",

tests/app/agent/tools/test_bigquery.py

@@ -12,11 +12,21 @@
 from app.agent.context import AgentContext
 from app.agent.tools.bigquery import (
     MAX_BYTES_BILLED,
+    MAX_CONTEXT_ROWS,
     decode_table_values,
     execute_bigquery_sql,
 )
 
 
+def _mock_result(rows: list[dict], total_rows: int | None = None) -> MagicMock:
+    """Stand in for BigQuery's RowIterator: iterable over `rows`, and carrying a
+    `total_rows` count (the full result size, which may exceed the fetched rows)."""
+    result = MagicMock()
+    result.__iter__.return_value = iter(rows)
+    result.total_rows = len(rows) if total_rows is None else total_rows
+    return result
+
+
 @pytest.fixture
 def mock_context() -> AgentContext:
     """The run context the agent injects into a tool
@@ -69,7 +79,9 @@ def test_successful_query(self, mocker: MockerFixture, mock_context: AgentContex
         mock_dry_run_query_job.statement_type = "SELECT"
 
         mock_query_job = MagicMock()
-        mock_query_job.result.return_value = [{"col1": "value1"}, {"col1": "value2"}]
+        mock_query_job.result.return_value = _mock_result(
+            [{"col1": "value1"}, {"col1": "value2"}]
+        )
         mock_query_job.destination.to_api_repr.return_value = {
             "projectId": "p",
             "datasetId": "d",
@@ -106,7 +118,7 @@ def test_successful_query_exposes_destination_table_on_artifact(
         mock_dry_run_query_job.statement_type = "SELECT"
 
         mock_query_job = MagicMock()
-        mock_query_job.result.return_value = [{"col1": "value1"}]
+        mock_query_job.result.return_value = _mock_result([{"col1": "value1"}])
         mock_query_job.destination.to_api_repr.return_value = {
             "projectId": "p",
             "datasetId": "d",
@@ -146,7 +158,7 @@ def test_successful_query_empty_result(
         mock_dry_run_query_job.statement_type = "SELECT"
 
         mock_query_job = MagicMock()
-        mock_query_job.result.return_value = []
+        mock_query_job.result.return_value = _mock_result([])
 
         mock_bigquery_client = MagicMock(spec=bq.Client)
         mock_bigquery_client.query.side_effect = [
@@ -170,6 +182,91 @@ def test_successful_query_empty_result(
         )
         assert message.artifact is None
 
+    def test_large_result_is_truncated_for_context(
+        self, mocker: MockerFixture, mock_context: AgentContext
+    ):
+        """A result larger than the cap only serializes a prefix, flags the truncation,
+        and still mints a download handle over the full (materialized) result."""
+        total = MAX_CONTEXT_ROWS + 500
+        all_rows = [{"n": i} for i in range(total)]
+
+        mock_dry_run_query_job = MagicMock()
+        mock_dry_run_query_job.statement_type = "SELECT"
+
+        mock_query_job = MagicMock()
+        mock_query_job.result.return_value = _mock_result(all_rows, total_rows=total)
+        mock_query_job.destination.to_api_repr.return_value = {
+            "projectId": "p",
+            "datasetId": "d",
+            "tableId": "t",
+        }
+
+        mock_bigquery_client = MagicMock(spec=bq.Client)
+        mock_bigquery_client.query.side_effect = [
+            mock_dry_run_query_job,
+            mock_query_job,
+        ]
+
+        mocker.patch(
+            "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client
+        )
+
+        message = _invoke_tool(
+            execute_bigquery_sql,
+            {"sql_query": "SELECT * FROM project.dataset.table", "slug": "resultado"},
+            context=mock_context,
+        )
+
+        output = json.loads(message.content)
+
+        assert output["row_count"] == total
+        assert len(output["rows"]) == MAX_CONTEXT_ROWS
+        assert output["rows"][0] == {"n": 0}
+        assert output["truncated"] is True
+        assert str(total) in output["truncation_note"]
+        # Full result is still downloadable.
+        assert re.fullmatch(r"qr_[0-9a-f]{32}", message.artifact["query_ref"])
+
+    def test_result_at_cap_is_not_flagged_truncated(
+        self, mocker: MockerFixture, mock_context: AgentContext
+    ):
+        """A result exactly at the cap returns every row and no truncation flag."""
+        all_rows = [{"n": i} for i in range(MAX_CONTEXT_ROWS)]
+
+        mock_dry_run_query_job = MagicMock()
+        mock_dry_run_query_job.statement_type = "SELECT"
+
+        mock_query_job = MagicMock()
+        mock_query_job.result.return_value = _mock_result(all_rows)
+        mock_query_job.destination.to_api_repr.return_value = {
+            "projectId": "p",
+            "datasetId": "d",
+            "tableId": "t",
+        }
+
+        mock_bigquery_client = MagicMock(spec=bq.Client)
+        mock_bigquery_client.query.side_effect = [
+            mock_dry_run_query_job,
+            mock_query_job,
+        ]
+
+        mocker.patch(
+            "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client
+        )
+
+        message = _invoke_tool(
+            execute_bigquery_sql,
+            {"sql_query": "SELECT * FROM project.dataset.table", "slug": "resultado"},
+            context=mock_context,
+        )
+
+        output = json.loads(message.content)
+
+        assert output["row_count"] == MAX_CONTEXT_ROWS
+        assert len(output["rows"]) == MAX_CONTEXT_ROWS
+        assert "truncated" not in output
+        assert "truncation_note" not in output
+
     def test_forbidden_statement_type(
         self, mocker: MockerFixture, mock_context: AgentContext
     ):