From 2b2e4564651813cb44a0749da77a7480fea7fb6d Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 21:50:53 +0900 Subject: [PATCH 01/12] feat(labs): add SnowflakeCortexAgent skeleton --- src/google/adk/labs/snowflake/README.md | 17 ++ src/google/adk/labs/snowflake/__init__.py | 46 ++++ .../labs/snowflake/_snowflake_cortex_agent.py | 198 +++++++++++++++ tests/unittests/labs/snowflake/__init__.py | 13 + .../snowflake/test_snowflake_cortex_agent.py | 227 ++++++++++++++++++ 5 files changed, 501 insertions(+) create mode 100644 src/google/adk/labs/snowflake/README.md create mode 100644 src/google/adk/labs/snowflake/__init__.py create mode 100644 src/google/adk/labs/snowflake/_snowflake_cortex_agent.py create mode 100644 tests/unittests/labs/snowflake/__init__.py create mode 100644 tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py diff --git a/src/google/adk/labs/snowflake/README.md b/src/google/adk/labs/snowflake/README.md new file mode 100644 index 00000000000..7461cc3012d --- /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 00000000000..a0d3138ebd6 --- /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/_snowflake_cortex_agent.py b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py new file mode 100644 index 00000000000..b9d11b7d22c --- /dev/null +++ b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py @@ -0,0 +1,198 @@ +# 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 + +from typing import Any +from typing import AsyncGenerator +from typing import Awaitable +from typing import Callable + +from pydantic import ConfigDict +from pydantic import Field +from typing_extensions import override + +from ...agents.base_agent import BaseAgent +from ...agents.invocation_context import InvocationContext +from ...agents.readonly_context import ReadonlyContext +from ...events.event import Event + +_STATE_KEY_PREFIX = '_snowflake_cortex_' + +_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.' +) + + +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. + The Snowflake thread and the last 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``. + + 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()}'} + + 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 an ``Authorization`` header. 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. + """ + + 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 truncated to its metadata + before it is recorded in a ``FunctionResponse`` event. Tool results are + persisted with the session, so this bounds how much a single SQL result set + 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. + """ + + @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 + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + self._validate_no_sub_agents() + raise NotImplementedError( + 'SnowflakeCortexAgent cannot run yet: the Cortex Agents Run API' + ' client is not implemented.' + ) + yield # AsyncGenerator requires having at least one yield statement diff --git a/tests/unittests/labs/snowflake/__init__.py b/tests/unittests/labs/snowflake/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /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_snowflake_cortex_agent.py b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py new file mode 100644 index 00000000000..a85820f0116 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py @@ -0,0 +1,227 @@ +# 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, and per-agent session state keys. +""" + +# `_state_key` is the only seam for the state key until the run loop lands and +# surfaces it in a `state_delta`. +# pylint: disable=protected-access + +from __future__ import annotations + +import functools +import json +from typing import AsyncGenerator + +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.cli.utils.graph_serialization import serialize_agent +from google.adk.events.event import Event +from google.adk.labs.snowflake import SnowflakeCortexAgent +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types as genai_types +from pydantic import ValidationError +import pytest + +_TOKEN = 'pat-secret-token-value' + + +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 _make_agent(name: str = 'cortex', **kwargs: object) -> SnowflakeCortexAgent: + """A minimal SnowflakeCortexAgent pointing at a fake account.""" + return SnowflakeCortexAgent( + name=name, + account_url='https://example.snowflakecomputing.com', + database='SALES_DB', + schema_name='ANALYTICS', + cortex_agent_name='SALES_AGENT', + header_provider=_HEADER_PROVIDER, + **kwargs, + ) + + +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) -> InvocationContext: + """A real InvocationContext rooted at `agent`.""" + session_service = InMemorySessionService() + return InvocationContext( + session_service=session_service, + invocation_id='inv_1', + agent=agent, + session=await session_service.create_session( + app_name='test_app', user_id='test_user' + ), + user_content=genai_types.Content( + role='user', parts=[genai_types.Part.from_text(text='hello')] + ), + run_config=RunConfig(), + ) + + +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 + + +@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'): + async for _ in agent.run_async(ctx): + pass + + +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'): + async for _ in cloned.run_async(ctx): + pass + + +async def test_running_is_not_implemented_yet(): + """Running a turn raises until the Snowflake client lands.""" + agent = _make_agent() + ctx = await _invocation_context(agent) + + with pytest.raises(NotImplementedError): + async for _ in agent.run_async(ctx): + pass + + +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 _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() From bae9061388c3d59f968d016c8759f8bd7e8a4265 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 21:51:03 +0900 Subject: [PATCH 02/12] feat(labs/snowflake): add SSE parser --- src/google/adk/labs/snowflake/_sse_parser.py | 245 +++++++++++++++++ .../labs/snowflake/test_sse_parser.py | 256 ++++++++++++++++++ 2 files changed, 501 insertions(+) create mode 100644 src/google/adk/labs/snowflake/_sse_parser.py create mode 100644 tests/unittests/labs/snowflake/test_sse_parser.py 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 00000000000..828bac76698 --- /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/test_sse_parser.py b/tests/unittests/labs/snowflake/test_sse_parser.py new file mode 100644 index 00000000000..9f3793afa56 --- /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 From f82c59c6402e0b033761d968ea6e8995b258b81a Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 22:03:46 +0900 Subject: [PATCH 03/12] feat(labs/snowflake): add Cortex event converter --- .../adk/labs/snowflake/_event_converter.py | 567 +++++++++++++++ .../labs/snowflake/test_event_converter.py | 665 ++++++++++++++++++ 2 files changed, 1232 insertions(+) create mode 100644 src/google/adk/labs/snowflake/_event_converter.py create mode 100644 tests/unittests/labs/snowflake/test_event_converter.py 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 00000000000..19467f500b5 --- /dev/null +++ b/src/google/adk/labs/snowflake/_event_converter.py @@ -0,0 +1,567 @@ +# 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 _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] = [] + 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_progress, + '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 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')) + ] + 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_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 a SQL result set is cut + # down to its shape (query id, column metadata) rather than stored whole. + 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 + 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/tests/unittests/labs/snowflake/test_event_converter.py b/tests/unittests/labs/snowflake/test_event_converter.py new file mode 100644 index 00000000000..c49263ac38b --- /dev/null +++ b/tests/unittests/labs/snowflake/test_event_converter.py @@ -0,0 +1,665 @@ +# 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_result_still_too_large_after_dropping_rows_loses_its_content(): + """When even the shape does 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_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_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 From 9a4c038e61e373d4bf947e1f6437489ce96cb2fe Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 22:12:34 +0900 Subject: [PATCH 04/12] feat(labs/snowflake): add Cortex REST client --- src/google/adk/labs/snowflake/_client.py | 427 +++++++++++++++ tests/unittests/labs/snowflake/test_client.py | 504 ++++++++++++++++++ 2 files changed, 931 insertions(+) create mode 100644 src/google/adk/labs/snowflake/_client.py create mode 100644 tests/unittests/labs/snowflake/test_client.py diff --git a/src/google/adk/labs/snowflake/_client.py b/src/google/adk/labs/snowflake/_client.py new file mode 100644 index 00000000000..4605ce039bb --- /dev/null +++ b/src/google/adk/labs/snowflake/_client.py @@ -0,0 +1,427 @@ +# 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: + logger.debug( + 'Best-effort cancel of a Snowflake Cortex run failed: %s', + type(e).__name__, + ) + 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/tests/unittests/labs/snowflake/test_client.py b/tests/unittests/labs/snowflake/test_client.py new file mode 100644 index 00000000000..738a602dd93 --- /dev/null +++ b/tests/unittests/labs/snowflake/test_client.py @@ -0,0 +1,504 @@ +# 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_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 From d1252d7aa177dc6318d87cebe2331e72ac10eb27 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 22:30:04 +0900 Subject: [PATCH 05/12] feat(labs/snowflake): run Cortex Agent turns and persist the thread cursor --- .../labs/snowflake/_snowflake_cortex_agent.py | 242 +++++++- .../snowflake/test_snowflake_cortex_agent.py | 569 ++++++++++++++++-- 2 files changed, 762 insertions(+), 49 deletions(-) diff --git a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py index b9d11b7d22c..c5168bd98fc 100644 --- a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py +++ b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py @@ -26,21 +26,34 @@ 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' @@ -54,15 +67,39 @@ ) +@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. - The Snowflake thread and the last 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``. + 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 @@ -77,7 +114,10 @@ class SnowflakeCortexAgent(BaseAgent): from google.adk.labs.snowflake import SnowflakeCortexAgent def bearer_headers(ctx: ReadonlyContext) -> dict[str, str]: - return {'Authorization': f'Bearer {load_snowflake_token()}'} + return { + 'Authorization': f'Bearer {load_snowflake_token()}', + 'X-Snowflake-Authorization-Token-Type': 'PROGRAMMATIC_ACCESS_TOKEN', + } root_agent = SnowflakeCortexAgent( name='sales_analyst', @@ -124,10 +164,21 @@ def bearer_headers(ctx: ReadonlyContext) -> dict[str, str]: ] = Field(exclude=True, repr=False) """Supplies the HTTP headers for each Snowflake request. - Typically an ``Authorization`` header. 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. + 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) @@ -161,6 +212,8 @@ def bearer_headers(ctx: ReadonlyContext) -> dict[str, str]: SSE mode. """ + _cortex_client: SnowflakeCortexClient | None = PrivateAttr(default=None) + @override def model_post_init(self, __context: Any) -> None: super().model_post_init(__context) @@ -186,13 +239,174 @@ def _state_key(self) -> str: # 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() - raise NotImplementedError( - 'SnowflakeCortexAgent cannot run yet: the Cortex Agents Run API' - ' client is not implemented.' + 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, ) - yield # AsyncGenerator requires having at least one yield statement + + 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.is_done and converter.has_final_response): + raise CortexTransportError( + 'The Snowflake stream ended before the run finished, so the answer' + ' is incomplete. The conversation cursor was left unchanged.', + timed_out=False, + ) + + state_delta: dict[str, Any] | None = None + if converter.assistant_message_id is not None: + # Only the assistant id may become the parent of the next turn; 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: + 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/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py index a85820f0116..22e18893f30 100644 --- a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py +++ b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py @@ -14,33 +14,44 @@ """Tests for SnowflakeCortexAgent. -Verifies the configuration surface: composition guards, credential exclusion -from ``repr`` and every serialization path, and per-agent session state keys. +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 only seam for the state key until the run loop lands and -# surfaces it in a `state_delta`. +# `_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]: @@ -54,19 +65,146 @@ def _bearer_headers(ctx: ReadonlyContext, *, token: str) -> dict[str, str]: _HEADER_PROVIDER = functools.partial(_bearer_headers, token=_TOKEN) -def _make_agent(name: str = 'cortex', **kwargs: object) -> SnowflakeCortexAgent: - """A minimal SnowflakeCortexAgent pointing at a fake account.""" - return SnowflakeCortexAgent( - name=name, - account_url='https://example.snowflakecomputing.com', - database='SALES_DB', - schema_name='ANALYTICS', - cortex_agent_name='SALES_AGENT', - header_provider=_HEADER_PROVIDER, - **kwargs, +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' +) -> 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': 'completed', + }, + )) + return _sse(*events) + _DONE + + +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, + ): + self.thread_id = thread_id + self.run_status = run_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(200) + 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.""" @@ -76,23 +214,53 @@ async def _run_async_impl( yield Event(invocation_id=ctx.invocation_id, author=self.name) -async def _invocation_context(agent: BaseAgent) -> InvocationContext: +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=await session_service.create_session( - app_name='test_app', user_id='test_user' + session=session, + user_content=( + genai_types.Content( + role='user', parts=[genai_types.Part.from_text(text=text)] + ) + if text is not None + else None ), - user_content=genai_types.Content( - role='user', parts=[genai_types.Part.from_text(text='hello')] - ), - run_config=RunConfig(), + 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() @@ -109,6 +277,7 @@ def test_defaults_are_the_documented_values(): 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']) @@ -143,8 +312,7 @@ async def test_sub_agents_added_after_construction_are_rejected_at_run(): ctx = await _invocation_context(agent) with pytest.raises(ValueError, match='sub_agents'): - async for _ in agent.run_async(ctx): - pass + await _run(agent, ctx) async def test_sub_agents_added_by_clone_are_rejected_at_run(): @@ -154,18 +322,7 @@ async def test_sub_agents_added_by_clone_are_rejected_at_run(): ctx = await _invocation_context(cloned) with pytest.raises(ValueError, match='sub_agents'): - async for _ in cloned.run_async(ctx): - pass - - -async def test_running_is_not_implemented_yet(): - """Running a turn raises until the Snowflake client lands.""" - agent = _make_agent() - ctx = await _invocation_context(agent) - - with pytest.raises(NotImplementedError): - async for _ in agent.run_async(ctx): - pass + await _run(cloned, ctx) def test_header_provider_is_hidden_from_repr(): @@ -185,6 +342,7 @@ def test_header_provider_is_excluded_from_model_dump(): 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' @@ -225,3 +383,344 @@ def test_state_key_is_scoped_by_agent_name(): 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_done_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 run finished'): + 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_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_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() From 24de518fbabc0a16b83a53b9e8bf57c3f75e0a67 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 22:36:35 +0900 Subject: [PATCH 06/12] docs(labs/snowflake): add SnowflakeCortexAgent guide and sample --- .../snowflake_cortex_agent/README.md | 110 ++++++ .../snowflake_cortex_agent/agent.py | 60 +++ docs/guides/README.md | 1 + .../snowflake/snowflake_cortex_agent/index.md | 343 ++++++++++++++++++ 4 files changed, 514 insertions(+) create mode 100644 contributing/samples/integrations/snowflake_cortex_agent/README.md create mode 100644 contributing/samples/integrations/snowflake_cortex_agent/agent.py create mode 100644 docs/guides/labs/snowflake/snowflake_cortex_agent/index.md 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 00000000000..e587a0bc720 --- /dev/null +++ b/contributing/samples/integrations/snowflake_cortex_agent/README.md @@ -0,0 +1,110 @@ +# 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=_required("SNOWFLAKE_ACCOUNT_URL"), + database=_required("SNOWFLAKE_DATABASE"), + schema_name=_required("SNOWFLAKE_SCHEMA"), + cortex_agent_name=_required("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: + +```python +def snowflake_headers(ctx: ReadonlyContext) -> dict[str, str]: + return { + "Authorization": f"Bearer {_required('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 00000000000..a0ec6e1c3bd --- /dev/null +++ b/contributing/samples/integrations/snowflake_cortex_agent/agent.py @@ -0,0 +1,60 @@ +# 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 + + +def _required(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError( + f"Set {name} in the environment or in a .env file next to agent.py." + ) + return value + + +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. + return { + "Authorization": f"Bearer {_required('SNOWFLAKE_TOKEN')}", + "X-Snowflake-Authorization-Token-Type": os.environ.get( + "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=_required("SNOWFLAKE_ACCOUNT_URL"), + database=_required("SNOWFLAKE_DATABASE"), + schema_name=_required("SNOWFLAKE_SCHEMA"), + cortex_agent_name=_required("SNOWFLAKE_CORTEX_AGENT"), + header_provider=snowflake_headers, +) diff --git a/docs/guides/README.md b/docs/guides/README.md index 454675fda77..f8d406aae6f 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 00000000000..2bc2445e696 --- /dev/null +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -0,0 +1,343 @@ +# 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 to their + shape, keeping the query id and column metadata but not the rows. +- When Snowflake sends its final `response` and the stream terminator, 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`. 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 run finished 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. 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 what happens to the Snowflake run when the ADK +consumer goes away mid-turn. With it enabled the agent sends one cancel request +for the run; the cancel is best effort and its failure is not raised, because +nobody is left to read the error. Disable it if you want abandoned runs to +finish and land in the thread anyway. + +`max_tool_result_bytes` protects the session store. Tool results are persisted +as `FunctionResponse` events and a SQL result set can be large, so a result +above the bound keeps its query id and column metadata and drops the rows; +tables and charts on the final event are bounded the same way. Raise it if your +application reads rows 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`. + +### 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 in the thread, and a + cancelled or abandoned run is still billed for 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. From 12c6225d7bbc6dd7c256fbbb541743e38ae0c290 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 23:14:28 +0900 Subject: [PATCH 07/12] fix(labs/snowflake): keep tool result shape when truncating --- .../snowflake/snowflake_cortex_agent/index.md | 18 ++++++---- .../adk/labs/snowflake/_event_converter.py | 31 +++++++++++++++-- .../labs/snowflake/_snowflake_cortex_agent.py | 8 ++--- .../labs/snowflake/test_event_converter.py | 34 +++++++++++++++++-- 4 files changed, 77 insertions(+), 14 deletions(-) diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md index 2bc2445e696..14d0d1392c5 100644 --- a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -85,8 +85,10 @@ ADK events: 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 to their - shape, keeping the query id and column metadata but not the rows. + 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` and the stream terminator, the agent yields one non-partial event with the answer text. Citations, warnings, suggested follow-up queries, token usage, tables and charts sit in @@ -158,10 +160,14 @@ nobody is left to read the error. Disable it if you want abandoned runs to finish and land in the thread anyway. `max_tool_result_bytes` protects the session store. Tool results are persisted -as `FunctionResponse` events and a SQL result set can be large, so a result -above the bound keeps its query id and column metadata and drops the rows; -tables and charts on the final event are bounded the same way. Raise it if your -application reads rows from the recorded events, lower it if session size +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 diff --git a/src/google/adk/labs/snowflake/_event_converter.py b/src/google/adk/labs/snowflake/_event_converter.py index 19467f500b5..0a1644e50f3 100644 --- a/src/google/adk/labs/snowflake/_event_converter.py +++ b/src/google/adk/labs/snowflake/_event_converter.py @@ -99,6 +99,25 @@ def _without_result_rows(item: Any) -> Any: 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') @@ -485,8 +504,10 @@ def _on_tool_result(self, name: str, payload: dict[str, Any]) -> list[Event]: def _bound_tool_result( self, response: dict[str, JsonValue] ) -> dict[str, JsonValue]: - # The response is persisted with the session, so a SQL result set is cut - # down to its shape (query id, column metadata) rather than stored whole. + # 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 @@ -501,6 +522,12 @@ def _bound_tool_result( } 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]: diff --git a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py index c5168bd98fc..f91f894682f 100644 --- a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py +++ b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py @@ -198,10 +198,10 @@ def bearer_headers(ctx: ReadonlyContext) -> dict[str, str]: 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 truncated to its metadata - before it is recorded in a ``FunctionResponse`` event. Tool results are - persisted with the session, so this bounds how much a single SQL result set - can grow it. + 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 diff --git a/tests/unittests/labs/snowflake/test_event_converter.py b/tests/unittests/labs/snowflake/test_event_converter.py index c49263ac38b..8ac046d1aa3 100644 --- a/tests/unittests/labs/snowflake/test_event_converter.py +++ b/tests/unittests/labs/snowflake/test_event_converter.py @@ -430,8 +430,38 @@ def test_oversized_result_rows_are_dropped_but_shape_is_kept(): assert response['content'][0]['json']['query_id'] == 'q1' -def test_result_still_too_large_after_dropping_rows_loses_its_content(): - """When even the shape does not fit, the content is emptied, not stored.""" +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( From 74f58192df098bd1dc98d1e34812d3a906714287 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 23:40:49 +0900 Subject: [PATCH 08/12] fix(labs/snowflake): make the sample importable without credentials --- .../snowflake_cortex_agent/README.md | 15 +++++--- .../snowflake_cortex_agent/agent.py | 38 +++++++++++++------ 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/contributing/samples/integrations/snowflake_cortex_agent/README.md b/contributing/samples/integrations/snowflake_cortex_agent/README.md index e587a0bc720..af1870db08b 100644 --- a/contributing/samples/integrations/snowflake_cortex_agent/README.md +++ b/contributing/samples/integrations/snowflake_cortex_agent/README.md @@ -77,10 +77,10 @@ provider: root_agent = SnowflakeCortexAgent( name="snowflake_cortex_analyst", description="Answers data questions by running a Snowflake Cortex Agent.", - account_url=_required("SNOWFLAKE_ACCOUNT_URL"), - database=_required("SNOWFLAKE_DATABASE"), - schema_name=_required("SNOWFLAKE_SCHEMA"), - cortex_agent_name=_required("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, ) ``` @@ -88,12 +88,15 @@ root_agent = SnowflakeCortexAgent( 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: +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 {_required('SNOWFLAKE_TOKEN')}", + "Authorization": f"Bearer {_env('SNOWFLAKE_TOKEN')}", "X-Snowflake-Authorization-Token-Type": os.environ.get( "SNOWFLAKE_TOKEN_TYPE", "PROGRAMMATIC_ACCESS_TOKEN" ), diff --git a/contributing/samples/integrations/snowflake_cortex_agent/agent.py b/contributing/samples/integrations/snowflake_cortex_agent/agent.py index a0ec6e1c3bd..3d1eb6bb81b 100644 --- a/contributing/samples/integrations/snowflake_cortex_agent/agent.py +++ b/contributing/samples/integrations/snowflake_cortex_agent/agent.py @@ -24,22 +24,38 @@ 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 _required(name: str) -> str: - value = os.environ.get(name) - if not value: +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( - f"Set {name} in the environment or in a .env file next to agent.py." + "Snowflake settings are missing:" + f" {', '.join(missing)}. Set them in the environment or in a .env" + " file next to agent.py." ) - return value 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 {_required('SNOWFLAKE_TOKEN')}", - "X-Snowflake-Authorization-Token-Type": os.environ.get( + "Authorization": f"Bearer {_env('SNOWFLAKE_TOKEN')}", + "X-Snowflake-Authorization-Token-Type": _env( "SNOWFLAKE_TOKEN_TYPE", "PROGRAMMATIC_ACCESS_TOKEN" ), } @@ -52,9 +68,9 @@ def snowflake_headers(ctx: ReadonlyContext) -> dict[str, str]: root_agent = SnowflakeCortexAgent( name="snowflake_cortex_analyst", description="Answers data questions by running a Snowflake Cortex Agent.", - account_url=_required("SNOWFLAKE_ACCOUNT_URL"), - database=_required("SNOWFLAKE_DATABASE"), - schema_name=_required("SNOWFLAKE_SCHEMA"), - cortex_agent_name=_required("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, ) From ab05048129469206a82bb409caadfe457d5bbdd9 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Fri, 4 Sep 2026 23:50:39 +0900 Subject: [PATCH 09/12] fix(labs/snowflake): commit the thread cursor only from a completed final response --- .../snowflake/snowflake_cortex_agent/index.md | 10 +++-- .../adk/labs/snowflake/_event_converter.py | 8 ++++ .../labs/snowflake/_snowflake_cortex_agent.py | 28 +++++++++---- .../labs/snowflake/test_event_converter.py | 9 +++++ .../snowflake/test_snowflake_cortex_agent.py | 40 ++++++++++++++++--- 5 files changed, 80 insertions(+), 15 deletions(-) diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md index 14d0d1392c5..065a0831350 100644 --- a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -89,14 +89,17 @@ ADK events: 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` and the stream terminator, the +- 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`. The ADK `Runner` applies it to the session, which +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 @@ -106,7 +109,8 @@ 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 run finished raises `CortexTransportError`, and a rejected request +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. diff --git a/src/google/adk/labs/snowflake/_event_converter.py b/src/google/adk/labs/snowflake/_event_converter.py index 0a1644e50f3..be18fd7c1ef 100644 --- a/src/google/adk/labs/snowflake/_event_converter.py +++ b/src/google/adk/labs/snowflake/_event_converter.py @@ -225,6 +225,14 @@ 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.""" diff --git a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py index f91f894682f..da4fe7054c1 100644 --- a/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py +++ b/src/google/adk/labs/snowflake/_snowflake_cortex_agent.py @@ -371,17 +371,26 @@ async def _run_async_impl( # 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.is_done and converter.has_final_response): + if not converter.has_final_response: raise CortexTransportError( - 'The Snowflake stream ended before the run finished, so the answer' - ' is incomplete. The conversation cursor was left unchanged.', + '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.assistant_message_id is not None: - # Only the assistant id may become the parent of the next turn; the - # user id would fork the thread. + 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( @@ -398,7 +407,12 @@ async def _cancel_abandoned_run( client: SnowflakeCortexClient, converter: CortexEventConverter, ) -> None: - if not self.cancel_on_disconnect or converter.is_done or converter.failed: + 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 diff --git a/tests/unittests/labs/snowflake/test_event_converter.py b/tests/unittests/labs/snowflake/test_event_converter.py index 8ac046d1aa3..8fb386c3806 100644 --- a/tests/unittests/labs/snowflake/test_event_converter.py +++ b/tests/unittests/labs/snowflake/test_event_converter.py @@ -646,6 +646,15 @@ def test_run_id_is_derived_when_the_payload_lacks_one(): 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() diff --git a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py index 22e18893f30..0a9431aec0a 100644 --- a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py +++ b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py @@ -77,7 +77,11 @@ def _sse(*events: tuple[str, Any]) -> bytes: def _run_stream( - *, assistant_message_id: int | None = 456, answer: str = 'Hello' + *, + 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]] = [ @@ -124,10 +128,10 @@ def _run_stream( { 'content': [{'text': answer, 'type': 'text'}], 'metadata': metadata, - 'status': 'completed', + 'status': status, }, )) - return _sse(*events) + _DONE + return _sse(*events) + (_DONE if done else b'') class _Chunks(httpx.AsyncByteStream): @@ -579,12 +583,12 @@ async def test_error_event_ends_the_turn_without_a_cursor_update(): assert all(not e.actions.state_delta for e in events) -async def test_stream_cut_before_done_fails_and_keeps_the_cursor(): +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 run finished'): + with pytest.raises(CortexTransportError, match='before the final response'): await _run(agent, await _invocation_context(agent)) @@ -598,6 +602,32 @@ async def test_http_error_on_run_is_raised(): 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) From 1119f3e50d43f1b894be6d782edfdbba5c573054 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Sat, 5 Sep 2026 00:15:07 +0900 Subject: [PATCH 10/12] fix(labs/snowflake): cover cancel edge cases and log refused cancels --- .../snowflake/snowflake_cortex_agent/index.md | 23 ++++++++----- src/google/adk/labs/snowflake/_client.py | 4 ++- tests/unittests/labs/snowflake/test_client.py | 14 ++++++++ .../snowflake/test_snowflake_cortex_agent.py | 33 ++++++++++++++++++- 4 files changed, 63 insertions(+), 11 deletions(-) diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md index 065a0831350..0051841a8eb 100644 --- a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -116,8 +116,11 @@ 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. Snowflake keeps whatever it had -already produced in the thread either way. +enabled, asks Snowflake once to cancel the run. In live tests Snowflake had +already marked the run completed by the time the cancel arrived, answering +`409 Agent run was already completed`: closing the stream is what ends a +foreground run. Snowflake keeps whatever it had already produced in the +thread either way. ## Configuration options @@ -157,11 +160,13 @@ 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 what happens to the Snowflake run when the ADK -consumer goes away mid-turn. With it enabled the agent sends one cancel request -for the run; the cancel is best effort and its failure is not raised, because -nobody is left to read the error. Disable it if you want abandoned runs to -finish and land in the thread anyway. +`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. Observed +behaviour is that Snowflake ends a foreground run as soon as the stream is +closed and answers the cancel with 409, so the request is a safeguard rather +than 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 @@ -339,8 +344,8 @@ Snowflake's error code, message and request id but never the request body. 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 in the thread, and a - cancelled or abandoned run is still billed for the work it did. + run is not resumed. Snowflake ends a foreground run when its stream closes, + keeps the partial output 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`. diff --git a/src/google/adk/labs/snowflake/_client.py b/src/google/adk/labs/snowflake/_client.py index 4605ce039bb..038e9c9f537 100644 --- a/src/google/adk/labs/snowflake/_client.py +++ b/src/google/adk/labs/snowflake/_client.py @@ -299,9 +299,11 @@ async def cancel_run(self, ctx: ReadonlyContext, run_id: str) -> bool: operation='cancel', ) except CortexClientError as e: + status = getattr(e, 'status_code', None) logger.debug( - 'Best-effort cancel of a Snowflake Cortex run failed: %s', + '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() diff --git a/tests/unittests/labs/snowflake/test_client.py b/tests/unittests/labs/snowflake/test_client.py index 738a602dd93..af022231b92 100644 --- a/tests/unittests/labs/snowflake/test_client.py +++ b/tests/unittests/labs/snowflake/test_client.py @@ -437,6 +437,20 @@ async def test_cancel_run_posts_to_the_cancel_endpoint(): 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( diff --git a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py index 0a9431aec0a..2a5b0173f98 100644 --- a/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py +++ b/tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py @@ -157,9 +157,11 @@ def __init__( 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] = [] @@ -186,7 +188,7 @@ def _handle(self, request: httpx.Request) -> httpx.Response: 200, stream=stream, headers={'content-type': 'text/event-stream'} ) if path.endswith('/cancel'): - return httpx.Response(200) + return httpx.Response(self.cancel_status) return httpx.Response(404) def paths(self, suffix: str) -> list[httpx.Request]: @@ -729,6 +731,35 @@ async def test_disconnect_without_the_option_does_not_cancel(): 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() From ec94fd00a66b2c63a61ed8b8542eb8599beda1d2 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Sat, 5 Sep 2026 00:19:15 +0900 Subject: [PATCH 11/12] feat(labs/snowflake): collect Cortex Analyst suggestions as suggested queries --- .../snowflake/snowflake_cortex_agent/index.md | 5 +- .../adk/labs/snowflake/_event_converter.py | 25 ++++++++- .../labs/snowflake/test_event_converter.py | 55 +++++++++++++++++++ 3 files changed, 83 insertions(+), 2 deletions(-) diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md index 0051841a8eb..0de37ba07f5 100644 --- a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -222,7 +222,10 @@ async def ask(session_id: str, question: str) -> tuple[str, list[dict]]: ``` The same dictionary holds `annotations` for citations, `warnings`, `tables`, -`charts`, `usage` and the Snowflake `run_id`. +`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 diff --git a/src/google/adk/labs/snowflake/_event_converter.py b/src/google/adk/labs/snowflake/_event_converter.py index be18fd7c1ef..94945ee3e7d 100644 --- a/src/google/adk/labs/snowflake/_event_converter.py +++ b/src/google/adk/labs/snowflake/_event_converter.py @@ -177,6 +177,8 @@ def __init__( 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 @@ -189,7 +191,7 @@ def __init__( '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_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, @@ -304,6 +306,14 @@ def final_event(self, *, state_delta: dict[str, Any] | None = None) -> Event: 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}' @@ -378,6 +388,19 @@ def _on_metadata(self, name: str, payload: dict[str, Any]) -> list[Event]: 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]: diff --git a/tests/unittests/labs/snowflake/test_event_converter.py b/tests/unittests/labs/snowflake/test_event_converter.py index 8fb386c3806..2dbeb68350c 100644 --- a/tests/unittests/labs/snowflake/test_event_converter.py +++ b/tests/unittests/labs/snowflake/test_event_converter.py @@ -509,6 +509,61 @@ def test_annotations_and_suggested_queries_land_on_the_final_event(): 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() From d9f66c120104550e5e152916c2e2a1b8c0379966 Mon Sep 17 00:00:00 2001 From: Park Won Gyun Date: Sat, 5 Sep 2026 00:23:59 +0900 Subject: [PATCH 12/12] docs(labs/snowflake): describe the cancel outcome without inferring its cause --- .../snowflake/snowflake_cortex_agent/index.md | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md index 0de37ba07f5..19871457279 100644 --- a/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md +++ b/docs/guides/labs/snowflake/snowflake_cortex_agent/index.md @@ -116,10 +116,10 @@ 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 Snowflake had -already marked the run completed by the time the cancel arrived, answering -`409 Agent run was already completed`: closing the stream is what ends a -foreground run. Snowflake keeps whatever it had already produced in the +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 @@ -163,10 +163,10 @@ 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. Observed -behaviour is that Snowflake ends a foreground run as soon as the stream is -closed and answers the cancel with 409, so the request is a safeguard rather -than the mechanism that stops the run. Disable it to skip that extra request. +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 @@ -347,8 +347,8 @@ Snowflake's error code, message and request id but never the request body. 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 ends a foreground run when its stream closes, - keeps the partial output in the thread, and bills the work it did. + 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`.