diff --git a/contributing/samples/integrations/snowflake_cortex_agent/README.md b/contributing/samples/integrations/snowflake_cortex_agent/README.md new file mode 100644 index 0000000000..af1870db08 --- /dev/null +++ b/contributing/samples/integrations/snowflake_cortex_agent/README.md @@ -0,0 +1,113 @@ +# Snowflake Cortex Analyst Agent + +## Overview + +This sample runs an existing [Snowflake Cortex Agent](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents) +as a native ADK root agent using `SnowflakeCortexAgent`. Snowflake runs the +agent loop and its server-side tools (Cortex Analyst, Cortex Search, SQL +execution); each ADK turn is sent to the Cortex Agents Run API and the run comes +back as standard ADK events: streamed text, the tool calls Snowflake made, and +one final answer whose citations, warnings, tables and charts are recorded as +event metadata. The Snowflake thread continues across turns through ADK session +state. + +`SnowflakeCortexAgent` is experimental and lives under `google.adk.labs`. See +the +[SnowflakeCortexAgent guide](../../../../docs/guides/labs/snowflake/snowflake_cortex_agent/index.md) +for the full setup, limitations, and API details. + +## Prerequisites + +- A Cortex Agent object in Snowflake that the token below may run. +- A Snowflake token for the REST API: a programmatic access token, an OAuth + access token, or a key-pair JWT. +- Environment variables, in the shell or in a `.env` file next to `agent.py` + (`adk web` and `adk run` load it): + +```text +SNOWFLAKE_ACCOUNT_URL=https://.snowflakecomputing.com +SNOWFLAKE_DATABASE=SALES_DB +SNOWFLAKE_SCHEMA=ANALYTICS +SNOWFLAKE_CORTEX_AGENT=SALES_AGENT +SNOWFLAKE_TOKEN= +SNOWFLAKE_TOKEN_TYPE=PROGRAMMATIC_ACCESS_TOKEN +``` + +`SNOWFLAKE_TOKEN_TYPE` is `PROGRAMMATIC_ACCESS_TOKEN`, `OAUTH`, `KEYPAIR_JWT` or +`WORKLOAD_IDENTITY_FEDERATION`, matching the token you supply. + +No extra package is needed: the integration uses the `httpx` client ADK already +depends on. + +## Sample Inputs + +- `What were total sales by region last quarter?` + + The Cortex Agent picks a semantic view, generates and runs SQL, and answers + from the result. The SQL tool call and its result appear as tool events, the + answer as the final event. + +- `Show only the mobile channel.` + + A follow-up in the same session continues the Snowflake thread, so the Cortex + Agent has the previous question and answer as context. + +- `Which product categories exist in the data?` + + A question the Cortex Agent can answer from the semantic model alone. + +## Graph + +The ADK agent fronts one Cortex Agent object, which owns its own tools: + +```mermaid +graph LR + ADK[snowflake_cortex_analyst
SnowflakeCortexAgent] -->|Run API| Cortex[SALES_AGENT
Cortex Agent object] + Cortex --> Analyst(Cortex Analyst) + Cortex --> Search(Cortex Search) + Cortex --> SQL(SQL execution) +``` + +## How To + +Point the agent at the Snowflake object and supply credentials through a header +provider: + +```python +root_agent = SnowflakeCortexAgent( + name="snowflake_cortex_analyst", + description="Answers data questions by running a Snowflake Cortex Agent.", + account_url=_env("SNOWFLAKE_ACCOUNT_URL"), + database=_env("SNOWFLAKE_DATABASE"), + schema_name=_env("SNOWFLAKE_SCHEMA"), + cortex_agent_name=_env("SNOWFLAKE_CORTEX_AGENT"), + header_provider=snowflake_headers, +) +``` + +The header provider is a plain function that receives the invocation's +`ReadonlyContext` and returns the HTTP headers for one Snowflake request. This +sample reads one service token from the environment on every call, so a rotated +token is picked up without a restart. Missing settings are reported on the first +request rather than at import, so the sample can be listed by `adk web` before +it is configured: + +```python +def snowflake_headers(ctx: ReadonlyContext) -> dict[str, str]: + _check_configured() # fails the first request, not the import + return { + "Authorization": f"Bearer {_env('SNOWFLAKE_TOKEN')}", + "X-Snowflake-Authorization-Token-Type": os.environ.get( + "SNOWFLAKE_TOKEN_TYPE", "PROGRAMMATIC_ACCESS_TOKEN" + ), + } +``` + +Run it with `adk web contributing/samples/integrations` and pick +`snowflake_cortex_agent`, or with `adk run`. Use SSE streaming to see the text +and reasoning deltas as they arrive; without it only the tool events and the +final answer are yielded. + +## Related Guides + +- [SnowflakeCortexAgent](../../../../docs/guides/labs/snowflake/snowflake_cortex_agent/index.md) - Setup, event mapping, thread continuity, security and limitations of the Snowflake Cortex Agent integration. diff --git a/contributing/samples/integrations/snowflake_cortex_agent/agent.py b/contributing/samples/integrations/snowflake_cortex_agent/agent.py new file mode 100644 index 0000000000..3d1eb6bb81 --- /dev/null +++ b/contributing/samples/integrations/snowflake_cortex_agent/agent.py @@ -0,0 +1,76 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Analytics assistant that runs a Snowflake Cortex Agent as an ADK root agent. + +Wraps an existing Cortex Agent object with `SnowflakeCortexAgent`. See the +guide at docs/guides/labs/snowflake/snowflake_cortex_agent/index.md for setup +and details. +""" + +import os + +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.labs.snowflake import SnowflakeCortexAgent + +_REQUIRED_ENV = ( + "SNOWFLAKE_ACCOUNT_URL", + "SNOWFLAKE_DATABASE", + "SNOWFLAKE_SCHEMA", + "SNOWFLAKE_CORTEX_AGENT", + "SNOWFLAKE_TOKEN", +) + + +def _env(name: str, default: str = "") -> str: + return os.environ.get(name, default).strip() + + +def _check_configured() -> None: + # Checked when the first request is made, not at import: `adk web` and the + # sample tests import every sample, with or without Snowflake credentials. + missing = [name for name in _REQUIRED_ENV if not _env(name)] + if missing: + raise RuntimeError( + "Snowflake settings are missing:" + f" {', '.join(missing)}. Set them in the environment or in a .env" + " file next to agent.py." + ) + + +def snowflake_headers(ctx: ReadonlyContext) -> dict[str, str]: + """Reads the token on every request so a rotated token is picked up.""" + del ctx # One service token for every user; see the guide for per-user auth. + _check_configured() + return { + "Authorization": f"Bearer {_env('SNOWFLAKE_TOKEN')}", + "X-Snowflake-Authorization-Token-Type": _env( + "SNOWFLAKE_TOKEN_TYPE", "PROGRAMMATIC_ACCESS_TOKEN" + ), + } + + +# 1. Point at the Cortex Agent object that already exists in Snowflake. The +# ADK name is separate from the Snowflake object name. +# 2. Credentials come from `header_provider`, never from a field, so they stay +# out of `repr`, the adk web agent graph and the session store. +root_agent = SnowflakeCortexAgent( + name="snowflake_cortex_analyst", + description="Answers data questions by running a Snowflake Cortex Agent.", + account_url=_env("SNOWFLAKE_ACCOUNT_URL"), + database=_env("SNOWFLAKE_DATABASE"), + schema_name=_env("SNOWFLAKE_SCHEMA"), + cortex_agent_name=_env("SNOWFLAKE_CORTEX_AGENT"), + header_provider=snowflake_headers, +) diff --git a/docs/guides/README.md b/docs/guides/README.md index 454675fda7..f8d406aae6 100644 --- a/docs/guides/README.md +++ b/docs/guides/README.md @@ -35,6 +35,7 @@ This directory contains specific developer guides for the ADK Python implementat ### Labs * [AntigravityAgent](labs/antigravity/index.md) - Runs a Google Antigravity SDK agent as an ADK agent node. +* [SnowflakeCortexAgent](labs/snowflake/snowflake_cortex_agent/index.md) - Runs a Snowflake Cortex Agent as an ADK root agent, streaming its run as ADK events. ### Live * [LiveRequestQueue](live/live_request_queue/index.md) - Streaming content, realtime audio, and stream control signals to live agents. diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md new file mode 100644 index 0000000000..1987145727 --- /dev/null +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -0,0 +1,361 @@ +# SnowflakeCortexAgent + +Runs an existing Snowflake Cortex Agent as an ADK root agent. Each turn goes to +the Cortex Agents Run API and comes back as ADK events: streamed text, the +server-side tool calls Snowflake made, and one final answer that carries +citations, warnings, tables and charts as metadata. + +## Introduction + +`SnowflakeCortexAgent` is a `BaseAgent` for a Cortex Agent object that already +exists in Snowflake. Snowflake owns the agent loop, the tools it can call +(Cortex Analyst, Cortex Search, SQL execution) and the conversation thread; ADK +owns the session, the events and whatever consumes them, such as `adk web` or +your own `Runner` loop. + +The integration exists because the other way to reach a Cortex Agent from ADK, +Snowflake's Managed MCP server through `McpToolset`, returns a whole run as one +tool result. That path cannot stream the answer as it is written, cannot show +which tools ran and with what, and has no place to keep the Snowflake thread +between turns. `SnowflakeCortexAgent` calls the REST API directly with the +`httpx` client ADK already depends on, so no extra package is required. + +The class lives under `google.adk.labs`, which means its API can change between +releases while it is experimental. + +## Get started + +The agent needs the location of the Cortex Agent object and a way to +authenticate. Credentials are supplied by a header provider, a function that +returns the HTTP headers for one request, so a token is never stored on the +agent: + +```python +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.labs.snowflake import SnowflakeCortexAgent + + +def snowflake_headers(ctx: ReadonlyContext) -> dict[str, str]: + return { + "Authorization": f"Bearer {load_snowflake_token()}", + "X-Snowflake-Authorization-Token-Type": "PROGRAMMATIC_ACCESS_TOKEN", + } + + +root_agent = SnowflakeCortexAgent( + name="sales_analyst", + account_url="https://.snowflakecomputing.com", + database="SALES_DB", + schema_name="ANALYTICS", + cortex_agent_name="SALES_AGENT", + header_provider=snowflake_headers, +) +``` + +`load_snowflake_token` stands for however your application obtains a Snowflake +token: a programmatic access token from a secret store, an OAuth access token, +or a key-pair JWT it mints itself. The `X-Snowflake-Authorization-Token-Type` +header names the kind of token; Snowflake accepts `PROGRAMMATIC_ACCESS_TOKEN`, +`OAUTH`, `KEYPAIR_JWT` and `WORKLOAD_IDENTITY_FEDERATION`. + +Run the agent as you would any other root agent, through `adk web`, `adk run` +or a `Runner`. Ask for SSE streaming in the `RunConfig` to receive text and +reasoning as it is produced. + +## How it works + +On each turn the agent reads a cursor from ADK session state under a key +scoped to its own `name`. On the first turn there is none, so it creates a +Snowflake thread and starts from message `0`. It then sends only the current +user message to the Run API, because Snowflake holds the earlier turns in the +thread; ADK session history is not re-sent. + +The run arrives as a stream of typed Cortex events and the agent maps them onto +ADK events: + +- Text and reasoning deltas become `partial=True` events, reasoning as + `thought` parts, and only when the `RunConfig` asks for SSE streaming. Partial + events are never persisted, so a consumer that did not ask for streaming + receives nothing it cannot use. +- Progress notices, tool status updates, warnings and event types the agent + does not know are forwarded the same way, as partial events with the Cortex + payload under `custom_metadata["snowflake_cortex"]`. +- Each server-side tool Snowflake ran is recorded as a `FunctionCall` event + authored by the agent, followed by a `FunctionResponse` event authored by the + tool name, paired by Snowflake's `tool_use_id`. These are real ADK tool + events: `adk web` renders them as tool calls, and ADK does not execute them + again, because tool execution only happens for calls an `LlmAgent` model + makes. Tool results larger than `max_tool_result_bytes` are reduced in + stages: first the SQL rows are dropped while the query id and column + metadata stay, then each content block is reduced to its type and the size + of each of its keys, so the record still says what came back. +- When Snowflake sends its final `response`, the + agent yields one non-partial event with the answer text. Citations, + warnings, suggested follow-up queries, token usage, tables and charts sit in + `custom_metadata["snowflake_cortex"]` on that event. The final `response` + is authoritative, so the answer comes from it rather than from the deltas. + +The final event also carries the new cursor, the thread id and the assistant +message id, as `state_delta`, but only when the final `response` reports +`status` `completed`; a `cancelled` or `timed_out` run leaves the cursor where +it was. The `[DONE]` terminator is optional: the final `response` is what +closes the turn. The ADK `Runner` applies it to the session, which +is why thread continuity needs the `Runner` rather than a direct +`run_async` loop. Only an assistant message id ever becomes the next parent; +using the user message id would fork the thread. The cursor also records the +account and object it belongs to. Pointing an existing session at a different +Cortex Agent raises a `ValueError` naming the state key to remove, rather than +silently mixing two conversations. + +Failures leave the cursor alone. A terminal `error` from Snowflake becomes a +non-partial event with `error_code` and `error_message`. A stream that ends +before the final `response` arrived raises `CortexTransportError`, and a +rejected request +raises `CortexApiError` with the HTTP status, Snowflake's error code and its +request id. The next turn continues from the last good message. + +If the consumer stops reading, for example when a browser tab closes, the +agent closes the connection to Snowflake and, with `cancel_on_disconnect` +enabled, asks Snowflake once to cancel the run. In live tests the cancel request was answered with +`409 Agent run was already completed`, both right after closing the stream +and for a run that had finished on its own, so the agent treats that answer +as a benign outcome. Snowflake keeps whatever it had already produced in the +thread either way. + +## Configuration options + +| Option | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `account_url` | `str` | (Required) | Base URL of the Snowflake account. | +| `database` | `str` | (Required) | Database that holds the Cortex Agent object. | +| `schema_name` | `str` | (Required) | Schema that holds the Cortex Agent object. | +| `cortex_agent_name` | `str` | (Required) | Name of the Cortex Agent object. | +| `header_provider` | `Callable[[ReadonlyContext], dict[str, str] \| Awaitable[dict[str, str]]]` | (Required) | Returns the HTTP headers for each Snowflake request. | +| `http_client` | `httpx.AsyncClient \| None` | `None` | A shared HTTP client to send requests through. | +| `timeout` | `float` | `900.0` | Seconds to wait on Snowflake before the turn fails. | +| `cancel_on_disconnect` | `bool` | `True` | Cancel the Snowflake run when the consumer stops reading. | +| `max_tool_result_bytes` | `int` | `32768` | Size bound for one recorded tool result, table or chart. | +| `include_thinking_in_final_event` | `bool` | `False` | Also persist the completed reasoning on the final event. | + +`account_url`, `database`, `schema_name` and `cortex_agent_name` locate the +Cortex Agent object. The field is `schema_name` rather than `schema` because +pydantic reserves that name. The ADK `name` of the agent is separate from +`cortex_agent_name`, so two ADK agents can front the same Snowflake object. + +`header_provider` is called with the invocation's `ReadonlyContext` before +every request and may be a plain function or a coroutine function. It must +return the `Authorization` header and, for tokens other than OAuth, the matching +`X-Snowflake-Authorization-Token-Type`. The agent adds content negotiation +headers itself. The provider is excluded from `repr`, `model_dump` and the +`adk web` agent graph, so a token bound into it does not leak through those +paths. + +`http_client` lets several agents pool connections, or lets you configure a +proxy or a certificate bundle once. When you pass one, you own it and close +it; when you leave it unset, the agent creates a client on first use and closes +it in `cleanup()`. + +`timeout` bounds how long the agent waits for Snowflake: to connect, and during +a run, between two chunks of the stream. Cortex Agent runs that plan, execute +SQL and summarize can take minutes, which is why the default is long. Lower it +when your Cortex Agent answers quickly and you would rather fail fast. + +`cancel_on_disconnect` decides whether the agent also sends a cancel request +for the run when the ADK consumer goes away mid-turn. The cancel is best +effort and its failure is not raised, because nobody is left to read the +error; a refusal is logged at debug level with its HTTP status. In live tests the +cancel was answered with 409 (run already completed) right after the stream +was closed, so treat the request as a safeguard rather than as the mechanism +that stops the run. Disable it to skip that extra request. + +`max_tool_result_bytes` protects the session store. Tool results are persisted +as `FunctionResponse` events and can be large, so a result above the bound is +cut down in stages: a SQL result keeps its query id and column metadata and +drops the rows; a result that is still too large, such as a semantic view +context, keeps each block's type and the byte size of each of its keys; only +if even that does not fit is the content emptied. `truncated` and +`original_bytes` on the response say when this happened. Tables and charts on +the final event are bounded the same way. Raise it if your application reads +rows or tool payloads from the recorded events, lower it if session size +matters more. + +`include_thinking_in_final_event` controls whether the completed reasoning text +is written into the persisted final event as a `thought` part. It is off by +default so reasoning does not reach the session store; streamed reasoning +deltas are unaffected because partial events are never persisted. + +## Advanced applications + +### Reading the run metadata + +The final event carries everything Snowflake reported about the run besides the +answer text. This example collects the final answer and its suggested follow-up +queries from a `Runner` loop: + +```python +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types + +session_service = InMemorySessionService() +runner = Runner( + app_name="analytics", agent=root_agent, session_service=session_service +) + + +async def ask(session_id: str, question: str) -> tuple[str, list[dict]]: + answer = "" + suggested: list[dict] = [] + async for event in runner.run_async( + user_id="analyst", + session_id=session_id, + new_message=types.Content( + role="user", parts=[types.Part.from_text(text=question)] + ), + ): + if event.partial or not event.is_final_response(): + continue + cortex = (event.custom_metadata or {}).get("snowflake_cortex", {}) + answer = "".join(part.text or "" for part in event.content.parts) + suggested = cortex.get("suggested_queries", []) + return answer, suggested +``` + +The same dictionary holds `annotations` for citations, `warnings`, `tables`, +`charts`, `usage` and the Snowflake `run_id`. Entries in `suggested_queries` +with `source` set to `cortex_analyst` are the questions Cortex Analyst proposed +when it could not answer the one it was given; the others come from the Cortex +Agent itself. + +### Authenticating as the end user + +Because the header provider receives the `ReadonlyContext`, it can mint a +token for the user of the invocation instead of using one service token. This +example looks up an OAuth access token by ADK user id in a store the +application maintains: + +```python +from google.adk.agents.readonly_context import ReadonlyContext + + +async def per_user_headers(ctx: ReadonlyContext) -> dict[str, str]: + token = await token_store.access_token_for(ctx.user_id) + return { + "Authorization": f"Bearer {token}", + "X-Snowflake-Authorization-Token-Type": "OAUTH", + } +``` + +The token itself should not be written into session state: state is persisted +by the session service and visible to anything that reads the session. + +### Sharing an HTTP client + +Pass an `httpx.AsyncClient` to pool connections or to route through a proxy. +The agent does not close a client it was given: + +```python +import httpx + +http_client = httpx.AsyncClient(proxy="http://proxy.internal:3128") + +root_agent = SnowflakeCortexAgent( + name="sales_analyst", + account_url="https://.snowflakecomputing.com", + database="SALES_DB", + schema_name="ANALYTICS", + cortex_agent_name="SALES_AGENT", + header_provider=snowflake_headers, + http_client=http_client, +) +``` + +Close the client yourself when the application shuts down. + +## Security + +The events this agent yields are recorded by the session service, returned by +session APIs and forwarded to memory services like any other ADK events. Two +of them deserve attention: + +- A `FunctionCall` event for `system_execute_sql` carries the generated SQL in + its `args`, and the matching `FunctionResponse` carries the result rows up to + `max_tool_result_bytes`. This mirrors how `AntigravityAgent` records tool + calls an external runtime made, and it is what lets `adk web` show the tool + trace. +- The final event carries tables and charts up to the same bound. + +If your deployment must not persist generated SQL, install a plugin that +redacts it before the event is stored. The `Runner` gives plugins the event +before persisting it and uses the returned event for both storage and the +caller, so the original text never reaches the session service: + +```python +from google.adk.plugins import BasePlugin +from google.adk.runners import Runner + + +class RedactSnowflakeSqlPlugin(BasePlugin): + """Replaces generated SQL in recorded tool calls before they are stored.""" + + def __init__(self): + super().__init__(name="redact_snowflake_sql") + + async def on_event_callback(self, *, invocation_context, event): + if event.partial or not event.content or not event.content.parts: + return None + changed = False + for part in event.content.parts: + call = part.function_call + if ( + call + and call.name == "system_execute_sql" + and call.args + and "sql" in call.args + ): + call.args = {**call.args, "sql": ""} + changed = True + return event if changed else None + + +runner = Runner( + app_name="analytics", + agent=root_agent, + session_service=session_service, + plugins=[RedactSnowflakeSqlPlugin()], +) +``` + +The same hook can drop result rows from `FunctionResponse` events or strip +tables from the final event. + +The agent itself never logs the user's question, generated SQL, tool payloads, +result rows, Snowflake thread or message ids, or request headers. The header +provider is excluded from `repr` and serialization, and errors quote +Snowflake's error code, message and request id but never the request body. + +## Limitations + +- **Root agent only.** Snowflake runs the agent loop and owns the thread, so the + agent cannot take part in another ADK agent's turn. Declaring `sub_agents` + or listing it in a parent's `sub_agents` raises a `ValueError` at + construction, and wrapping it in `AgentTool` is not supported. +- **Text in, one Cortex Agent object.** Only the text parts of the user message + are sent; images and files are not. The Cortex Agent must already exist in + Snowflake; this version does not create or configure agents. +- **Server-side tools only.** A Cortex Agent configured with a client-side tool + or one that asks for a permission decision cannot be run: the turn fails + with `UnsupportedCortexEventError`. +- **No reconnection.** If the connection drops mid-run, the turn fails and the + run is not resumed. Snowflake keeps the partial output of an abandoned run in + the thread and bills the work it did. +- **One turn per session at a time.** Two concurrent turns on one session + would both continue from the same message and fork the thread. Serialize + turns per session in the application before calling the `Runner`. +- **The cursor is bound to its configuration.** Changing the account, database, + schema or Cortex Agent for an existing session fails on the next turn until + the state key named in the error is removed or a new session is used. + +## Related samples + +* [Snowflake Cortex analyst](../../../../../contributing/samples/integrations/snowflake_cortex_agent/agent.py) - Runs a Cortex Agent object as an ADK root agent with credentials from the environment. diff --git a/src/google/adk/labs/snowflake/README.md b/src/google/adk/labs/snowflake/README.md new file mode 100644 index 0000000000..7461cc3012 --- /dev/null +++ b/src/google/adk/labs/snowflake/README.md @@ -0,0 +1,17 @@ +# Snowflake Cortex Agents Integration (Experimental) + +This folder contains an experimental integration that runs an existing +Snowflake Cortex Agent as an ADK root agent. Like everything under +`google.adk.labs`, it may change or be removed without notice. + +`SnowflakeCortexAgent` calls the Cortex Agents Run REST API directly with the +`httpx` client that ADK already depends on, so there is no extra package to +install. + +```python +from google.adk.labs.snowflake import SnowflakeCortexAgent +``` + +See the +[SnowflakeCortexAgent guide](../../../../../docs/guides/labs/snowflake/snowflake_cortex_agent/index.md) +for setup, configuration, limitations, and API details. diff --git a/src/google/adk/labs/snowflake/__init__.py b/src/google/adk/labs/snowflake/__init__.py new file mode 100644 index 0000000000..a0d3138ebd --- /dev/null +++ b/src/google/adk/labs/snowflake/__init__.py @@ -0,0 +1,46 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Snowflake Cortex Agents integration. + +Runs an existing, named Snowflake Cortex Agent as an ADK ``BaseAgent`` node. +The agent calls the Cortex Agents Run REST API directly with the ``httpx`` +client that ADK already depends on, so no extra package is required. + +Experimental: like everything under ``google.adk.labs``, this API may change +or be removed without notice. + +Example: + ```python + from google.adk.labs.snowflake import SnowflakeCortexAgent + + def bearer_headers(ctx): + return {'Authorization': f'Bearer {load_snowflake_token()}'} + + root_agent = SnowflakeCortexAgent( + name='sales_analyst', + account_url='https://.snowflakecomputing.com', + database='SALES_DB', + schema_name='ANALYTICS', + cortex_agent_name='SALES_AGENT', + header_provider=bearer_headers, + ) + ``` +""" + +from ._snowflake_cortex_agent import SnowflakeCortexAgent + +__all__ = [ + 'SnowflakeCortexAgent', +] diff --git a/src/google/adk/labs/snowflake/_client.py b/src/google/adk/labs/snowflake/_client.py new file mode 100644 index 0000000000..038e9c9f53 --- /dev/null +++ b/src/google/adk/labs/snowflake/_client.py @@ -0,0 +1,429 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HTTP client for the Snowflake Cortex Agents REST API. + +Thin ``httpx``-based access to the three calls the agent needs: creating a +thread, running the agent as an SSE stream, and cancelling a run. Credentials +come from a caller-supplied header provider on every request, so this module +never holds a token. Failures surface as typed errors that carry Snowflake's +error code and request id, never the request payload. +""" + +from __future__ import annotations + +import contextlib +import inspect +import logging +import re +from typing import Any +from typing import AsyncIterator +from typing import Awaitable +from typing import Callable +from typing import TYPE_CHECKING +from urllib.parse import quote + +import httpx + +from ._sse_parser import iter_sse_events +from ._sse_parser import SseEvent + +if TYPE_CHECKING: + from ...agents.readonly_context import ReadonlyContext + +logger = logging.getLogger('google_adk.' + __name__) + +HeaderProvider = Callable[ + ['ReadonlyContext'], 'dict[str, str] | Awaitable[dict[str, str]]' +] +"""Supplies the HTTP headers, typically ``Authorization``, for one request.""" + +_MAX_SNOWFLAKE_ID = 10**38 - 1 +_DECIMAL_RE = re.compile(r'[0-9]+') +_DEFAULT_ORIGIN_APPLICATION = 'google_adk' +_SSE_MEDIA_TYPE = 'text/event-stream' +_JSON_MEDIA_TYPE = 'application/json' +_MAX_ERROR_MESSAGE_CHARS = 200 +_AUTH_HINT = ( + ' Check that header_provider returns a valid Authorization header and,' + ' for tokens other than OAuth, the matching' + ' X-Snowflake-Authorization-Token-Type.' +) + + +class CortexClientError(Exception): + """Base class for failures talking to the Cortex Agents REST API.""" + + +class CortexApiError(CortexClientError): + """Snowflake answered a request with an error, or with an unusable body.""" + + def __init__( + self, + message: str, + *, + status_code: int, + snowflake_code: str | None = None, + request_id: str | None = None, + ): + super().__init__(message) + self.status_code = status_code + """The HTTP status Snowflake answered with.""" + self.snowflake_code = snowflake_code + """Snowflake's own error code from the response body, if any.""" + self.request_id = request_id + """Snowflake's request id, for quoting to Snowflake support.""" + + +class CortexTransportError(CortexClientError): + """A request never completed: connection failure, timeout, or a cut stream.""" + + def __init__(self, message: str, *, timed_out: bool): + super().__init__(message) + self.timed_out = timed_out + """Whether the failure was a timeout rather than a broken connection.""" + + +def _parse_snowflake_id(value: str | int, field: str, *, minimum: int) -> int: + """Converts a stored id into the integer Snowflake expects, strictly. + + The value is never echoed into the error: Snowflake ids stay out of logs. + """ + if isinstance(value, bool): + number = None + elif isinstance(value, int): + number = value + elif isinstance(value, str) and _DECIMAL_RE.fullmatch(value): + number = int(value) + else: + number = None + if number is None or not minimum <= number <= _MAX_SNOWFLAKE_ID: + raise ValueError( + f'{field} must be a decimal integer between {minimum} and 10^38-1;' + ' the stored Snowflake cursor is not usable.' + ) + return number + + +def _media_type(response: httpx.Response) -> str: + content_type: str = response.headers.get('content-type', '') + return content_type.partition(';')[0].strip().lower() + + +class SnowflakeCortexClient: + """Calls the Cortex Agents REST API for one configured Cortex Agent object. + + Each request asks ``header_provider`` for its headers, so credentials can be + minted per invocation and are never stored here. The client can share an + ``httpx.AsyncClient`` with the application or own one of its own, which + ``aclose`` releases. + """ + + def __init__( + self, + *, + account_url: str, + database: str, + schema_name: str, + cortex_agent_name: str, + header_provider: HeaderProvider, + timeout: float, + http_client: httpx.AsyncClient | None = None, + origin_application: str = _DEFAULT_ORIGIN_APPLICATION, + ): + """Initializes the client. + + Args: + account_url: Base URL of the Snowflake account. + database: Database holding the Cortex Agent object. + schema_name: Schema holding the Cortex Agent object. + cortex_agent_name: Name of the Cortex Agent object. + header_provider: Supplies the headers of each request, typically + ``Authorization`` and ``X-Snowflake-Authorization-Token-Type``. May be + sync or async. + timeout: Seconds to wait on Snowflake for a connection and, during a + run, between two SSE chunks. + http_client: An ``httpx.AsyncClient`` to send through. When omitted, the + client creates its own and closes it in ``aclose``. + origin_application: Label Snowflake stores on threads this client + creates. Snowflake accepts at most 16 UTF-8 bytes. + """ + base = account_url.rstrip('/') + self._threads_url = f'{base}/api/v2/cortex/threads' + self._run_url = ( + f'{base}/api/v2/databases/{quote(database, safe="")}' + f'/schemas/{quote(schema_name, safe="")}' + f'/agents/{quote(cortex_agent_name, safe="")}:run' + ) + self._cancel_url = f'{base}/api/v2/cortex/agent/runs/{{run_id}}/cancel' + self._header_provider = header_provider + # The read timeout is what bounds a stream that goes quiet between + # chunks; connecting should never take as long as a run may. + self._timeout = httpx.Timeout(timeout, connect=min(timeout, 30.0)) + self._http_client = http_client + self._owns_http_client = http_client is None + self._origin_application = origin_application + + async def create_thread(self, ctx: ReadonlyContext) -> str: + """Creates a Snowflake thread for a new conversation. + + Args: + ctx: The invocation's read-only context, passed to ``header_provider``. + + Returns: + The new thread id as a decimal string. + + Raises: + CortexApiError: Snowflake rejected the request or returned no thread id. + CortexTransportError: Snowflake could not be reached in time. + """ + response = await self._send( + ctx, + self._threads_url, + json_body={'origin_application': self._origin_application}, + accept=_JSON_MEDIA_TYPE, + operation='create thread', + ) + try: + payload = response.json() + except ValueError: + payload = None + thread_id = payload.get('thread_id') if isinstance(payload, dict) else None + if ( + isinstance(thread_id, bool) + or not isinstance(thread_id, int) + or not 1 <= thread_id <= _MAX_SNOWFLAKE_ID + ): + raise CortexApiError( + 'Snowflake created a thread but the response carried no usable' + ' thread_id.', + status_code=response.status_code, + request_id=_request_id(response, payload), + ) + return str(thread_id) + + @contextlib.asynccontextmanager + async def run( + self, + ctx: ReadonlyContext, + *, + thread_id: str | int, + parent_message_id: str | int, + text: str, + ) -> AsyncIterator[AsyncIterator[SseEvent]]: + """Runs the Cortex Agent on one user message and streams its events. + + Use as ``async with client.run(...) as events``. Leaving the block closes + the HTTP response, which is what stops Snowflake's stream when the + consumer gives up early. + + Args: + ctx: The invocation's read-only context, passed to ``header_provider``. + thread_id: The Snowflake thread to append to. + parent_message_id: The assistant message to continue from; ``0`` for + the first turn of a thread. + text: The user's message. + + Yields: + The run's events in stream order, ending with the ``done`` event when + Snowflake finished normally. + + Raises: + ValueError: ``thread_id`` or ``parent_message_id`` is not a Snowflake id. + CortexApiError: Snowflake rejected the run or did not answer with an + event stream. + CortexTransportError: Snowflake could not be reached, timed out, or the + stream was cut before the run finished. + """ + body = { + 'thread_id': _parse_snowflake_id(thread_id, 'thread_id', minimum=1), + 'parent_message_id': _parse_snowflake_id( + parent_message_id, 'parent_message_id', minimum=0 + ), + 'messages': [ + {'role': 'user', 'content': [{'type': 'text', 'text': text}]} + ], + 'stream': True, + } + response = await self._send( + ctx, + self._run_url, + json_body=body, + accept=_SSE_MEDIA_TYPE, + operation='run', + stream=True, + ) + try: + media_type = _media_type(response) + if media_type != _SSE_MEDIA_TYPE: + await response.aread() + raise CortexApiError( + 'Snowflake answered the run with' + f' {media_type or "no"} content instead of {_SSE_MEDIA_TYPE}.', + status_code=response.status_code, + request_id=_request_id(response, None), + ) + yield self._events(response) + finally: + await response.aclose() + + async def cancel_run(self, ctx: ReadonlyContext, run_id: str) -> bool: + """Asks Snowflake to cancel a run, best effort. + + Args: + ctx: The invocation's read-only context, passed to ``header_provider``. + run_id: The run to cancel. + + Returns: + Whether Snowflake acknowledged the cancel. A failure is logged by type + only and never raised: by the time this is called the consumer has + already gone. + """ + try: + response = await self._send( + ctx, + self._cancel_url.format(run_id=quote(run_id, safe='')), + json_body=None, + accept=_JSON_MEDIA_TYPE, + operation='cancel', + ) + except CortexClientError as e: + status = getattr(e, 'status_code', None) + logger.debug( + 'Best-effort cancel of a Snowflake Cortex run was refused: %s%s', + type(e).__name__, + f' (HTTP {status})' if status is not None else '', + ) + return False + await response.aclose() + return True + + async def aclose(self) -> None: + """Closes the HTTP client if this instance created it.""" + if self._owns_http_client and self._http_client is not None: + await self._http_client.aclose() + self._http_client = None + + def _client(self) -> httpx.AsyncClient: + if self._http_client is None: + self._http_client = httpx.AsyncClient(timeout=self._timeout) + return self._http_client + + async def _headers( + self, ctx: ReadonlyContext, *, accept: str + ) -> dict[str, str]: + provided = self._header_provider(ctx) + if inspect.isawaitable(provided): + provided = await provided + headers = dict(provided or {}) + # Content negotiation is the client's to decide: a provider that sets + # `Accept: application/json` would otherwise turn the run into a single + # JSON body. `identity` keeps proxies from buffering the event stream. + headers.update({ + 'Content-Type': _JSON_MEDIA_TYPE, + 'Accept': accept, + 'Accept-Encoding': 'identity', + }) + return headers + + async def _send( + self, + ctx: ReadonlyContext, + url: str, + *, + json_body: dict[str, Any] | None, + accept: str, + operation: str, + stream: bool = False, + ) -> httpx.Response: + client = self._client() + request = client.build_request( + 'POST', + url, + json=json_body, + headers=await self._headers(ctx, accept=accept), + timeout=self._timeout, + ) + try: + response = await client.send(request, stream=stream) + except httpx.TimeoutException as e: + raise CortexTransportError( + f'Snowflake did not answer the {operation} request in time.', + timed_out=True, + ) from e + except httpx.RequestError as e: + raise CortexTransportError( + f'Snowflake could not be reached for the {operation} request.', + timed_out=False, + ) from e + if response.is_success: + return response + try: + await response.aread() + raise _api_error(response, operation) + finally: + await response.aclose() + + async def _events(self, response: httpx.Response) -> AsyncIterator[SseEvent]: + try: + async for event in iter_sse_events(response.aiter_bytes()): + yield event + except httpx.TimeoutException as e: + raise CortexTransportError( + 'Snowflake stopped sending events before the run finished.', + timed_out=True, + ) from e + except httpx.RequestError as e: + raise CortexTransportError( + 'The connection to Snowflake dropped before the run finished.', + timed_out=False, + ) from e + + +def _request_id(response: httpx.Response, payload: Any) -> str | None: + header: str | None = response.headers.get('x-snowflake-request-id') + if header: + return header + if isinstance(payload, dict) and payload.get('request_id') is not None: + return str(payload['request_id']) + return None + + +def _api_error(response: httpx.Response, operation: str) -> CortexApiError: + try: + payload = response.json() + except ValueError: + payload = None + code = payload.get('code') if isinstance(payload, dict) else None + message = payload.get('message') if isinstance(payload, dict) else None + detail = ( + str(message)[:_MAX_ERROR_MESSAGE_CHARS] + if message + else 'no error message in the response' + ) + text = ( + f'Snowflake rejected the {operation} request with HTTP' + f' {response.status_code}' + ) + if code is not None: + text += f' (code {code})' + text += f': {detail}.' + if response.status_code in (401, 403): + text += _AUTH_HINT + return CortexApiError( + text, + status_code=response.status_code, + snowflake_code=str(code) if code is not None else None, + request_id=_request_id(response, payload), + ) diff --git a/src/google/adk/labs/snowflake/_event_converter.py b/src/google/adk/labs/snowflake/_event_converter.py new file mode 100644 index 0000000000..94945ee3e7 --- /dev/null +++ b/src/google/adk/labs/snowflake/_event_converter.py @@ -0,0 +1,625 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Translates Cortex Agents Run API events into ADK events. + +Kept separate from the agent so the mapping rules stay readable and testable +without a network or a session. One ``CortexEventConverter`` accumulates a +single run: it deduplicates repeated deltas, pairs server-side tool use with +tool results, bounds what gets persisted, and builds the one final event from +the run's authoritative ``response`` payload. + +Scope: progress and unknown events as partial metadata events, thinking and +text deltas as partial parts (both only in SSE streaming mode), server-side +tool use and results as ``FunctionCall`` / ``FunctionResponse`` events, +annotations, warnings, tables, charts and suggested queries as final-event +metadata, ``error`` as a terminal error event, and ``done`` as the transport +terminator. +""" + +from __future__ import annotations + +import json +from typing import Any +from typing import Callable +from typing import TYPE_CHECKING + +from google.genai import types as genai_types +from pydantic import JsonValue + +from ...events.event import Event +from ...events.event_actions import EventActions +from ._sse_parser import SseEvent +from ._sse_parser import SseParseError + +if TYPE_CHECKING: + from ...agents.invocation_context import InvocationContext + +METADATA_KEY = 'snowflake_cortex' +"""The ``custom_metadata`` key under which every Cortex detail is namespaced.""" + +_DEFAULT_ERROR_CODE = 'SNOWFLAKE_CORTEX_ERROR' +_DEFAULT_ERROR_MESSAGE = 'The Snowflake Cortex Agent run failed.' +_DEFAULT_TOOL_ERROR = {'message': 'The Snowflake tool call failed.'} + +_Handler = Callable[[str, dict[str, Any]], list[Event]] + + +class UnsupportedCortexEventError(RuntimeError): + """The run asked for a capability this integration does not support yet. + + Raised when Cortex requests client-side tool execution or a permission + decision, neither of which this agent can answer. The run is left to time + out on the Snowflake side; the caller sees the turn fail. + """ + + +def _json_size(value: Any) -> int: + return len( + json.dumps(value, separators=(',', ':'), ensure_ascii=False).encode( + 'utf-8' + ) + ) + + +def _as_dict(value: Any) -> dict[str, Any]: + return value if isinstance(value, dict) else {} + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + return value if isinstance(value, list) else [value] + + +def _content_index(payload: dict[str, Any]) -> int: + value = payload.get('content_index') + return value if isinstance(value, int) else 0 + + +def _without_result_rows(item: Any) -> Any: + """Drops ``result_set.data`` from one tool-result content block.""" + if not isinstance(item, dict): + return item + body = item.get('json') + if not isinstance(body, dict) or not isinstance(body.get('result_set'), dict): + return item + result_set = {k: v for k, v in body['result_set'].items() if k != 'data'} + return {**item, 'json': {**body, 'result_set': result_set}} + + +def _content_shape(item: Any) -> dict[str, JsonValue]: + """Reduces one tool-result content block to its type and key sizes.""" + if not isinstance(item, dict): + return {'type': type(item).__name__, 'bytes': _json_size(item)} + shape: dict[str, JsonValue] = {'type': item.get('type')} + body = item.get('json') + if isinstance(body, dict): + shape['json_keys'] = { + str(key): _json_size(value) for key, value in body.items() + } + else: + shape['keys'] = { + str(key): _json_size(value) + for key, value in item.items() + if key != 'type' + } + return shape + + +def _tool_error(content: list[Any]) -> JsonValue: + for item in content: + body = _as_dict(item).get('json') + if isinstance(body, dict) and 'error' in body: + error: JsonValue = body['error'] + return error + return dict(_DEFAULT_TOOL_ERROR) + + +class CortexEventConverter: + """Accumulates one Cortex run and maps its events onto ADK events. + + Create one per run, pass every parsed ``SseEvent`` to ``convert`` in stream + order, and once ``is_done`` and ``has_final_response`` hold, ask + ``final_event`` for the single non-partial event that carries the answer. + """ + + def __init__( + self, + *, + ctx: InvocationContext, + author: str, + streaming: bool, + max_tool_result_bytes: int, + include_thinking_in_final_event: bool, + thread_id: str | None = None, + ): + """Initializes the converter for one run. + + Args: + ctx: The invocation context whose id and branch every event carries. + author: The ADK agent name to author model events with. + streaming: Whether the consumer asked for SSE streaming. Partial events + are produced only when this is true. + max_tool_result_bytes: Size bound for one recorded tool result, table or + chart; larger payloads are reduced to their metadata. + include_thinking_in_final_event: Whether the final event also carries the + completed reasoning as a ``thought`` part. + thread_id: The Snowflake thread the run belongs to, when known, so a + ``run_id`` can be derived if the final payload does not name one. + """ + self._ctx = ctx + self._author = author + self._streaming = streaming + self._max_bytes = max_tool_result_bytes + self._include_thinking = include_thinking_in_final_event + self._thread_id = thread_id + self._user_message_id: str | None = None + self._assistant_message_id: str | None = None + self._seen_deltas: set[tuple[str, int, int]] = set() + self._seen_tool_calls: set[str] = set() + self._seen_tool_results: set[str] = set() + self._text_blocks: dict[int, str] = {} + self._thinking_blocks: dict[int, str] = {} + self._annotations: list[Any] = [] + self._warnings: list[Any] = [] + self._tables: list[Any] = [] + self._charts: list[Any] = [] + self._suggested_queries: list[Any] = [] + # (tool_use_id, suggestion index) -> text assembled from deltas. + self._analyst_suggestions: dict[tuple[str, int], str] = {} + self._final_response: dict[str, Any] | None = None + self._error: dict[str, Any] | None = None + self._done = False + self._handlers: dict[str, _Handler] = { + 'metadata': self._on_metadata, + 'response.status': self._on_progress, + 'response.thinking.delta': self._on_thinking_delta, + 'response.thinking': self._on_thinking, + 'response.text.delta': self._on_text_delta, + 'response.text': self._on_text, + 'response.tool_use': self._on_tool_use, + 'response.tool_result.status': self._on_progress, + 'response.tool_result.analyst.delta': self._on_analyst_delta, + 'response.tool_result': self._on_tool_result, + 'response.text.annotation': self._on_annotation, + 'response.warning': self._on_warning, + 'response.table': self._on_table, + 'response.chart': self._on_chart, + 'response.suggested_queries': self._on_suggested_queries, + 'response': self._on_response, + 'error': self._on_error, + } + + @property + def thread_id(self) -> str | None: + """The Snowflake thread id, from the caller or the final payload.""" + return self._thread_id + + @property + def user_message_id(self) -> str | None: + """The id Snowflake gave the user message of this run, if seen.""" + return self._user_message_id + + @property + def assistant_message_id(self) -> str | None: + """The id Snowflake gave the assistant message of this run, if seen.""" + return self._assistant_message_id + + @property + def is_done(self) -> bool: + """Whether the transport terminator has been seen.""" + return self._done + + @property + def has_final_response(self) -> bool: + """Whether the authoritative final ``response`` has been seen.""" + return self._final_response is not None + + @property + def final_status(self) -> str | None: + """The ``status`` of the final ``response``, if it has been seen.""" + if self._final_response is None: + return None + status = self._final_response.get('status') + return str(status) if status is not None else None + + @property + def failed(self) -> bool: + """Whether a terminal ``error`` event has been seen.""" + return self._error is not None + + def convert(self, sse_event: SseEvent) -> list[Event]: + """Maps one Cortex event onto zero or more ADK events. + + Args: + sse_event: The next event of the run, in stream order. + + Returns: + The ADK events to yield for it, possibly none. + + Raises: + SseParseError: The event's data is not a JSON object. + UnsupportedCortexEventError: The run needs client-side tool execution or + a permission decision. + """ + if sse_event.is_done: + self._done = True + return [] + payload = sse_event.json_data() + if not isinstance(payload, dict): + raise SseParseError( + f'SSE event {sse_event.event!r} carries JSON that is not an object.' + ) + handler = self._handlers.get(sse_event.event) + if handler is None: + return self._partial_metadata( + {'unknown': {'event': sse_event.event, 'data': payload}} + ) + return handler(sse_event.event, payload) + + def final_event(self, *, state_delta: dict[str, Any] | None = None) -> Event: + """Builds the single non-partial event that closes a successful run. + + Args: + state_delta: Session state to commit with the event, typically the + thread cursor. Applied by the ADK ``Runner``. + + Returns: + A model event carrying the answer text, the run's metadata under + ``custom_metadata['snowflake_cortex']``, and ``state_delta``. + + Raises: + ValueError: No final ``response`` was received. + """ + if self._final_response is None: + raise ValueError( + 'The Cortex run ended without a final response event, so there is' + ' no answer to record.' + ) + content = _as_list(self._final_response.get('content')) + metadata = _as_dict(self._final_response.get('metadata')) + + parts: list[genai_types.Part] = [] + if self._include_thinking: + thinking = self._final_text(content, 'thinking', self._thinking_blocks) + if thinking: + parts.append(genai_types.Part(text=thinking, thought=True)) + text = self._final_text(content, 'text', self._text_blocks) + if text: + parts.append(genai_types.Part.from_text(text=text)) + + suggested = self._suggested_queries or [ + query + for block in content + if _as_dict(block).get('type') == 'suggested_queries' + for query in _as_list(block.get('suggested_queries')) + ] + # Questions Cortex Analyst proposes when it cannot answer arrive as + # per-index text deltas; they are follow-ups too, so they join the list + # with their origin marked. + suggested = list(suggested) + [ + {'query': text, 'source': 'cortex_analyst', 'tool_use_id': tool_use_id} + for (tool_use_id, _), text in sorted(self._analyst_suggestions.items()) + if text.strip() + ] + run_id = metadata.get('run_id') + if run_id is None and self._thread_id and self._user_message_id: + run_id = f'{self._thread_id}-{self._user_message_id}' + + return Event( + invocation_id=self._ctx.invocation_id, + author=self._author, + branch=self._ctx.branch, + content=genai_types.Content(role='model', parts=parts), + custom_metadata={ + METADATA_KEY: { + 'run_id': run_id, + 'status': self._final_response.get('status'), + 'annotations': list(self._annotations), + 'warnings': list(self._warnings), + 'suggested_queries': suggested, + 'usage': metadata.get('usage'), + 'tables': list(self._tables), + 'charts': list(self._charts), + } + }, + actions=EventActions(state_delta=state_delta or {}), + ) + + def _final_text( + self, content: list[Any], block_type: str, buffered: dict[int, str] + ) -> str: + # The final `response` aggregates the run, so its blocks win over the + # deltas; the buffer only covers a payload that omits them. + texts = [] + for block in content: + block = _as_dict(block) + if block.get('type') != block_type: + continue + value = block.get(block_type) + if isinstance(value, dict): + value = value.get('text') + if isinstance(value, str): + texts.append(value) + if texts: + return ''.join(texts) + return ''.join(buffered[index] for index in sorted(buffered)) + + def _event(self, **kwargs: Any) -> Event: + return Event( + invocation_id=self._ctx.invocation_id, + author=self._author, + branch=self._ctx.branch, + **kwargs, + ) + + def _partial_metadata(self, body: dict[str, Any]) -> list[Event]: + if not self._streaming: + return [] + return [self._event(partial=True, custom_metadata={METADATA_KEY: body})] + + def _on_metadata(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + metadata = _as_dict(payload.get('metadata')) or payload + message_id = metadata.get('message_id') + if message_id is not None: + # Decimal strings: Snowflake ids are NUMBER(38,0) and would lose + # precision in JSON consumers that read them as doubles. + if metadata.get('role') == 'user': + self._user_message_id = str(message_id) + elif metadata.get('role') == 'assistant': + self._assistant_message_id = str(message_id) + if metadata.get('thread_id') is not None: + self._thread_id = str(metadata['thread_id']) + return [] + + def _on_progress(self, name: str, payload: dict[str, Any]) -> list[Event]: + return self._partial_metadata({'event': name, 'data': payload}) + + def _on_analyst_delta( + self, name: str, payload: dict[str, Any] + ) -> list[Event]: + suggestion = _as_dict(_as_dict(payload.get('delta')).get('suggestions')) + text = suggestion.get('delta') + index = suggestion.get('index') + if isinstance(text, str) and text and isinstance(index, int): + key = (str(payload.get('tool_use_id') or ''), index) + self._analyst_suggestions[key] = ( + self._analyst_suggestions.get(key, '') + text + ) + return self._on_progress(name, payload) + + def _on_thinking_delta( + self, name: str, payload: dict[str, Any] + ) -> list[Event]: + return self._on_delta(name, payload, self._thinking_blocks, thought=True) + + def _on_text_delta(self, name: str, payload: dict[str, Any]) -> list[Event]: + return self._on_delta(name, payload, self._text_blocks, thought=False) + + def _on_delta( + self, + name: str, + payload: dict[str, Any], + blocks: dict[int, str], + *, + thought: bool, + ) -> list[Event]: + index = _content_index(payload) + sequence = payload.get('sequence_number') + if isinstance(sequence, int): + # Snowflake may resend a delta with the same sequence number; appending + # it twice would duplicate text in the buffer and on screen. + key = (name, index, sequence) + if key in self._seen_deltas: + return [] + self._seen_deltas.add(key) + text = payload.get('text') + if not isinstance(text, str) or not text: + return [] + blocks[index] = blocks.get(index, '') + text + if not self._streaming: + return [] + part = ( + genai_types.Part(text=text, thought=True) + if thought + else genai_types.Part.from_text(text=text) + ) + return [ + self._event( + partial=True, + content=genai_types.Content(role='model', parts=[part]), + ) + ] + + def _on_thinking(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + return self._on_block(payload, self._thinking_blocks) + + def _on_text(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + return self._on_block(payload, self._text_blocks) + + def _on_block( + self, payload: dict[str, Any], blocks: dict[int, str] + ) -> list[Event]: + # The completed block is authoritative for its index; the deltas that + # built it up may have been reordered or duplicated on the wire. + text = payload.get('text') + if isinstance(text, str): + blocks[_content_index(payload)] = text + return [] + + def _tool_identity(self, payload: dict[str, Any]) -> tuple[str, str]: + tool_use_id = payload.get('tool_use_id') + if tool_use_id is None: + tool_use_id = f"{payload.get('type')}-{payload.get('sequence_number')}" + name = payload.get('name') or payload.get('type') or 'unknown_tool' + return str(tool_use_id), str(name) + + def _on_tool_use(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + if payload.get('client_side_execute') or 'permission' in payload: + raise UnsupportedCortexEventError( + 'The Cortex Agent asked for a client-side tool execution or a' + ' permission decision, which SnowflakeCortexAgent does not support' + ' yet. Configure the Cortex Agent with server-side tools only.' + ) + tool_use_id, tool_name = self._tool_identity(payload) + if tool_use_id in self._seen_tool_calls: + return [] + self._seen_tool_calls.add(tool_use_id) + args = payload.get('input') + if not isinstance(args, dict): + args = {} if args is None else {'input': args} + return [ + self._event( + content=genai_types.Content( + role='model', + parts=[ + genai_types.Part( + function_call=genai_types.FunctionCall( + name=tool_name, args=args, id=tool_use_id + ) + ) + ], + ) + ) + ] + + def _on_tool_result(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + tool_use_id, tool_name = self._tool_identity(payload) + if tool_use_id in self._seen_tool_results: + return [] + self._seen_tool_results.add(tool_use_id) + status = str(payload.get('status') or 'unknown') + content = _as_list(payload.get('content')) + response: dict[str, JsonValue] = {'status': status, 'content': content} + if status == 'error': + response['error'] = _tool_error(content) + return [ + Event( + invocation_id=self._ctx.invocation_id, + # Authored by the tool so session history attributes the response + # to it, mirroring ADK's own function-response events. + author=tool_name, + branch=self._ctx.branch, + content=genai_types.Content( + role='user', + parts=[ + genai_types.Part( + function_response=genai_types.FunctionResponse( + name=tool_name, + id=tool_use_id, + response=self._bound_tool_result(response), + ) + ) + ], + ), + ) + ] + + def _bound_tool_result( + self, response: dict[str, JsonValue] + ) -> dict[str, JsonValue]: + # The response is persisted with the session, so it is cut down in stages + # rather than stored whole: first the SQL rows go (query id and column + # metadata stay), then each block is reduced to its type and the sizes of + # its keys, so the record still says what the tool returned. + size = _json_size(response) + if size <= self._max_bytes: + return response + content = [ + _without_result_rows(item) for item in _as_list(response['content']) + ] + bounded: dict[str, JsonValue] = { + **response, + 'content': content, + 'truncated': True, + 'original_bytes': size, + } + if _json_size(bounded) <= self._max_bytes: + return bounded + shaped: dict[str, JsonValue] = { + **bounded, + 'content': [_content_shape(item) for item in content], + } + if _json_size(shaped) <= self._max_bytes: + return shaped + return {**bounded, 'content': []} + + def _on_annotation(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + self._annotations.append(payload) + return [] + + def _on_warning(self, name: str, payload: dict[str, Any]) -> list[Event]: + self._warnings.append(payload) + return self._partial_metadata({'event': name, 'data': payload}) + + def _on_table(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + table = payload + if _json_size(table) > self._max_bytes: + result_set = { + k: v + for k, v in _as_dict(table.get('result_set')).items() + if k != 'data' + } + table = {**table, 'result_set': result_set, 'truncated': True} + self._tables.append(table) + return [] + + def _on_chart(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + chart = payload + if _json_size(chart) > self._max_bytes: + chart = {k: v for k, v in chart.items() if k != 'chart_spec'} + chart['truncated'] = True + self._charts.append(chart) + return [] + + def _on_suggested_queries( + self, name: str, payload: dict[str, Any] + ) -> list[Event]: + del name + self._suggested_queries.extend(_as_list(payload.get('suggested_queries'))) + return [] + + def _on_response(self, name: str, payload: dict[str, Any]) -> list[Event]: + del name + self._final_response = payload + metadata = _as_dict(payload.get('metadata')) + if ( + self._assistant_message_id is None + and metadata.get('assistant_message_id') is not None + ): + self._assistant_message_id = str(metadata['assistant_message_id']) + if self._thread_id is None and metadata.get('thread_id') is not None: + self._thread_id = str(metadata['thread_id']) + return [] + + def _on_error(self, name: str, payload: dict[str, Any]) -> list[Event]: + self._error = payload + code = payload.get('code') or payload.get('error_code') + message = payload.get('message') or payload.get('error_message') + return [ + self._event( + error_code=str(code) if code else _DEFAULT_ERROR_CODE, + error_message=str(message) if message else _DEFAULT_ERROR_MESSAGE, + custom_metadata={METADATA_KEY: {'event': name, 'data': payload}}, + ) + ] diff --git a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py new file mode 100644 index 0000000000..da4fe7054c --- /dev/null +++ b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py @@ -0,0 +1,426 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runs a Snowflake Cortex Agent as an ADK agent. + +Wraps an existing, named Snowflake Cortex Agent object as a native ADK +``BaseAgent`` node. Snowflake runs the agent loop and owns the conversation +thread; this node sends each ADK turn to the Cortex Agents Run API and +projects the resulting SSE stream onto ADK events. + +Because the loop and the thread live in Snowflake, a ``SnowflakeCortexAgent`` +must run as an ADK root agent: it accepts no ``sub_agents`` and refuses to be +adopted by a parent agent. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import hashlib +import logging +from typing import Any +from typing import AsyncGenerator +from typing import Awaitable +from typing import Callable + +import httpx +from pydantic import ConfigDict +from pydantic import Field +from pydantic import PrivateAttr +from typing_extensions import override + +from ...agents._streaming_mode import StreamingMode +from ...agents.base_agent import BaseAgent +from ...agents.invocation_context import InvocationContext +from ...agents.readonly_context import ReadonlyContext +from ...events.event import Event +from ._client import CortexTransportError +from ._client import SnowflakeCortexClient +from ._event_converter import CortexEventConverter + +logger = logging.getLogger('google_adk.' + __name__) + +_STATE_KEY_PREFIX = '_snowflake_cortex_' +_CURSOR_SCHEMA_VERSION = 1 + +_SUB_AGENTS_NOT_SUPPORTED_MESSAGE = ( + 'SnowflakeCortexAgent does not support sub_agents: the agent loop runs' + ' inside Snowflake, where an ADK sub-agent cannot be reached.' +) + +_PARENT_NOT_SUPPORTED_MESSAGE = ( + 'SnowflakeCortexAgent must run as an ADK root agent and cannot be a' + ' sub-agent: Snowflake runs the agent loop and owns the conversation' + ' thread, so it cannot take part in the turn of an ADK parent.' +) + + +@dataclasses.dataclass(frozen=True) +class _Cursor: + """Where the next turn continues in the Snowflake thread.""" + + thread_id: str + parent_message_id: str + + def to_state(self, fingerprint: str) -> dict[str, Any]: + return { + 'schema_version': _CURSOR_SCHEMA_VERSION, + 'resource_fingerprint': fingerprint, + 'thread_id': self.thread_id, + 'parent_message_id': self.parent_message_id, + } + + +class SnowflakeCortexAgent(BaseAgent): + """Runs a Snowflake Cortex Agent as an ADK agent node. + + Each ADK turn sends the user's message to an existing Cortex Agent object + through the Cortex Agents Run API and streams the run back as ADK events: + partial text and reasoning deltas in SSE streaming mode, server-side tool + calls and results as ``FunctionCall`` / ``FunctionResponse`` events, and one + final event carrying the answer with citations, warnings, tables, charts and + suggested queries under ``custom_metadata['snowflake_cortex']``. + + The Snowflake thread and the last assistant message id are kept in ADK + session state under a key scoped to this agent's ``name``, so a conversation + continues across turns and survives a restart. Persisting that cursor needs + the ADK ``Runner``, which is what applies a yielded event's ``state_delta``. + The cursor also records which account and Cortex Agent object it belongs + to; pointing an existing session at a different one fails loudly rather + than mixing two conversations. + + Credentials are supplied per request by ``header_provider`` rather than + stored on the agent, and the provider is excluded from ``repr`` and + serialization. + + Must be an ADK root agent: ``sub_agents`` are rejected and a parent cannot + adopt it. + + Example: + ```python + from google.adk.agents.readonly_context import ReadonlyContext + from google.adk.labs.snowflake import SnowflakeCortexAgent + + def bearer_headers(ctx: ReadonlyContext) -> dict[str, str]: + return { + 'Authorization': f'Bearer {load_snowflake_token()}', + 'X-Snowflake-Authorization-Token-Type': 'PROGRAMMATIC_ACCESS_TOKEN', + } + + root_agent = SnowflakeCortexAgent( + name='sales_analyst', + account_url='https://.snowflakecomputing.com', + database='SALES_DB', + schema_name='ANALYTICS', + cortex_agent_name='SALES_AGENT', + header_provider=bearer_headers, + ) + ``` + """ + + model_config = ConfigDict( + arbitrary_types_allowed=True, + use_attribute_docstrings=True, + extra='forbid', + ) + + account_url: str + """Base URL of the Snowflake account. + + For example ``https://.snowflakecomputing.com``, without a trailing + slash. + """ + + database: str + """Database that holds the Cortex Agent object.""" + + schema_name: str + """Schema that holds the Cortex Agent object. + + Named ``schema_name`` because ``schema`` is reserved by pydantic. + """ + + cortex_agent_name: str + """Name of the Cortex Agent object in Snowflake. + + Distinct from ``name``, which identifies this node within ADK; two ADK + agents may point at the same Snowflake object. + """ + + header_provider: Callable[ + [ReadonlyContext], dict[str, str] | Awaitable[dict[str, str]] + ] = Field(exclude=True, repr=False) + """Supplies the HTTP headers for each Snowflake request. + + Typically ``Authorization`` and, for tokens other than OAuth, the matching + ``X-Snowflake-Authorization-Token-Type``. Called with the + ``ReadonlyContext`` of the current invocation and may be sync or async. + Excluded from serialization and ``repr`` so that a token never reaches the + ``adk web`` agent graph, logs, or a session store. + """ + + http_client: httpx.AsyncClient | None = Field( + default=None, exclude=True, repr=False + ) + """An ``httpx.AsyncClient`` to send Snowflake requests through. + + Share one to pool connections across agents or to configure proxies and + certificates. When omitted the agent creates its own and closes it in + ``cleanup``. Excluded from serialization: it is runtime wiring. + """ + + timeout: float = Field(default=900.0, gt=0) + """Seconds to wait on Snowflake before the turn fails with a timeout. + + Cortex Agent runs that plan, execute SQL and summarize can take minutes, so + the default is deliberately long. + """ + + cancel_on_disconnect: bool = True + """Whether to cancel the Snowflake run when the ADK consumer stops reading. + + Best effort: the cancel is attempted, not guaranteed, and Snowflake keeps + whatever partial output it already produced in the thread either way. + """ + + max_tool_result_bytes: int = Field(default=32 * 1024, gt=0) + """Upper bound on the serialized size of one recorded tool result. + + A server-side tool result larger than this is cut down before it is + recorded in a ``FunctionResponse`` event: SQL rows go first, then each + block is reduced to its type and key sizes. Tool results are persisted + with the session, so this bounds how much a single result can grow it. + """ + + include_thinking_in_final_event: bool = False + """Whether the final event also carries the completed reasoning text. + + Off by default so that reasoning is not written to the session store with + the final event. Reasoning deltas are still streamed as partial events in + SSE mode. + """ + + _cortex_client: SnowflakeCortexClient | None = PrivateAttr(default=None) + + @override + def model_post_init(self, __context: Any) -> None: + super().model_post_init(__context) + self._validate_no_sub_agents() + + def _validate_no_sub_agents(self) -> None: + # Called again on entry to `_run_async_impl` because `sub_agents` can be + # mutated or `model_copy`-ed after construction, bypassing + # `model_post_init`. + if self.sub_agents: + raise ValueError(_SUB_AGENTS_NOT_SUPPORTED_MESSAGE) + + def __setattr__(self, name: str, value: Any) -> None: + # `BaseAgent` adopts a child by assigning `parent_agent` from the parent's + # `model_post_init`, so refusing the assignment fails the parent's + # construction at its `sub_agents=[...]` declaration rather than a turn. + if name == 'parent_agent' and value is not None: + raise ValueError(_PARENT_NOT_SUPPORTED_MESSAGE) + super().__setattr__(name, value) + + def _state_key(self) -> str: + # Scoped by agent name so two `SnowflakeCortexAgent`s in one ADK session + # do not continue each other's Snowflake thread. + return _STATE_KEY_PREFIX + self.name + + def _resource_fingerprint(self) -> str: + # The account is part of it: a thread id only means something within the + # account that issued it, whatever the object is called. + material = '|'.join([ + self.account_url.rstrip('/'), + self.database, + self.schema_name, + self.cortex_agent_name, + ]) + return 'sha256:' + hashlib.sha256(material.encode('utf-8')).hexdigest() + + def _read_cursor(self, stored: Any) -> _Cursor | None: + """Validates the stored cursor, or returns None when there is none yet.""" + if stored is None: + return None + key = self._state_key() + remedy = f' Remove session state key {key!r} or start a new session.' + if ( + not isinstance(stored, dict) + or stored.get('schema_version') != _CURSOR_SCHEMA_VERSION + ): + raise ValueError( + f'Session state key {key!r} does not hold a SnowflakeCortexAgent' + ' cursor this version understands.' + + remedy + ) + if stored.get('resource_fingerprint') != self._resource_fingerprint(): + raise ValueError( + f'Session state key {key!r} holds a Snowflake thread that belongs' + ' to a different account, database, schema or Cortex Agent than' + ' this agent is configured with; continuing it would mix two' + ' conversations.' + + remedy + ) + thread_id = stored.get('thread_id') + parent_message_id = stored.get('parent_message_id') + if not isinstance(thread_id, str) or not isinstance(parent_message_id, str): + raise ValueError( + f'Session state key {key!r} holds a cursor without string' + ' thread_id and parent_message_id.' + + remedy + ) + return _Cursor(thread_id=thread_id, parent_message_id=parent_message_id) + + def _user_text(self, ctx: InvocationContext) -> str: + parts = ( + ctx.user_content.parts + if ctx.user_content and ctx.user_content.parts + else [] + ) + text = '\n'.join(part.text for part in parts if part.text) + if not text.strip(): + raise ValueError( + 'SnowflakeCortexAgent needs a text message: this version sends only' + ' text to the Cortex Agents Run API, and the current user content' + ' has none.' + ) + return text + + def _get_client(self) -> SnowflakeCortexClient: + if self._cortex_client is None: + self._cortex_client = SnowflakeCortexClient( + account_url=self.account_url, + database=self.database, + schema_name=self.schema_name, + cortex_agent_name=self.cortex_agent_name, + header_provider=self.header_provider, + timeout=self.timeout, + http_client=self.http_client, + ) + return self._cortex_client + + async def cleanup(self) -> None: + """Closes the HTTP client this agent created. + + A shared ``http_client`` is left open for its owner to close. + """ + if self._cortex_client is not None: + await self._cortex_client.aclose() + self._cortex_client = None + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + self._validate_no_sub_agents() + text = self._user_text(ctx) + cursor = self._read_cursor(ctx.session.state.get(self._state_key())) + readonly_ctx = ReadonlyContext(ctx) + client = self._get_client() + if cursor is None: + cursor = _Cursor( + thread_id=await client.create_thread(readonly_ctx), + parent_message_id='0', + ) + + streaming = bool( + ctx.run_config and ctx.run_config.streaming_mode == StreamingMode.SSE + ) + converter = CortexEventConverter( + ctx=ctx, + author=self.name, + streaming=streaming, + max_tool_result_bytes=self.max_tool_result_bytes, + include_thinking_in_final_event=self.include_thinking_in_final_event, + thread_id=cursor.thread_id, + ) + + try: + async with client.run( + readonly_ctx, + thread_id=cursor.thread_id, + parent_message_id=cursor.parent_message_id, + text=text, + ) as events: + async for sse_event in events: + for event in converter.convert(sse_event): + yield event + if converter.is_done or converter.failed: + # Nothing useful follows the terminator; stop reading rather than + # wait for Snowflake to close the connection. + break + except (GeneratorExit, asyncio.CancelledError): + # The consumer is gone. `run()` has already closed the upstream + # response; awaiting here is fine, yielding would not be. + await self._cancel_abandoned_run(readonly_ctx, client, converter) + raise + + if converter.failed: + # The terminal error event has been yielded; the cursor stays where it + # was so the next turn continues from the last good message. + return + if not converter.has_final_response: + raise CortexTransportError( + 'The Snowflake stream ended before the final response arrived, so' + ' the answer is incomplete. The conversation cursor was left' + ' unchanged.', + timed_out=False, + ) + if not converter.is_done: + # `[DONE]` is a compatibility sentinel; the final `response` is what + # closes a run, so a stream that ends right after it is complete. + logger.debug('Snowflake closed the run stream without a [DONE] event.') + + state_delta: dict[str, Any] | None = None + if ( + converter.final_status == 'completed' + and converter.assistant_message_id is not None + ): + # Only a completed run's assistant message may become the parent of the + # next turn: a cancelled or timed-out run can still store a partial + # assistant message, and the user id would fork the thread. + state_delta = { + self._state_key(): ( + _Cursor( + thread_id=cursor.thread_id, + parent_message_id=converter.assistant_message_id, + ).to_state(self._resource_fingerprint()) + ) + } + yield converter.final_event(state_delta=state_delta) + + async def _cancel_abandoned_run( + self, + readonly_ctx: ReadonlyContext, + client: SnowflakeCortexClient, + converter: CortexEventConverter, + ) -> None: + if ( + not self.cancel_on_disconnect + or converter.is_done + or converter.failed + or converter.has_final_response + ): + return + if converter.thread_id is None or converter.user_message_id is None: + # Snowflake names a run `{thread_id}-{user_message_id}`, so until the + # user message is acknowledged there is nothing to cancel by. + return + run_id = f'{converter.thread_id}-{converter.user_message_id}' + try: + await client.cancel_run(readonly_ctx, run_id) + except Exception: # pylint: disable=broad-exception-caught + # Best effort only: the consumer that would care has already gone. + logger.debug('Cancelling an abandoned Snowflake Cortex run failed.') diff --git a/src/google/adk/labs/snowflake/_sse_parser.py b/src/google/adk/labs/snowflake/_sse_parser.py new file mode 100644 index 0000000000..828bac7669 --- /dev/null +++ b/src/google/adk/labs/snowflake/_sse_parser.py @@ -0,0 +1,245 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Incremental parser for the Server-Sent Events wire format. + +The Cortex Agents Run API streams ``text/event-stream``. Network chunks line up +neither with event boundaries nor with UTF-8 code points, so the parser is fed +raw bytes and emits each event once its terminating blank line has arrived. +Framing follows the SSE specification: ``data:`` lines join with LF, a line +ends at LF, CR or CRLF, a leading byte-order mark is dropped, comment lines are +skipped, and an event cut off by the end of the stream is discarded. The parser +knows nothing about the JSON inside ``data``. +""" + +from __future__ import annotations + +import codecs +import dataclasses +import json +from typing import Any +from typing import AsyncGenerator +from typing import AsyncIterable + +_DEFAULT_MAX_EVENT_BYTES = 16 * 1024 * 1024 +_DEFAULT_EVENT_NAME = 'message' +_DONE_EVENT_NAME = 'done' +_DONE_DATA = '[DONE]' +_BOM = '\ufeff' + + +class SseParseError(ValueError): + """The byte stream is not a usable event stream. + + Raised for an event larger than the parser's limit, and for ``data`` that is + not JSON when JSON was asked for. The message never quotes the payload, which + can hold a user's query or generated SQL. + """ + + +@dataclasses.dataclass(frozen=True) +class SseEvent: + """One event from a ``text/event-stream`` response.""" + + data: str + """The ``data`` lines of the event joined with LF.""" + + event: str = _DEFAULT_EVENT_NAME + """The ``event`` field, or ``'message'`` when the server sent none.""" + + id: str | None = None + """The most recent ``id`` field seen on the stream, if any.""" + + @property + def is_done(self) -> bool: + """Whether this event ends the stream normally. + + Cortex sends ``event: done`` together with ``data: [DONE]``. Either half on + its own counts, so a server that drops one of them still terminates + cleanly. + """ + return self.event == _DONE_EVENT_NAME or self.data.strip() == _DONE_DATA + + def json_data(self) -> Any: + """Parses ``data`` as JSON. + + Returns: + The decoded JSON value. + + Raises: + SseParseError: ``data`` is not valid JSON. + """ + try: + return json.loads(self.data) + except json.JSONDecodeError as e: + raise SseParseError( + f'SSE event {self.event!r} carries data that is not JSON (error at' + f' position {e.pos}).' + ) from e + + +class SseParser: + """Turns byte chunks into ``SseEvent``s, holding partial input in between. + + Feed every chunk to ``feed`` and the events it completes come back in order; + call ``close`` at the end of the stream to flush the decoder and drop any + unterminated event. + """ + + def __init__(self, *, max_event_bytes: int = _DEFAULT_MAX_EVENT_BYTES): + """Initializes the parser. + + Args: + max_event_bytes: Upper bound on the bytes buffered for one event. A + stream that exceeds it raises ``SseParseError`` from ``feed``, so a + single oversized event cannot grow memory without limit. + """ + self._max_event_bytes = max_event_bytes + # 'replace' rather than 'strict': one bad byte inside a multi-megabyte + # result set should not fail the whole turn, and the payload is JSON text + # whose structure survives a replacement character. + self._decoder = codecs.getincrementaldecoder('utf-8')(errors='replace') + self._pending = '' + self._at_start = True + self._buffered_bytes = 0 + self._event_name: str | None = None + self._data_lines: list[str] = [] + self._last_id: str | None = None + + def feed(self, chunk: bytes) -> list[SseEvent]: + """Consumes one chunk and returns the events it completed, in order. + + Args: + chunk: The next bytes of the response body. + + Returns: + The events whose terminating blank line arrived in this chunk. + + Raises: + SseParseError: The event being buffered exceeds ``max_event_bytes``. + """ + self._buffered_bytes += len(chunk) + events = self._consume(self._decoder.decode(chunk), final=False) + if self._buffered_bytes > self._max_event_bytes: + raise SseParseError( + 'SSE event exceeds the maximum size of' + f' {self._max_event_bytes} bytes.' + ) + return events + + def close(self) -> list[SseEvent]: + """Ends the stream and returns any events completed by the final bytes. + + An event the stream ended in the middle of, before its blank line, is + discarded as the SSE specification requires: a truncated final event is + not evidence that the run finished. + """ + events = self._consume(self._decoder.decode(b'', final=True), final=True) + self._pending = '' + self._reset_event() + return events + + def _consume(self, text: str, *, final: bool) -> list[SseEvent]: + if self._at_start and text: + text = text.removeprefix(_BOM) + self._at_start = False + self._pending += text + events: list[SseEvent] = [] + while (line := self._pop_line(final=final)) is not None: + event = self._process_line(line) + if event is not None: + events.append(event) + return events + + def _pop_line(self, *, final: bool) -> str | None: + pending = self._pending + cr = pending.find('\r') + lf = pending.find('\n') + if cr == -1 and lf == -1: + return None + if cr != -1 and (lf == -1 or cr < lf): + end = cr + if end + 1 == len(pending) and not final: + # A CR at the very end may be the first half of a CRLF split across + # chunks; wait for the next chunk rather than dispatch on it now and + # read the LF as an extra blank line later. + return None + skip = 2 if pending[end + 1 : end + 2] == '\n' else 1 + else: + end = lf + skip = 1 + self._pending = pending[end + skip :] + return pending[:end] + + def _process_line(self, line: str) -> SseEvent | None: + if not line: + return self._dispatch() + if line.startswith(':'): + return None + field, sep, value = line.partition(':') + if sep and value.startswith(' '): + value = value[1:] + if field == 'event': + self._event_name = value + elif field == 'data': + self._data_lines.append(value) + elif field == 'id' and '\x00' not in value: + self._last_id = value + # `retry` and unknown fields are ignored, as the specification requires. + return None + + def _dispatch(self) -> SseEvent | None: + if not self._data_lines: + # A block without data is not an event; its event name is dropped too. + self._reset_event() + return None + event = SseEvent( + data='\n'.join(self._data_lines), + event=self._event_name or _DEFAULT_EVENT_NAME, + id=self._last_id, + ) + self._reset_event() + return event + + def _reset_event(self) -> None: + self._event_name = None + self._data_lines = [] + # Whatever follows the blank line already belongs to the next event. + self._buffered_bytes = len(self._pending.encode('utf-8')) + + +async def iter_sse_events( + chunks: AsyncIterable[bytes], + *, + max_event_bytes: int = _DEFAULT_MAX_EVENT_BYTES, +) -> AsyncGenerator[SseEvent, None]: + """Yields the events of a byte stream as they complete. + + Args: + chunks: The response body, such as ``httpx.Response.aiter_bytes()``. + max_event_bytes: See ``SseParser``. + + Yields: + Each complete event, in stream order. Reading stops when ``chunks`` does; + deciding that a ``done`` event ends the stream is the caller's job. + + Raises: + SseParseError: An event exceeds ``max_event_bytes``. + """ + parser = SseParser(max_event_bytes=max_event_bytes) + async for chunk in chunks: + for event in parser.feed(chunk): + yield event + for event in parser.close(): + yield event diff --git a/tests/unittests/labs/snowflake/__init__.py b/tests/unittests/labs/snowflake/__init__.py new file mode 100644 index 0000000000..58d482ea38 --- /dev/null +++ b/tests/unittests/labs/snowflake/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittests/labs/snowflake/test_client.py b/tests/unittests/labs/snowflake/test_client.py new file mode 100644 index 0000000000..af022231b9 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_client.py @@ -0,0 +1,518 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Snowflake Cortex REST client. + +Drives ``SnowflakeCortexClient`` against an ``httpx.MockTransport`` standing in +for Snowflake, and verifies the requests it sends, the events it streams, and +the errors it raises when Snowflake or the network misbehaves. +""" + +from __future__ import annotations + +import json +from typing import Any +from typing import AsyncIterator +from typing import Callable +from unittest.mock import MagicMock + +from google.adk.labs.snowflake._client import CortexApiError +from google.adk.labs.snowflake._client import CortexTransportError +from google.adk.labs.snowflake._client import SnowflakeCortexClient +from google.adk.labs.snowflake._sse_parser import SseEvent +import httpx +import pytest + +_TOKEN = 'pat-secret-token-value' +_ACCOUNT_URL = 'https://acct.snowflakecomputing.com' +_RUN_PATH = ( + '/api/v2/databases/SALES_DB/schemas/ANALYTICS/agents/SALES_AGENT:run' +) +_STREAM = ( + b'event: metadata\n' + b'data: {"metadata":{"role":"user","message_id":455}}\n\n' + b'event: response.text.delta\n' + b'data: {"content_index":0,"sequence_number":1,"text":"Hi"}\n\n' + b'event: done\ndata: [DONE]\n\n' +) +# Where the first complete event ends: a cut here loses nothing. +_FIRST_EVENT_END = _STREAM.index(b'\n\n') + 2 + +_Handler = Callable[[httpx.Request], httpx.Response] + + +def _bearer_headers(ctx: Any) -> dict[str, str]: + del ctx + return { + 'Authorization': f'Bearer {_TOKEN}', + 'X-Snowflake-Authorization-Token-Type': 'PROGRAMMATIC_ACCESS_TOKEN', + } + + +class _Chunks(httpx.AsyncByteStream): + """A scripted response body: some chunks, then optionally a failure.""" + + def __init__( + self, chunks: list[bytes], *, then_raise: Exception | None = None + ): + self._chunks = chunks + self._then_raise = then_raise + self.closed = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + if self._then_raise is not None: + raise self._then_raise + + async def aclose(self) -> None: + self.closed = True + + +def _sse_response(body: httpx.AsyncByteStream) -> httpx.Response: + return httpx.Response( + 200, stream=body, headers={'content-type': 'text/event-stream'} + ) + + +def _make_client( + handler: _Handler, + *, + header_provider: Callable[..., Any] = _bearer_headers, + http_client: httpx.AsyncClient | None = None, +) -> tuple[SnowflakeCortexClient, list[httpx.Request]]: + """A client wired to `handler`, plus the list of requests it received.""" + requests: list[httpx.Request] = [] + + def _recording(request: httpx.Request) -> httpx.Response: + requests.append(request) + return handler(request) + + client = SnowflakeCortexClient( + account_url=_ACCOUNT_URL, + database='SALES_DB', + schema_name='ANALYTICS', + cortex_agent_name='SALES_AGENT', + header_provider=header_provider, + timeout=5.0, + http_client=http_client + or httpx.AsyncClient(transport=httpx.MockTransport(_recording)), + ) + return client, requests + + +def _ctx() -> MagicMock: + return MagicMock(name='ReadonlyContext') + + +# --- create_thread ------------------------------------------------------------ + + +async def test_create_thread_posts_origin_and_returns_the_id_as_text(): + """The Threads API call carries the auth headers and yields a string id.""" + client, requests = _make_client( + lambda request: httpx.Response(200, json={'thread_id': 1234567890}) + ) + + thread_id = await client.create_thread(_ctx()) + + (request,) = requests + assert thread_id == '1234567890' + assert request.url == f'{_ACCOUNT_URL}/api/v2/cortex/threads' + assert json.loads(request.content) == {'origin_application': 'google_adk'} + assert request.headers['authorization'] == f'Bearer {_TOKEN}' + assert request.headers['accept'] == 'application/json' + + +async def test_async_header_provider_is_awaited(): + """A coroutine provider works the same as a plain callable.""" + + async def _async_headers(ctx: Any) -> dict[str, str]: + del ctx + return {'Authorization': 'Bearer async-token'} + + client, requests = _make_client( + lambda request: httpx.Response(200, json={'thread_id': 7}), + header_provider=_async_headers, + ) + + await client.create_thread(_ctx()) + + assert requests[0].headers['authorization'] == 'Bearer async-token' + + +async def test_header_provider_receives_the_context(): + """The provider is handed the invocation's context to mint headers from.""" + seen: list[Any] = [] + + def _provider(ctx: Any) -> dict[str, str]: + seen.append(ctx) + return {} + + client, _ = _make_client( + lambda request: httpx.Response(200, json={'thread_id': 7}), + header_provider=_provider, + ) + ctx = _ctx() + + await client.create_thread(ctx) + + assert seen == [ctx] + + +@pytest.mark.parametrize( + 'payload', [{'thread_id': 'abc'}, {'thread_id': 0}, {'thread_id': True}, {}] +) +async def test_create_thread_rejects_an_unusable_id(payload: dict[str, Any]): + """A 2xx without a positive integer thread_id is an API error.""" + client, _ = _make_client(lambda request: httpx.Response(200, json=payload)) + + with pytest.raises(CortexApiError, match='no usable thread_id'): + await client.create_thread(_ctx()) + + +async def test_create_thread_surfaces_snowflake_error_details_without_token(): + """An error carries status, code and request id but never the credential.""" + client, _ = _make_client( + lambda request: httpx.Response( + 401, + json={ + 'code': '390144', + 'message': 'JWT token is invalid', + 'request_id': 'req-1', + }, + ) + ) + + with pytest.raises(CortexApiError) as info: + await client.create_thread(_ctx()) + + error = info.value + assert (error.status_code, error.snowflake_code, error.request_id) == ( + 401, + '390144', + 'req-1', + ) + assert 'JWT token is invalid' in str(error) + assert 'X-Snowflake-Authorization-Token-Type' in str(error) + assert _TOKEN not in str(error) + + +# --- run ---------------------------------------------------------------------- + + +async def test_run_posts_the_documented_body_and_streams_events(): + """The run request carries ids, the message and stream=true; events flow.""" + body = _Chunks([_STREAM[:37], _STREAM[37:]]) + client, requests = _make_client(lambda request: _sse_response(body)) + + async with client.run( + _ctx(), thread_id='1234567890', parent_message_id='0', text='hello' + ) as events: + received = [event async for event in events] + + (request,) = requests + assert request.url == f'{_ACCOUNT_URL}{_RUN_PATH}' + assert json.loads(request.content) == { + 'thread_id': 1234567890, + 'parent_message_id': 0, + 'messages': [ + {'role': 'user', 'content': [{'type': 'text', 'text': 'hello'}]} + ], + 'stream': True, + } + assert request.headers['accept'] == 'text/event-stream' + assert request.headers['accept-encoding'] == 'identity' + assert [e.event for e in received] == [ + 'metadata', + 'response.text.delta', + 'done', + ] + assert received[-1].is_done + assert body.closed + + +async def test_run_encodes_identifiers_and_tolerates_a_trailing_slash(): + """Object names are URL-encoded; a trailing slash on the account is fine.""" + requests: list[httpx.Request] = [] + + def _handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return _sse_response(_Chunks([b'event: done\ndata: [DONE]\n\n'])) + + client = SnowflakeCortexClient( + account_url=f'{_ACCOUNT_URL}/', + database='MY DB', + schema_name='S/1', + cortex_agent_name='AGENT', + header_provider=_bearer_headers, + timeout=5.0, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_handler)), + ) + + async with client.run( + _ctx(), thread_id=1, parent_message_id=0, text='x' + ) as events: + async for _ in events: + pass + + assert str(requests[0].url) == ( + f'{_ACCOUNT_URL}/api/v2/databases/MY%20DB/schemas/S%2F1/agents/AGENT:run' + ) + + +@pytest.mark.parametrize( + 'ids', + [ + {'thread_id': '0', 'parent_message_id': '0'}, + {'thread_id': '12a', 'parent_message_id': '0'}, + {'thread_id': '1', 'parent_message_id': '-1'}, + {'thread_id': '1', 'parent_message_id': ' 1'}, + {'thread_id': str(10**38), 'parent_message_id': '0'}, + {'thread_id': True, 'parent_message_id': '0'}, + ], +) +async def test_run_rejects_invalid_ids_before_sending(ids: dict[str, Any]): + """A cursor that is not a strict Snowflake id never reaches the network.""" + client, requests = _make_client(lambda request: _sse_response(_Chunks([]))) + + with pytest.raises(ValueError, match='decimal integer') as info: + async with client.run(_ctx(), text='x', **ids): + pass + + assert requests == [] + assert '12a' not in str(info.value) + + +@pytest.mark.parametrize('status', [401, 403, 429, 500, 503]) +async def test_run_raises_on_error_status_before_yielding(status: int): + """A non-2xx answer fails the run with its status before any event.""" + client, _ = _make_client( + lambda request: httpx.Response(status, json={'message': 'nope'}) + ) + + with pytest.raises(CortexApiError) as info: + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ): + pytest.fail('the stream must not open on an error status') + + assert info.value.status_code == status + assert 'nope' in str(info.value) + + +async def test_run_rejects_a_non_event_stream_answer(): + """A 200 that is not text/event-stream cannot be a run.""" + client, _ = _make_client( + lambda request: httpx.Response(200, json={'content': []}) + ) + + with pytest.raises(CortexApiError, match='instead of text/event-stream'): + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ): + pytest.fail('the stream must not open without an event stream') + + +async def test_connect_timeout_is_a_transport_error(): + """Not reaching Snowflake in time is reported as a timeout.""" + + def _timeout(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectTimeout('slow', request=request) + + client, _ = _make_client(_timeout) + + with pytest.raises(CortexTransportError, match='in time') as info: + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ): + pass + + assert info.value.timed_out is True + + +async def test_connection_failure_is_a_transport_error(): + """A refused connection is reported as unreachable, not as a timeout.""" + + def _refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('refused', request=request) + + client, _ = _make_client(_refuse) + + with pytest.raises( + CortexTransportError, match='could not be reached' + ) as info: + await client.create_thread(_ctx()) + + assert info.value.timed_out is False + + +async def test_stream_cut_mid_run_is_a_transport_error(): + """A connection dropping between chunks fails the run, not just ends it.""" + body = _Chunks([_STREAM[:37]], then_raise=httpx.ReadError('dropped')) + client, _ = _make_client(lambda request: _sse_response(body)) + + with pytest.raises(CortexTransportError, match='dropped') as info: + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ) as events: + async for _ in events: + pass + + assert info.value.timed_out is False + assert body.closed + + +async def test_read_timeout_mid_run_is_a_timeout(): + """Snowflake going quiet between chunks is reported as a timeout.""" + body = _Chunks([_STREAM[:37]], then_raise=httpx.ReadTimeout('quiet')) + client, _ = _make_client(lambda request: _sse_response(body)) + + with pytest.raises(CortexTransportError) as info: + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ) as events: + async for _ in events: + pass + + assert info.value.timed_out is True + + +async def test_stream_ending_without_done_is_left_to_the_caller(): + """A clean close before `[DONE]` yields what arrived and stops.""" + body = _Chunks([_STREAM[:_FIRST_EVENT_END]]) + client, _ = _make_client(lambda request: _sse_response(body)) + + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ) as events: + received = [event async for event in events] + + assert [e.event for e in received] == ['metadata'] + assert not any(e.is_done for e in received) + + +async def test_leaving_the_run_early_closes_the_upstream_response(): + """A consumer that stops reading releases the Snowflake connection.""" + body = _Chunks([_STREAM]) + client, _ = _make_client(lambda request: _sse_response(body)) + + async with client.run( + _ctx(), thread_id='1', parent_message_id='0', text='x' + ) as events: + async for event in events: + first = event + break + + assert isinstance(first, SseEvent) + assert body.closed + + +# --- cancel_run and lifecycle ------------------------------------------------- + + +async def test_cancel_run_posts_to_the_cancel_endpoint(): + """A cancel is a POST to the documented run cancel path.""" + client, requests = _make_client(lambda request: httpx.Response(200)) + + cancelled = await client.cancel_run(_ctx(), '123-455') + + (request,) = requests + assert cancelled is True + assert request.method == 'POST' + assert ( + request.url == f'{_ACCOUNT_URL}/api/v2/cortex/agent/runs/123-455/cancel' + ) + assert request.headers['authorization'] == f'Bearer {_TOKEN}' + + +async def test_cancel_of_a_finished_run_is_reported_not_raised(): + """Snowflake answering 409 for a run that already ended is not an error.""" + client, requests = _make_client( + lambda request: httpx.Response( + 409, json={'code': '390201', 'message': 'run already completed'} + ) + ) + + cancelled = await client.cancel_run(_ctx(), '123-455') + + assert cancelled is False + assert len(requests) == 1 + + +async def test_cancel_run_failure_is_reported_not_raised(): + """Cancel is best effort: an error status or a dead network is `False`.""" + client, _ = _make_client( + lambda request: httpx.Response(500, json={'message': 'busy'}) + ) + + def _refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('refused', request=request) + + unreachable, _ = _make_client(_refuse) + + assert await client.cancel_run(_ctx(), '1-1') is False + assert await unreachable.cancel_run(_ctx(), '1-1') is False + + +async def test_aclose_leaves_a_shared_http_client_open(): + """The application's client is the application's to close.""" + shared = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={'thread_id': 1}) + ) + ) + client, _ = _make_client( + lambda request: httpx.Response(200), http_client=shared + ) + await client.create_thread(_ctx()) + + await client.aclose() + + assert shared.is_closed is False + await shared.aclose() + + +async def test_aclose_closes_an_owned_http_client( + monkeypatch: pytest.MonkeyPatch, +): + """Without a shared client, the one created on demand is closed.""" + created: list[httpx.AsyncClient] = [] + real_async_client = httpx.AsyncClient + + def _tracking(**kwargs: Any) -> httpx.AsyncClient: + instance = real_async_client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={'thread_id': 1}) + ), + **kwargs, + ) + created.append(instance) + return instance + + monkeypatch.setattr(httpx, 'AsyncClient', _tracking) + client = SnowflakeCortexClient( + account_url=_ACCOUNT_URL, + database='D', + schema_name='S', + cortex_agent_name='A', + header_provider=_bearer_headers, + timeout=5.0, + ) + await client.create_thread(_ctx()) + + await client.aclose() + + (owned,) = created + assert owned.is_closed is True diff --git a/tests/unittests/labs/snowflake/test_event_converter.py b/tests/unittests/labs/snowflake/test_event_converter.py new file mode 100644 index 0000000000..2dbeb68350 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_event_converter.py @@ -0,0 +1,759 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Cortex event converter. + +Verifies that Cortex Agents Run API events map onto ADK events as the design +prescribes: deltas only in SSE mode and without duplicates, server-side tool +use paired with its result, persisted payloads bounded, and one final event +rebuilt from the authoritative ``response``. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import MagicMock + +from google.adk.events.event import Event +from google.adk.labs.snowflake._event_converter import CortexEventConverter +from google.adk.labs.snowflake._event_converter import METADATA_KEY +from google.adk.labs.snowflake._event_converter import UnsupportedCortexEventError +from google.adk.labs.snowflake._sse_parser import SseEvent +from google.adk.labs.snowflake._sse_parser import SseParseError +import pytest + +_TOOL_USE_ID = 'tool-use-1' + + +def _sse(name: str, payload: Any) -> SseEvent: + return SseEvent(event=name, data=json.dumps(payload)) + + +def _make_converter( + *, + streaming: bool = True, + max_tool_result_bytes: int = 32 * 1024, + include_thinking: bool = False, + thread_id: str | None = None, +) -> CortexEventConverter: + ctx = MagicMock() + ctx.invocation_id = 'inv_1' + ctx.branch = 'main' + return CortexEventConverter( + ctx=ctx, + author='cortex', + streaming=streaming, + max_tool_result_bytes=max_tool_result_bytes, + include_thinking_in_final_event=include_thinking, + thread_id=thread_id, + ) + + +def _tool_use(**overrides: Any) -> SseEvent: + payload = { + 'client_side_execute': False, + 'content_index': 20, + 'input': {'semantic_model': 'SV', 'sql': 'SELECT 1'}, + 'name': 'system_execute_sql', + 'sequence_number': 946, + 'tool_use_id': _TOOL_USE_ID, + 'type': 'system_execute_sql', + } + payload.update(overrides) + return _sse('response.tool_use', payload) + + +def _tool_result( + *, status: str = 'success', content: list[Any] | None = None +) -> SseEvent: + return _sse( + 'response.tool_result', + { + 'content': ( + [{'json': {'query_id': 'q1', 'result_set': {}}, 'type': 'json'}] + if content is None + else content + ), + 'content_index': 21, + 'name': 'system_execute_sql', + 'sequence_number': 952, + 'status': status, + 'tool_use_id': _TOOL_USE_ID, + 'type': 'system_execute_sql', + }, + ) + + +def _final_response(**overrides: Any) -> SseEvent: + payload = { + 'content': [ + {'thinking': {'text': 'reasoning'}, 'type': 'thinking'}, + {'text': 'The answer.', 'type': 'text'}, + ], + 'metadata': { + 'assistant_message_id': 456, + 'run_id': 'run-1', + 'thread_id': 123, + 'usage': {'tokens_consumed': []}, + 'user_message_id': 455, + }, + 'role': 'assistant', + 'status': 'completed', + } + payload.update(overrides) + return _sse('response', payload) + + +def _cortex(event: Event) -> dict[str, Any]: + assert event.custom_metadata is not None + return event.custom_metadata[METADATA_KEY] + + +# --- ids and terminators ----------------------------------------------------- + + +def test_metadata_records_message_ids_as_strings_without_events(): + """Thread message ids are kept per role as decimal strings.""" + converter = _make_converter() + + events = converter.convert( + _sse('metadata', {'metadata': {'role': 'user', 'message_id': 123}}) + ) + converter.convert( + _sse('metadata', {'metadata': {'role': 'assistant', 'message_id': 456}}) + ) + + assert events == [] + assert converter.user_message_id == '123' + assert converter.assistant_message_id == '456' + + +@pytest.mark.parametrize( + 'terminator', + [SseEvent(event='done', data='[DONE]'), SseEvent(data='[DONE]')], +) +def test_done_marks_the_run_complete_without_an_event(terminator: SseEvent): + """The transport terminator only flips `is_done`.""" + converter = _make_converter() + + assert converter.convert(terminator) == [] + assert converter.is_done + + +def test_non_object_json_is_rejected(): + """A payload that is JSON but not an object is a protocol error.""" + converter = _make_converter() + + with pytest.raises(SseParseError, match='not an object'): + converter.convert(SseEvent(event='response.status', data='[1, 2]')) + + +def test_invalid_json_is_rejected(): + """Data that is not JSON at all is a protocol error.""" + converter = _make_converter() + + with pytest.raises(SseParseError, match='not JSON'): + converter.convert(SseEvent(event='response.status', data='nope')) + + +# --- progress and unknown events --------------------------------------------- + + +def test_status_streams_as_partial_metadata_in_sse_mode(): + """Progress is forwarded as a partial, non-persisted metadata event.""" + converter = _make_converter(streaming=True) + payload = {'message': 'Planning', 'sequence_number': 1, 'status': 'planning'} + + (event,) = converter.convert(_sse('response.status', payload)) + + assert event.partial is True + assert event.content is None + assert _cortex(event) == {'event': 'response.status', 'data': payload} + + +def test_status_is_silent_without_streaming(): + """A consumer that did not ask for SSE gets no partial events.""" + converter = _make_converter(streaming=False) + + events = converter.convert( + _sse('response.status', {'status': 'planning', 'sequence_number': 1}) + ) + + assert events == [] + + +@pytest.mark.parametrize( + 'name', + ['response.tool_result.status', 'response.tool_result.analyst.delta'], +) +def test_tool_progress_streams_as_partial_metadata(name: str): + """Tool execution progress is forwarded like any other progress.""" + converter = _make_converter(streaming=True) + + (event,) = converter.convert(_sse(name, {'tool_use_id': _TOOL_USE_ID})) + + assert event.partial is True + assert _cortex(event)['event'] == name + + +def test_unknown_event_is_forwarded_and_does_not_stop_the_run(): + """An event the converter does not know is passed through as `unknown`.""" + converter = _make_converter(streaming=True) + + (event,) = converter.convert(_sse('response.new_thing', {'x': 1})) + later = converter.convert( + _sse('response.text.delta', {'content_index': 0, 'text': 'ok'}) + ) + + assert _cortex(event) == { + 'unknown': {'event': 'response.new_thing', 'data': {'x': 1}} + } + assert len(later) == 1 + + +def test_unknown_event_is_silent_without_streaming(): + """Unknown events are progress-only and never persisted.""" + converter = _make_converter(streaming=False) + + assert converter.convert(_sse('response.new_thing', {'x': 1})) == [] + + +# --- deltas and completed blocks --------------------------------------------- + + +def test_text_deltas_stream_in_order_as_partial_text_parts(): + """Each text delta becomes one partial model text event.""" + converter = _make_converter(streaming=True) + + events = converter.convert( + _sse( + 'response.text.delta', + {'content_index': 1, 'sequence_number': 13, 'text': '17과'}, + ) + ) + converter.convert( + _sse( + 'response.text.delta', + {'content_index': 1, 'sequence_number': 14, 'text': ' 23'}, + ) + ) + + assert [e.content.parts[0].text for e in events] == ['17과', ' 23'] + assert all(e.partial for e in events) + assert all(e.content.role == 'model' for e in events) + + +def test_thinking_deltas_are_thought_parts(): + """Reasoning deltas are marked as thoughts so UIs can fold them.""" + converter = _make_converter(streaming=True) + + (event,) = converter.convert( + _sse( + 'response.thinking.delta', + {'content_index': 0, 'sequence_number': 2, 'text': 'hmm'}, + ) + ) + + assert event.partial is True + assert event.content.parts[0].thought is True + assert event.content.parts[0].text == 'hmm' + + +def test_duplicate_sequence_number_is_dropped(): + """A resent delta neither streams again nor doubles the buffered text.""" + converter = _make_converter(streaming=True) + delta = _sse( + 'response.text.delta', + {'content_index': 1, 'sequence_number': 13, 'text': 'once'}, + ) + + events = converter.convert(delta) + converter.convert(delta) + converter.convert(_final_response(content=[])) + + assert len(events) == 1 + assert converter.final_event().content.parts[0].text == 'once' + + +def test_deltas_without_streaming_still_feed_the_final_answer(): + """Without SSE nothing streams, but the buffer backs the final event.""" + converter = _make_converter(streaming=False) + + events = converter.convert( + _sse('response.text.delta', {'content_index': 1, 'text': 'buffered'}) + ) + converter.convert(_final_response(content=[])) + + assert events == [] + assert converter.final_event().content.parts[0].text == 'buffered' + + +def test_completed_block_replaces_its_deltas(): + """`response.text` is authoritative for its content index.""" + converter = _make_converter(streaming=False) + converter.convert( + _sse('response.text.delta', {'content_index': 1, 'text': 'drafty'}) + ) + + converter.convert( + _sse('response.text', {'content_index': 1, 'text': 'final'}) + ) + converter.convert(_final_response(content=[])) + + assert converter.final_event().content.parts[0].text == 'final' + + +# --- server-side tools -------------------------------------------------------- + + +def test_tool_use_becomes_a_model_function_call(): + """A server-side tool call is recorded as a real `FunctionCall`.""" + converter = _make_converter() + + (event,) = converter.convert(_tool_use()) + + call = event.content.parts[0].function_call + assert event.partial is None + assert event.author == 'cortex' + assert event.content.role == 'model' + assert (call.id, call.name) == (_TOOL_USE_ID, 'system_execute_sql') + assert call.args == {'semantic_model': 'SV', 'sql': 'SELECT 1'} + assert not event.is_final_response() + + +def test_duplicate_tool_use_is_dropped(): + """The same `tool_use_id` is recorded once.""" + converter = _make_converter() + + events = converter.convert(_tool_use()) + converter.convert(_tool_use()) + + assert len(events) == 1 + + +def test_non_dict_tool_input_is_wrapped(): + """`FunctionCall.args` must be a dict, so a scalar input is wrapped.""" + converter = _make_converter() + + (event,) = converter.convert(_tool_use(input='raw')) + + assert event.content.parts[0].function_call.args == {'input': 'raw'} + + +@pytest.mark.parametrize( + 'overrides', + [{'client_side_execute': True}, {'permission': {'kind': 'ask'}}], +) +def test_client_side_tools_and_permissions_are_unsupported(overrides: dict): + """Client-side execution and permission prompts end the turn explicitly.""" + converter = _make_converter() + + with pytest.raises(UnsupportedCortexEventError, match='not support'): + converter.convert(_tool_use(**overrides)) + + +def test_tool_result_becomes_a_function_response_from_the_tool(): + """The result is authored by the tool and paired by `tool_use_id`.""" + converter = _make_converter() + converter.convert(_tool_use()) + + (event,) = converter.convert(_tool_result()) + + response = event.content.parts[0].function_response + assert event.author == 'system_execute_sql' + assert event.content.role == 'user' + assert (response.id, response.name) == (_TOOL_USE_ID, 'system_execute_sql') + assert response.response == { + 'status': 'success', + 'content': [ + {'json': {'query_id': 'q1', 'result_set': {}}, 'type': 'json'} + ], + } + assert not event.is_final_response() + + +def test_duplicate_tool_result_is_dropped(): + """The same `tool_use_id` is answered once.""" + converter = _make_converter() + + events = converter.convert(_tool_result()) + converter.convert(_tool_result()) + + assert len(events) == 1 + + +def test_failed_tool_result_carries_the_error(): + """A tool error surfaces under `error` next to the raw content.""" + converter = _make_converter() + error = {'error_code': 'SQL_COMPILATION_ERROR', 'message': 'bad sql'} + + (event,) = converter.convert( + _tool_result( + status='error', + content=[{'json': {'error': error, 'sql': 'SELECT'}, 'type': 'json'}], + ) + ) + + response = event.content.parts[0].function_response.response + assert response['status'] == 'error' + assert response['error'] == error + + +def test_oversized_result_rows_are_dropped_but_shape_is_kept(): + """Past the size bound only the query id and column metadata survive.""" + converter = _make_converter(max_tool_result_bytes=400) + rows = [[str(i), 'Country A', 'Air Purifier'] for i in range(50)] + result_set = {'data': rows, 'resultSetMetaData': {'numRows': 50}} + + (event,) = converter.convert( + _tool_result( + content=[{ + 'json': {'query_id': 'q1', 'result_set': result_set}, + 'type': 'json', + }] + ) + ) + + response = event.content.parts[0].function_response.response + assert response['truncated'] is True + assert response['original_bytes'] > 400 + assert response['content'][0]['json']['result_set'] == { + 'resultSetMetaData': {'numRows': 50} + } + assert response['content'][0]['json']['query_id'] == 'q1' + + +def test_oversized_result_without_rows_keeps_its_key_sizes(): + """A big non-SQL result is reduced to block types and per-key sizes.""" + converter = _make_converter(max_tool_result_bytes=300) + semantic_context = { + 'semantic_model_name': 'sv_overview', + 'tables': [ + {'name': f't{i}', 'columns': ['a', 'b', 'c']} for i in range(40) + ], + 'verified_queries': [{'sql': 'SELECT 1'}] * 20, + } + + (event,) = converter.convert( + _tool_result(content=[{'json': semantic_context, 'type': 'json'}]) + ) + + response = event.content.parts[0].function_response.response + (block,) = response['content'] + assert response['truncated'] is True + assert block['type'] == 'json' + assert set(block['json_keys']) == { + 'semantic_model_name', + 'tables', + 'verified_queries', + } + assert ( + block['json_keys']['tables'] > block['json_keys']['semantic_model_name'] + ) + assert 'sv_overview' not in json.dumps(response) + + +def test_result_still_too_large_after_shaping_loses_its_content(): + """When even the key sizes do not fit, the content is emptied, not stored.""" + converter = _make_converter(max_tool_result_bytes=64) + + (event,) = converter.convert( + _tool_result(content=[{'text': 'x' * 500, 'type': 'text'}]) + ) + + response = event.content.parts[0].function_response.response + assert response['truncated'] is True + assert response['content'] == [] + assert response['status'] == 'success' + + +# --- final-event metadata ---------------------------------------------------- + + +def test_warning_streams_and_is_kept_for_the_final_event(): + """Warnings are both forwarded live and recorded on the final event.""" + converter = _make_converter(streaming=True) + warning = {'message': 'Semantic view is stale'} + + (partial,) = converter.convert(_sse('response.warning', warning)) + converter.convert(_final_response()) + + assert partial.partial is True + assert _cortex(converter.final_event())['warnings'] == [warning] + + +def test_annotations_and_suggested_queries_land_on_the_final_event(): + """Citations and follow-up questions are final-event metadata only.""" + converter = _make_converter(streaming=True) + annotation = {'type': 'cortex_search_citation', 'index': 0} + + events = converter.convert(_sse('response.text.annotation', annotation)) + events += converter.convert( + _sse( + 'response.suggested_queries', + {'suggested_queries': [{'query': 'And next year?'}]}, + ) + ) + converter.convert(_final_response()) + + metadata = _cortex(converter.final_event()) + assert events == [] + assert metadata['annotations'] == [annotation] + assert metadata['suggested_queries'] == [{'query': 'And next year?'}] + + +def test_analyst_suggestion_deltas_are_assembled_into_suggested_queries(): + """Per-index suggestion deltas from Cortex Analyst become follow-ups.""" + converter = _make_converter(streaming=True) + deltas = [ + {'index': 0, 'delta': 'Which region '}, + {'index': 0, 'delta': 'sold most?'}, + {'index': 1, 'delta': 'Compare years'}, + ] + + streamed = [] + for delta in deltas: + streamed += converter.convert( + _sse( + 'response.tool_result.analyst.delta', + { + 'tool_use_id': _TOOL_USE_ID, + 'tool_type': 'cortex_analyst_text2sql', + 'delta': {'suggestions': delta}, + }, + ) + ) + converter.convert(_final_response()) + + assert len(streamed) == 3 and all(e.partial for e in streamed) + assert _cortex(converter.final_event())['suggested_queries'] == [ + { + 'query': 'Which region sold most?', + 'source': 'cortex_analyst', + 'tool_use_id': _TOOL_USE_ID, + }, + { + 'query': 'Compare years', + 'source': 'cortex_analyst', + 'tool_use_id': _TOOL_USE_ID, + }, + ] + + +def test_analyst_text_deltas_do_not_become_suggestions(): + """Only the `suggestions` delta feeds the list; text and SQL deltas do not.""" + converter = _make_converter() + converter.convert( + _sse( + 'response.tool_result.analyst.delta', + { + 'tool_use_id': _TOOL_USE_ID, + 'delta': {'text': 'Looking', 'sql': 'SELECT 1'}, + }, + ) + ) + converter.convert(_final_response()) + + assert _cortex(converter.final_event())['suggested_queries'] == [] + + +def test_suggested_queries_fall_back_to_the_final_response(): + """Without a dedicated event, the final payload's block is used.""" + converter = _make_converter() + converter.convert( + _final_response( + content=[ + {'text': 'A', 'type': 'text'}, + { + 'suggested_queries': [{'query': 'B?'}], + 'type': 'suggested_queries', + }, + ] + ) + ) + + assert _cortex(converter.final_event())['suggested_queries'] == [ + {'query': 'B?'} + ] + + +def test_oversized_table_keeps_metadata_only(): + """A big table is reduced to its query id and column metadata.""" + converter = _make_converter(max_tool_result_bytes=200) + table = { + 'content_index': 2, + 'query_id': 'q1', + 'result_set': { + 'data': [['x'] * 10] * 30, + 'resultSetMetaData': {'numRows': 30}, + }, + 'title': 'By country', + } + + converter.convert(_sse('response.table', table)) + converter.convert(_final_response()) + + (bounded,) = _cortex(converter.final_event())['tables'] + assert bounded['truncated'] is True + assert bounded['result_set'] == {'resultSetMetaData': {'numRows': 30}} + assert bounded['title'] == 'By country' + + +def test_oversized_chart_drops_its_spec(): + """A big chart keeps everything but the serialized spec.""" + converter = _make_converter(max_tool_result_bytes=100) + + converter.convert( + _sse('response.chart', {'content_index': 3, 'chart_spec': 'x' * 500}) + ) + converter.convert(_final_response()) + + (bounded,) = _cortex(converter.final_event())['charts'] + assert bounded == {'content_index': 3, 'truncated': True} + + +def test_small_table_and_chart_are_kept_whole(): + """Under the bound, presentations are recorded as sent.""" + converter = _make_converter() + table = {'content_index': 2, 'result_set': {'data': [['1']]}} + chart = {'content_index': 3, 'chart_spec': '{}'} + + converter.convert(_sse('response.table', table)) + converter.convert(_sse('response.chart', chart)) + converter.convert(_final_response()) + + metadata = _cortex(converter.final_event()) + assert metadata['tables'] == [table] + assert metadata['charts'] == [chart] + + +# --- final event -------------------------------------------------------------- + + +def test_final_event_is_rebuilt_from_the_final_response(): + """The answer comes from `response.content`, not from the deltas.""" + converter = _make_converter(streaming=True) + converter.convert( + _sse('response.text.delta', {'content_index': 1, 'text': 'drafty'}) + ) + converter.convert(_final_response()) + + event = converter.final_event() + + assert event.partial is None + assert event.author == 'cortex' + assert event.content.role == 'model' + assert [p.text for p in event.content.parts] == ['The answer.'] + assert event.is_final_response() + assert _cortex(event)['run_id'] == 'run-1' + assert _cortex(event)['status'] == 'completed' + assert _cortex(event)['usage'] == {'tokens_consumed': []} + + +def test_final_event_omits_thinking_by_default(): + """Reasoning is not persisted with the final event unless asked for.""" + converter = _make_converter(include_thinking=False) + converter.convert(_final_response()) + + parts = converter.final_event().content.parts + + assert [p.thought for p in parts] == [None] + + +def test_final_event_includes_thinking_when_enabled(): + """With the option on, completed reasoning precedes the answer.""" + converter = _make_converter(include_thinking=True) + converter.convert(_final_response()) + + parts = converter.final_event().content.parts + + assert [(p.thought, p.text) for p in parts] == [ + (True, 'reasoning'), + (None, 'The answer.'), + ] + + +def test_final_event_carries_the_state_delta(): + """The cursor handed in is committed through `EventActions`.""" + converter = _make_converter() + converter.convert(_final_response()) + + event = converter.final_event(state_delta={'_snowflake_cortex_x': {'a': 1}}) + + assert event.actions.state_delta == {'_snowflake_cortex_x': {'a': 1}} + + +def test_run_id_is_derived_when_the_payload_lacks_one(): + """`{thread_id}-{user_message_id}` is the documented run id shape.""" + converter = _make_converter(thread_id='123') + converter.convert( + _sse('metadata', {'metadata': {'role': 'user', 'message_id': 455}}) + ) + converter.convert(_final_response(metadata={})) + + assert _cortex(converter.final_event())['run_id'] == '123-455' + + +def test_final_status_is_exposed_once_the_final_response_arrives(): + """The run's terminal status is readable so callers can gate the cursor.""" + converter = _make_converter() + + assert converter.final_status is None + converter.convert(_final_response(status='cancelled')) + assert converter.final_status == 'cancelled' + + +def test_final_response_backfills_ids(): + """Ids missing from `metadata` events are taken from the final payload.""" + converter = _make_converter() + + converter.convert(_final_response()) + + assert converter.thread_id == '123' + assert converter.assistant_message_id == '456' + + +def test_final_event_requires_a_final_response(): + """Without the authoritative payload there is nothing to record.""" + converter = _make_converter() + + assert not converter.has_final_response + with pytest.raises(ValueError, match='without a final response'): + converter.final_event() + + +# --- terminal error ----------------------------------------------------------- + + +def test_error_becomes_a_terminal_error_event(): + """A stream-level `error` ends the run with code and message.""" + converter = _make_converter() + payload = {'code': 'STREAM_TIMEOUT', 'message': 'took too long'} + + (event,) = converter.convert(_sse('error', payload)) + + assert event.partial is None + assert (event.error_code, event.error_message) == ( + 'STREAM_TIMEOUT', + 'took too long', + ) + assert _cortex(event) == {'event': 'error', 'data': payload} + assert event.is_final_response() + assert converter.failed + + +def test_error_without_fields_uses_defaults(): + """An empty error payload still produces a usable error event.""" + converter = _make_converter() + + (event,) = converter.convert(_sse('error', {})) + + assert event.error_code == 'SNOWFLAKE_CORTEX_ERROR' + assert event.error_message diff --git a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py new file mode 100644 index 0000000000..2a5b0173f9 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py @@ -0,0 +1,787 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for SnowflakeCortexAgent. + +Verifies the configuration surface (composition guards, credential exclusion +from ``repr`` and every serialization path, per-agent state keys) and the run +loop against a scripted Snowflake behind ``httpx.MockTransport``: thread +creation, cursor commit and continuation, SSE gating, tool trace, failure +paths that leave the cursor alone, and cancellation on disconnect. +""" + +# `_state_key` is the documented shape of the session state key; the tests +# check it directly once and otherwise observe it through `state_delta`. +# pylint: disable=protected-access + +from __future__ import annotations + +import functools +import json +from typing import Any +from typing import AsyncGenerator +from typing import AsyncIterator + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.agents.run_config import RunConfig +from google.adk.agents.run_config import StreamingMode +from google.adk.cli.utils.graph_serialization import serialize_agent +from google.adk.events.event import Event +from google.adk.labs.snowflake import SnowflakeCortexAgent +from google.adk.labs.snowflake._client import CortexApiError +from google.adk.labs.snowflake._client import CortexTransportError +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types as genai_types +import httpx +from pydantic import ValidationError +import pytest + +_TOKEN = 'pat-secret-token-value' +_ACCOUNT_URL = 'https://example.snowflakecomputing.com' + + +def _bearer_headers(ctx: ReadonlyContext, *, token: str) -> dict[str, str]: + del ctx + return {'Authorization': f'Bearer {token}'} + + +# A `functools.partial` rather than a closure, because its `repr` prints the +# bound token. That makes "the token is absent" a real check: it would appear +# if the field were not excluded from `repr` and serialization. +_HEADER_PROVIDER = functools.partial(_bearer_headers, token=_TOKEN) + + +def _sse(*events: tuple[str, Any]) -> bytes: + """Encodes `(event name, JSON payload)` pairs as one SSE byte stream.""" + return b''.join( + f'event: {name}\ndata: {json.dumps(payload)}\n\n'.encode('utf-8') + for name, payload in events + ) + + +_DONE = b'event: done\ndata: [DONE]\n\n' + + +def _run_stream( + *, + assistant_message_id: int | None = 456, + answer: str = 'Hello', + status: str = 'completed', + done: bool = True, +) -> bytes: + """A complete run: status, two text deltas, one SQL tool call, the answer.""" + events: list[tuple[str, Any]] = [ + ('metadata', {'metadata': {'role': 'user', 'message_id': 455}}), + ('response.status', {'status': 'planning', 'sequence_number': 1}), + ( + 'response.text.delta', + {'content_index': 0, 'sequence_number': 2, 'text': answer[:3]}, + ), + ( + 'response.text.delta', + {'content_index': 0, 'sequence_number': 3, 'text': answer[3:]}, + ), + ( + 'response.tool_use', + { + 'client_side_execute': False, + 'input': {'sql': 'SELECT 1'}, + 'name': 'system_execute_sql', + 'tool_use_id': 't1', + 'type': 'system_execute_sql', + }, + ), + ( + 'response.tool_result', + { + 'content': [{'json': {'query_id': 'q1'}, 'type': 'json'}], + 'name': 'system_execute_sql', + 'status': 'success', + 'tool_use_id': 't1', + 'type': 'system_execute_sql', + }, + ), + ] + metadata: dict[str, Any] = {'run_id': 'run-1', 'user_message_id': 455} + if assistant_message_id is not None: + events.append(( + 'metadata', + {'metadata': {'role': 'assistant', 'message_id': assistant_message_id}}, + )) + metadata['assistant_message_id'] = assistant_message_id + events.append(( + 'response', + { + 'content': [{'text': answer, 'type': 'text'}], + 'metadata': metadata, + 'status': status, + }, + )) + return _sse(*events) + (_DONE if done else b'') + + +class _Chunks(httpx.AsyncByteStream): + """A scripted response body that records whether it was closed.""" + + def __init__(self, body: bytes): + self._body = body + self.closed = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._body + + async def aclose(self) -> None: + self.closed = True + + +class _FakeSnowflake: + """A scripted Snowflake behind `httpx.MockTransport`, recording requests.""" + + def __init__( + self, + *, + run_body: bytes | None = None, + thread_id: int = 123, + run_status: int = 200, + cancel_status: int = 200, + ): + self.thread_id = thread_id + self.run_status = run_status + self.cancel_status = cancel_status + self.run_bodies = [run_body if run_body is not None else _run_stream()] + self.streams: list[_Chunks] = [] + self.requests: list[httpx.Request] = [] + + def http_client(self) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(self._handle)) + + def _handle(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + path = request.url.path + if path.endswith('/cortex/threads'): + return httpx.Response(200, json={'thread_id': self.thread_id}) + if path.endswith(':run'): + if self.run_status != 200: + return httpx.Response(self.run_status, json={'message': 'nope'}) + body = ( + self.run_bodies.pop(0) + if len(self.run_bodies) > 1 + else (self.run_bodies[0]) + ) + stream = _Chunks(body) + self.streams.append(stream) + return httpx.Response( + 200, stream=stream, headers={'content-type': 'text/event-stream'} + ) + if path.endswith('/cancel'): + return httpx.Response(self.cancel_status) + return httpx.Response(404) + + def paths(self, suffix: str) -> list[httpx.Request]: + return [r for r in self.requests if r.url.path.endswith(suffix)] + + +def _make_agent( + name: str = 'cortex', **overrides: object +) -> SnowflakeCortexAgent: + """A minimal SnowflakeCortexAgent pointing at a fake account.""" + fields: dict[str, object] = { + 'name': name, + 'account_url': _ACCOUNT_URL, + 'database': 'SALES_DB', + 'schema_name': 'ANALYTICS', + 'cortex_agent_name': 'SALES_AGENT', + 'header_provider': _HEADER_PROVIDER, + } + fields.update(overrides) + return SnowflakeCortexAgent(**fields) + + +class _StubChild(BaseAgent): + """A runnable ADK child agent.""" + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + yield Event(invocation_id=ctx.invocation_id, author=self.name) + + +async def _invocation_context( + agent: BaseAgent, + *, + text: str | None = 'hello', + state: dict[str, Any] | None = None, + streaming_mode: StreamingMode = StreamingMode.SSE, +) -> InvocationContext: + """A real InvocationContext rooted at `agent`.""" + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name='test_app', user_id='test_user', state=state + ) + return InvocationContext( + session_service=session_service, + invocation_id='inv_1', + agent=agent, + session=session, + user_content=( + genai_types.Content( + role='user', parts=[genai_types.Part.from_text(text=text)] + ) + if text is not None + else None + ), + run_config=RunConfig(streaming_mode=streaming_mode), + ) + + +async def _run(agent: BaseAgent, ctx: InvocationContext) -> list[Event]: + return [event async for event in agent.run_async(ctx)] + + +def _final(events: list[Event]) -> Event: + (final,) = [ + e + for e in events + if not e.partial + and e.content + and e.content.role == 'model' + and not e.get_function_calls() + ] + return final + + +# --- configuration ------------------------------------------------------------ + + +def test_standalone_agent_is_allowed(): + """An agent with neither parent nor children constructs cleanly.""" + agent = _make_agent() + + assert agent.parent_agent is None + assert agent.sub_agents == [] + + +def test_defaults_are_the_documented_values(): + """Options not passed take the documented defaults.""" + agent = _make_agent() + + assert agent.timeout == 900.0 + assert agent.cancel_on_disconnect is True + assert agent.max_tool_result_bytes == 32 * 1024 + assert agent.include_thinking_in_final_event is False + assert agent.http_client is None + + +@pytest.mark.parametrize('field', ['timeout', 'max_tool_result_bytes']) +def test_non_positive_bounds_are_rejected(field: str): + """A zero timeout or result size limit fails validation.""" + with pytest.raises(ValidationError, match='greater than 0'): + _make_agent(**{field: 0}) + + +def test_sub_agents_are_rejected(): + """Declaring `sub_agents` fails at construction.""" + child = _StubChild(name='reviewer') + + with pytest.raises(ValueError, match='sub_agents'): + _make_agent(sub_agents=[child]) + + +def test_using_as_sub_agent_is_rejected(): + """A parent listing this agent in `sub_agents` fails to construct.""" + agent = _make_agent() + + with pytest.raises(ValueError, match='root agent'): + BaseAgent(name='parent', sub_agents=[agent]) + + assert agent.parent_agent is None + + +async def test_sub_agents_added_after_construction_are_rejected_at_run(): + """Mutating `sub_agents` past validation still fails, at the first turn.""" + agent = _make_agent() + agent.sub_agents.append(_StubChild(name='late')) + ctx = await _invocation_context(agent) + + with pytest.raises(ValueError, match='sub_agents'): + await _run(agent, ctx) + + +async def test_sub_agents_added_by_clone_are_rejected_at_run(): + """A clone given `sub_agents` skips construction checks but cannot run.""" + agent = _make_agent() + cloned = agent.clone(update={'sub_agents': [_StubChild(name='late')]}) + ctx = await _invocation_context(cloned) + + with pytest.raises(ValueError, match='sub_agents'): + await _run(cloned, ctx) + + +def test_header_provider_is_hidden_from_repr(): + """`repr` shows neither the provider nor the token it carries.""" + agent = _make_agent() + + text = repr(agent) + + assert 'header_provider' not in text + assert _TOKEN not in text + + +def test_header_provider_is_excluded_from_model_dump(): + """`model_dump` omits the provider and the token it carries.""" + agent = _make_agent() + + dumped = agent.model_dump() + + assert 'header_provider' not in dumped + assert 'http_client' not in dumped + assert _TOKEN not in str(dumped) + assert dumped['cortex_agent_name'] == 'SALES_AGENT' + + +def test_header_provider_is_hidden_from_the_adk_web_agent_graph(): + """The `adk web` agent graph omits the provider and the token it carries.""" + agent = _make_agent() + + serialized = json.dumps(serialize_agent(agent), default=str) + + assert 'header_provider' not in serialized + assert _TOKEN not in serialized + + +async def test_header_provider_stays_callable_on_the_instance(): + """Exclusion from output leaves the provider itself in place.""" + agent = _make_agent() + ctx = await _invocation_context(agent) + + headers = agent.header_provider(ReadonlyContext(ctx)) + + assert headers == {'Authorization': f'Bearer {_TOKEN}'} + + +def test_clone_keeps_the_header_provider(): + """A clone can still authenticate: exclusion is from output, not copies.""" + agent = _make_agent() + + cloned = agent.clone(update={'name': 'copy'}) + + assert cloned.header_provider is agent.header_provider + + +def test_state_key_is_scoped_by_agent_name(): + """Two agents with different names keep separate Snowflake threads.""" + first = _make_agent(name='first') + second = _make_agent(name='second') + + assert first._state_key() != second._state_key() + assert first._state_key() == _make_agent(name='first')._state_key() + + +# --- first and second turn ---------------------------------------------------- + + +async def test_first_turn_creates_a_thread_and_commits_the_cursor(): + """Turn one creates a thread, runs from message 0 and stores the cursor.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + ctx = await _invocation_context(agent) + + events = await _run(agent, ctx) + + (thread_request,) = snowflake.paths('/cortex/threads') + (run_request,) = snowflake.paths(':run') + assert thread_request.headers['authorization'] == f'Bearer {_TOKEN}' + run_body = json.loads(run_request.content) + assert (run_body['thread_id'], run_body['parent_message_id']) == (123, 0) + assert run_body['messages'][0]['content'][0]['text'] == 'hello' + final = _final(events) + assert final.content.parts[0].text == 'Hello' + assert final.custom_metadata['snowflake_cortex']['run_id'] == 'run-1' + cursor = final.actions.state_delta['_snowflake_cortex_cortex'] + assert cursor['schema_version'] == 1 + assert cursor['resource_fingerprint'].startswith('sha256:') + assert (cursor['thread_id'], cursor['parent_message_id']) == ('123', '456') + + +async def test_user_message_id_is_never_stored(): + """Only the assistant message may be the next parent.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + assert '455' not in json.dumps(_final(events).actions.state_delta) + + +async def test_second_turn_continues_the_stored_thread(): + """With a cursor in state, no thread is created and the parent advances.""" + snowflake = _FakeSnowflake(run_body=_run_stream(assistant_message_id=789)) + agent = _make_agent(http_client=snowflake.http_client()) + first_turn = _FakeSnowflake() + seed = _make_agent(http_client=first_turn.http_client()) + stored = _final( + await _run(seed, await _invocation_context(seed)) + ).actions.state_delta + ctx = await _invocation_context(agent, text='and then?', state=stored) + + events = await _run(agent, ctx) + + assert snowflake.paths('/cortex/threads') == [] + run_body = json.loads(snowflake.paths(':run')[0].content) + assert (run_body['thread_id'], run_body['parent_message_id']) == (123, 456) + cursor = _final(events).actions.state_delta['_snowflake_cortex_cortex'] + assert (cursor['thread_id'], cursor['parent_message_id']) == ('123', '789') + + +async def test_runner_persists_the_cursor_between_turns(): + """Through the Runner the state delta lands in the session for turn two.""" + snowflake = _FakeSnowflake() + snowflake.run_bodies = [ + _run_stream(assistant_message_id=456), + _run_stream(assistant_message_id=789), + ] + agent = _make_agent(http_client=snowflake.http_client()) + session_service = InMemorySessionService() + session = await session_service.create_session(app_name='app', user_id='u') + runner = Runner(app_name='app', agent=agent, session_service=session_service) + + for text in ('first', 'second'): + async for _ in runner.run_async( + user_id='u', + session_id=session.id, + new_message=genai_types.Content( + role='user', parts=[genai_types.Part.from_text(text=text)] + ), + ): + pass + + session = await session_service.get_session( + app_name='app', user_id='u', session_id=session.id + ) + assert len(snowflake.paths('/cortex/threads')) == 1 + parents = [ + json.loads(r.content)['parent_message_id'] + for r in snowflake.paths(':run') + ] + assert parents == [0, 456] + cursor = session.state['_snowflake_cortex_cortex'] + assert (cursor['thread_id'], cursor['parent_message_id']) == ('123', '789') + + +async def test_two_agents_in_one_session_keep_separate_cursors(): + """Each agent's cursor lives under its own key.""" + first = _make_agent(name='first', http_client=_FakeSnowflake().http_client()) + second = _make_agent( + name='second', + http_client=_FakeSnowflake(thread_id=999).http_client(), + ) + state = _final( + await _run(first, await _invocation_context(first)) + ).actions.state_delta + state.update( + _final( + await _run(second, await _invocation_context(second, state=state)) + ).actions.state_delta + ) + + assert state['_snowflake_cortex_first']['thread_id'] == '123' + assert state['_snowflake_cortex_second']['thread_id'] == '999' + + +# --- streaming and tool trace ------------------------------------------------- + + +async def test_sse_mode_streams_partial_events(): + """With SSE streaming the deltas and progress arrive as partial events.""" + agent = _make_agent(http_client=_FakeSnowflake().http_client()) + ctx = await _invocation_context(agent, streaming_mode=StreamingMode.SSE) + + events = await _run(agent, ctx) + + partial_text = [ + e.content.parts[0].text + for e in events + if e.partial and e.content and e.content.parts + ] + assert partial_text == ['Hel', 'lo'] + assert any(e.partial and e.custom_metadata for e in events) + + +async def test_none_mode_yields_only_persisted_events(): + """Without SSE streaming nothing partial is yielded, the answer still is.""" + agent = _make_agent(http_client=_FakeSnowflake().http_client()) + ctx = await _invocation_context(agent, streaming_mode=StreamingMode.NONE) + + events = await _run(agent, ctx) + + assert not any(e.partial for e in events) + assert _final(events).content.parts[0].text == 'Hello' + + +async def test_tool_trace_is_recorded_as_function_call_and_response(): + """Server-side tool use shows up as a call by the agent and a response.""" + agent = _make_agent(http_client=_FakeSnowflake().http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + (call_event,) = [e for e in events if e.get_function_calls()] + (response_event,) = [e for e in events if e.get_function_responses()] + assert call_event.author == 'cortex' + assert call_event.get_function_calls()[0].args == {'sql': 'SELECT 1'} + assert response_event.author == 'system_execute_sql' + assert response_event.get_function_responses()[0].id == 't1' + assert not call_event.partial and not response_event.partial + + +async def test_reading_stops_at_done(): + """Bytes after `[DONE]` are never read, so a chatty server cannot stall.""" + body = _run_stream() + _sse(('response.status', {'status': 'late'})) * 3 + snowflake = _FakeSnowflake(run_body=body) + agent = _make_agent(http_client=snowflake.http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + assert not any( + e.custom_metadata + and e.custom_metadata['snowflake_cortex'].get('data', {}).get('status') + == 'late' + for e in events + ) + assert snowflake.streams[0].closed + + +# --- failures leave the cursor alone ------------------------------------------ + + +async def test_error_event_ends_the_turn_without_a_cursor_update(): + """A terminal `error` is surfaced as an error event and nothing is stored.""" + body = ( + _sse( + ('metadata', {'metadata': {'role': 'user', 'message_id': 455}}), + ('error', {'code': 'STREAM_TIMEOUT', 'message': 'took too long'}), + ) + + _DONE + ) + agent = _make_agent(http_client=_FakeSnowflake(run_body=body).http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + (error_event,) = [e for e in events if e.error_code] + assert error_event.error_code == 'STREAM_TIMEOUT' + assert all(not e.actions.state_delta for e in events) + + +async def test_stream_cut_before_the_final_response_fails_and_keeps_the_cursor(): + """A stream that ends early is a failure, not a truncated answer.""" + body = _run_stream()[: _run_stream().index(b'event: response\n')] + agent = _make_agent(http_client=_FakeSnowflake(run_body=body).http_client()) + + with pytest.raises(CortexTransportError, match='before the final response'): + await _run(agent, await _invocation_context(agent)) + + +async def test_http_error_on_run_is_raised(): + """Snowflake refusing the run surfaces as an API error with its status.""" + agent = _make_agent(http_client=_FakeSnowflake(run_status=401).http_client()) + + with pytest.raises(CortexApiError) as info: + await _run(agent, await _invocation_context(agent)) + + assert info.value.status_code == 401 + + +async def test_final_response_without_done_still_completes_the_turn(): + """`[DONE]` is a compatibility sentinel; the final `response` closes the turn.""" + body = _run_stream(done=False) + agent = _make_agent(http_client=_FakeSnowflake(run_body=body).http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + final = _final(events) + assert final.content.parts[0].text == 'Hello' + cursor = final.actions.state_delta['_snowflake_cortex_cortex'] + assert cursor['parent_message_id'] == '456' + + +@pytest.mark.parametrize('status', ['cancelled', 'timed_out']) +async def test_non_completed_final_status_does_not_commit_the_cursor(status): + """Only a `completed` run may become the parent of the next turn.""" + body = _run_stream(status=status) + agent = _make_agent(http_client=_FakeSnowflake(run_body=body).http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + final = _final(events) + assert final.custom_metadata['snowflake_cortex']['status'] == status + assert final.actions.state_delta == {} + + +async def test_missing_assistant_id_skips_the_cursor_update(): + """Without an assistant message id there is no safe parent to store.""" + body = _run_stream(assistant_message_id=None) + agent = _make_agent(http_client=_FakeSnowflake(run_body=body).http_client()) + + events = await _run(agent, await _invocation_context(agent)) + + assert _final(events).actions.state_delta == {} + + +async def test_cursor_from_another_cortex_agent_is_refused(): + """A cursor whose fingerprint differs is never continued.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + other = _make_agent(cortex_agent_name='OTHER_AGENT') + stored = { + '_snowflake_cortex_cortex': { + 'schema_version': 1, + 'resource_fingerprint': other._resource_fingerprint(), + 'thread_id': '123', + 'parent_message_id': '456', + } + } + ctx = await _invocation_context(agent, state=stored) + + with pytest.raises(ValueError, match='different account') as info: + await _run(agent, ctx) + + assert snowflake.requests == [] + assert '123' not in str(info.value) + + +@pytest.mark.parametrize( + 'stored', + ['garbage', {'schema_version': 2}, {'schema_version': 1, 'thread_id': 1}], +) +async def test_malformed_cursor_is_refused(stored: Any): + """State that is not a cursor this version understands is an error.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + if isinstance(stored, dict) and 'thread_id' in stored: + stored = { + **stored, + 'resource_fingerprint': agent._resource_fingerprint(), + 'parent_message_id': '0', + } + ctx = await _invocation_context( + agent, state={'_snowflake_cortex_cortex': stored} + ) + + with pytest.raises(ValueError, match='_snowflake_cortex_cortex'): + await _run(agent, ctx) + + assert snowflake.requests == [] + + +async def test_missing_user_text_is_rejected_before_any_request(): + """A turn without text cannot be sent to Snowflake.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + ctx = await _invocation_context(agent, text=None) + + with pytest.raises(ValueError, match='text message'): + await _run(agent, ctx) + + assert snowflake.requests == [] + + +# --- disconnect --------------------------------------------------------------- + + +async def test_disconnect_closes_upstream_and_cancels_the_run(): + """A consumer that stops reading releases Snowflake and cancels the run.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + generator = agent.run_async(await _invocation_context(agent)) + + first = await generator.__anext__() + await generator.aclose() + + assert first.partial is True + assert snowflake.streams[0].closed + (cancel,) = snowflake.paths('/cancel') + assert cancel.url.path.endswith('/runs/123-455/cancel') + + +async def test_disconnect_without_the_option_does_not_cancel(): + """`cancel_on_disconnect=False` only closes the connection.""" + snowflake = _FakeSnowflake() + agent = _make_agent( + cancel_on_disconnect=False, http_client=snowflake.http_client() + ) + generator = agent.run_async(await _invocation_context(agent)) + + await generator.__anext__() + await generator.aclose() + + assert snowflake.streams[0].closed + assert snowflake.paths('/cancel') == [] + + +async def test_disconnect_before_the_user_message_is_acknowledged_does_not_cancel(): + """Without a user message id there is no run id, so nothing is cancelled.""" + body = _sse(('response.status', {'status': 'planning', 'sequence_number': 1})) + body += _run_stream() + snowflake = _FakeSnowflake(run_body=body) + agent = _make_agent(http_client=snowflake.http_client()) + generator = agent.run_async(await _invocation_context(agent)) + + first = await generator.__anext__() + await generator.aclose() + + assert first.custom_metadata['snowflake_cortex']['event'] == 'response.status' + assert snowflake.streams[0].closed + assert snowflake.paths('/cancel') == [] + + +async def test_disconnect_cancel_rejected_by_snowflake_is_swallowed(): + """A 409 from the cancel endpoint (run already over) does not surface.""" + snowflake = _FakeSnowflake(cancel_status=409) + agent = _make_agent(http_client=snowflake.http_client()) + generator = agent.run_async(await _invocation_context(agent)) + + await generator.__anext__() + await generator.aclose() + + assert len(snowflake.paths('/cancel')) == 1 + assert snowflake.streams[0].closed + + +async def test_disconnect_after_done_does_not_cancel(): + """Nothing is cancelled once Snowflake has finished the run.""" + snowflake = _FakeSnowflake() + agent = _make_agent(http_client=snowflake.http_client()) + + await _run(agent, await _invocation_context(agent)) + + assert snowflake.paths('/cancel') == [] + + +# --- lifecycle ---------------------------------------------------------------- + + +async def test_cleanup_leaves_a_shared_http_client_open(): + """The application's client is the application's to close.""" + snowflake = _FakeSnowflake() + shared = snowflake.http_client() + agent = _make_agent(http_client=shared) + await _run(agent, await _invocation_context(agent)) + + await agent.cleanup() + + assert shared.is_closed is False + await _run(agent, await _invocation_context(agent)) + await shared.aclose() diff --git a/tests/unittests/labs/snowflake/test_sse_parser.py b/tests/unittests/labs/snowflake/test_sse_parser.py new file mode 100644 index 0000000000..9f3793afa5 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_sse_parser.py @@ -0,0 +1,256 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Snowflake SSE parser. + +Verifies that ``text/event-stream`` bytes become events regardless of how the +network splits them, and that the Cortex terminal markers are recognized. +""" + +from __future__ import annotations + +from typing import AsyncIterator +from typing import Sequence + +from google.adk.labs.snowflake._sse_parser import iter_sse_events +from google.adk.labs.snowflake._sse_parser import SseEvent +from google.adk.labs.snowflake._sse_parser import SseParseError +from google.adk.labs.snowflake._sse_parser import SseParser +import pytest + +_TWO_EVENTS = ( + b'event: metadata\n' + b'data: {"metadata":{"role":"user","message_id":123}}\n' + b'\n' + b'event: response.text.delta\n' + b'data: {"content_index":1,"sequence_number":13,"text":"17\xea\xb3\xbc"}\n' + b'\n' +) + + +def _parse(stream: bytes, *, chunk_size: int | None = None) -> list[SseEvent]: + """Feeds `stream` in `chunk_size` pieces and closes the parser.""" + parser = SseParser() + events: list[SseEvent] = [] + size = chunk_size or len(stream) + for start in range(0, len(stream), size): + events.extend(parser.feed(stream[start : start + size])) + events.extend(parser.close()) + return events + + +async def _chunks(pieces: Sequence[bytes]) -> AsyncIterator[bytes]: + for piece in pieces: + yield piece + + +def test_single_event_carries_its_name_and_data(): + """An `event:` and `data:` pair becomes one event.""" + events = _parse(b'event: response.status\ndata: {"status":"planning"}\n\n') + + assert events == [ + SseEvent(event='response.status', data='{"status":"planning"}') + ] + + +def test_event_name_defaults_to_message(): + """A block with only `data:` uses the specification's default name.""" + events = _parse(b'data: hello\n\n') + + assert events == [SseEvent(event='message', data='hello')] + + +@pytest.mark.parametrize('chunk_size', [1, 2, 3, 7, 16, None]) +def test_chunk_boundaries_do_not_change_the_events(chunk_size: int | None): + """Any split of the byte stream yields the same events in the same order.""" + events = _parse(_TWO_EVENTS, chunk_size=chunk_size) + + assert [e.event for e in events] == ['metadata', 'response.text.delta'] + assert events[1].json_data()['text'] == '17과' + + +def test_multibyte_character_split_across_chunks_is_decoded(): + """A UTF-8 sequence cut between chunks is reassembled, not replaced.""" + parser = SseParser() + head, tail = b'data: 17\xea', b'\xb3\xbc\n\n' + + events = parser.feed(head) + parser.feed(tail) + parser.close() + + assert events == [SseEvent(data='17과')] + + +def test_invalid_utf8_is_replaced_rather_than_fatal(): + """A bad byte becomes U+FFFD so the rest of the stream still parses.""" + events = _parse(b'data: a\xffb\n\n') + + assert events == [SseEvent(data='a\ufffdb')] + + +@pytest.mark.parametrize('newline', [b'\n', b'\r\n', b'\r']) +def test_every_line_ending_frames_events(newline: bytes): + """LF, CRLF and bare CR all end lines and blank lines.""" + stream = newline.join([b'event: a', b'data: 1', b'', b'data: 2', b'', b'']) + + events = _parse(stream) + + assert events == [SseEvent(event='a', data='1'), SseEvent(data='2')] + + +def test_crlf_split_across_chunks_is_one_line_ending(): + """A CR ending one chunk and an LF starting the next do not add a line.""" + parser = SseParser() + + events = ( + parser.feed(b'event: a\r') + + parser.feed(b'\ndata: 1\r') + + parser.feed(b'\n\r') + + parser.feed(b'\n') + + parser.close() + ) + + assert events == [SseEvent(event='a', data='1')] + + +def test_multi_line_data_is_joined_with_lf(): + """Several `data:` lines form one payload separated by LF.""" + events = _parse(b'data: {\ndata: "a": 1\ndata: }\n\n') + + assert events == [SseEvent(data='{\n "a": 1\n}')] + assert events[0].json_data() == {'a': 1} + + +def test_one_leading_space_after_the_colon_is_dropped(): + """`data:x`, `data: x` and `data: x` differ only by the extra spaces.""" + events = _parse(b'data:x\n\ndata: x\n\ndata: x\n\n') + + assert [e.data for e in events] == ['x', 'x', ' x'] + + +def test_comments_retry_and_unknown_fields_are_ignored(): + """Only `event`, `data` and `id` shape the event.""" + events = _parse( + b': keep-alive\nretry: 3000\nfoo: bar\nid: 7\nevent: a\ndata: 1\n\n' + ) + + assert events == [SseEvent(event='a', data='1', id='7')] + + +def test_block_without_data_is_dropped_with_its_name(): + """An `event:` line followed by a blank line is not an event.""" + events = _parse(b'event: orphan\n\ndata: 1\n\n') + + assert events == [SseEvent(event='message', data='1')] + + +def test_leading_byte_order_mark_is_ignored(): + """A BOM at the start of the stream does not become part of a field name.""" + events = _parse(b'\xef\xbb\xbfevent: a\ndata: 1\n\n') + + assert events == [SseEvent(event='a', data='1')] + + +def test_unterminated_final_event_is_discarded(): + """An event cut off before its blank line is not dispatched at close.""" + events = _parse(b'event: a\ndata: 1\n\nevent: response\ndata: {"x":1}') + + assert events == [SseEvent(event='a', data='1')] + + +def test_unknown_event_names_pass_through_in_order(): + """The parser preserves names it does not know about.""" + events = _parse( + b'event: response.new_thing\ndata: 1\n\nevent: z\ndata: 2\n\n' + ) + + assert [e.event for e in events] == ['response.new_thing', 'z'] + + +def test_done_event_with_done_data_is_terminal(): + """`event: done` plus `data: [DONE]` is how Cortex ends a run.""" + (event,) = _parse(b'event: done\ndata: [DONE]\n\n') + + assert event.is_done + + +@pytest.mark.parametrize( + 'stream', + [b'event: done\ndata: {}\n\n', b'data: [DONE]\n\n', b'data: [DONE] \n\n'], +) +def test_either_done_marker_alone_is_terminal(stream: bytes): + """The name or the data alone marks the end, for servers that send one.""" + (event,) = _parse(stream) + + assert event.is_done + + +def test_ordinary_events_are_not_terminal(): + """A `response` event with a completed status still is not the end.""" + (event,) = _parse(b'event: response\ndata: {"status":"completed"}\n\n') + + assert not event.is_done + + +def test_error_event_is_parsed_like_any_other(): + """A terminal `error` event keeps its name and JSON body.""" + (event,) = _parse( + b'event: error\ndata: {"code":"STREAM_TIMEOUT","message":"m"}\n\n' + ) + + assert event.event == 'error' + assert event.json_data() == {'code': 'STREAM_TIMEOUT', 'message': 'm'} + + +def test_non_json_data_raises_without_quoting_the_payload(): + """`json_data` fails clearly and keeps the payload out of the message.""" + (event,) = _parse(b'event: response.text\ndata: SELECT secret FROM t\n\n') + + with pytest.raises(SseParseError, match='not JSON') as info: + event.json_data() + + assert 'secret' not in str(info.value) + + +def test_oversized_event_raises(): + """An event past `max_event_bytes` fails instead of growing memory.""" + parser = SseParser(max_event_bytes=32) + + with pytest.raises(SseParseError, match='maximum size'): + parser.feed(b'data: ' + b'x' * 64) + + +def test_size_limit_applies_per_event_not_per_stream(): + """Many small events do not add up against the limit.""" + parser = SseParser(max_event_bytes=32) + + events = parser.feed(b'data: 1\n\n' * 20) + + assert len(events) == 20 + + +async def test_async_iteration_yields_events_across_chunks(): + """`iter_sse_events` drives the parser over an async byte source.""" + pieces = [ + _TWO_EVENTS[:20], + _TWO_EVENTS[20:], + b'event: done\ndata: [DONE]\n\n', + ] + + events = [e async for e in iter_sse_events(_chunks(pieces))] + + assert [e.event for e in events] == [ + 'metadata', + 'response.text.delta', + 'done', + ] + assert events[-1].is_done