Skip to content

feat(labs): add SnowflakeCortexAgent (Cortex Agents Run API SSE → ADK events) - #7015

Draft
miewone wants to merge 12 commits into
google:mainfrom
miewone:feat/snowflake-cortex-agent
Draft

feat(labs): add SnowflakeCortexAgent (Cortex Agents Run API SSE → ADK events)#7015
miewone wants to merge 12 commits into
google:mainfrom
miewone:feat/snowflake-cortex-agent

Conversation

@miewone

@miewone miewone commented Sep 4, 2026

Copy link
Copy Markdown

Please ensure you have read the contribution guide before creating a pull request.

Draft opened alongside #7014 for early design feedback. Scope, whether labs/snowflake is the right first landing (or integrations/), and event representation are discussed there; this PR will follow the outcome.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

2. Or, if no issue exists, describe the change:

Problem:

Snowflake Cortex Agents stream typed SSE events (status, thinking/text deltas, server-side tool use/result, citations, warnings, tables/charts, final response, thread metadata). Through Managed MCP + McpToolset the whole run collapses into a single FunctionResponse, so ADK applications cannot stream progressively, observe tool use, keep citations structured, or continue a Snowflake thread across turns.

Solution:

SnowflakeCortexAgent(BaseAgent) under google.adk.labs.snowflake calls the Cortex Agents Run REST API directly (httpx, no new dependency), parses SSE incrementally, and yields ADK events following the AntigravityAgent precedent: SSE-gated partial deltas, FunctionCall / FunctionResponse events for server-side tool trace, one non-partial final event carrying the answer, namespaced custom_metadata, and the thread cursor as state_delta. Root-agent only in this PR; composition guards reject sub_agents and parent registration.

Files:

  • src/google/adk/labs/snowflake/ (_snowflake_cortex_agent.py, _client.py, _sse_parser.py, _event_converter.py, README.md)
  • tests/unittests/labs/snowflake/
  • contributing/samples/integrations/snowflake_cortex_agent/
  • docs/guides/labs/snowflake/snowflake_cortex_agent/index.md (listed in docs/guides/README.md)

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.
pytest tests/unittests/labs/snowflake -q

Coverage (mock SSE, no Snowflake access):

  • SSE parser: arbitrary byte chunk boundaries, split UTF-8 sequences, LF/CRLF event boundaries, multi-line data:, non-JSON data, [DONE] / done, terminal error, unknown events in order, buffer size limits
  • Event converter: delta order and duplicate sequence_number, thinking gating by streaming_mode, tool_use/tool_result correlation by tool_use_id, dedupe, result truncation at max_tool_result_bytes, client_side_execute=true / permission rejected, final response aggregation, annotations/warnings/tables/charts/suggested queries metadata (including Cortex Analyst suggestion deltas assembled per index), unknown event passthrough
  • REST client (httpx.MockTransport): thread creation request/response and unusable thread_id; run request body (thread_id, parent_message_id, messages, stream), Accept: text/event-stream, URL-encoded object names; strict Snowflake id validation before any request; 401/403/429/5xx and non-event-stream answers raise CortexApiError with status, Snowflake code and request id (never the token); connect/read timeouts and dropped connections raise CortexTransportError; leaving the stream early closes the upstream response; cancel is best effort; a shared httpx.AsyncClient is left open
  • Agent: first turn creates a thread once and sends parent_message_id=0; second turn reuses the thread with the last assistant id; the Runner persists the cursor between turns; user metadata id is never stored; cursor unchanged on terminal error, on a stream cut before the final response (CortexTransportError), when the final status is not completed, and when no assistant id was seen; [DONE] is optional once the final response arrived; per-agent state keys; fingerprint mismatch and malformed cursors fail closed before any request, without quoting ids; sub_agents and parent_agent rejected, also after clone(); header_provider absent from repr, model_dump and the adk web agent graph; SSE mode streams partial deltas, NONE yields only persisted events; tool trace recorded as FunctionCall / FunctionResponse; reading stops at [DONE]; disconnect closes the upstream response and cancels the run ({thread_id}-{user_message_id}) unless cancel_on_disconnect=False; no cancel before the user message id is known (no run id yet) or after the run finished; a 409 from the cancel endpoint is swallowed; cleanup() leaves a shared client open

pytest tests/unittests/labs/snowflake -q: 150 passed (2026-09-04, mock SSE only). mypy src/google/adk/labs/snowflake/: no issues. tox (py310, py311, py312, py313, py314, full tests/unittests): all environments OK.

Manual End-to-End (E2E) Tests:

Setup: a Snowflake account with a Cortex Agent object, a semantic view and an OAuth/PAT token. Sample: contributing/samples/integrations/snowflake_cortex_agent.

export SNOWFLAKE_ACCOUNT_URL=... SNOWFLAKE_DATABASE=... SNOWFLAKE_SCHEMA=... \
       SNOWFLAKE_CORTEX_AGENT=... SNOWFLAKE_TOKEN=...
adk web contributing/samples/integrations

Verified on 2026-09-04 against a real Snowflake account (Key Pair JWT auth, X-Snowflake-Role/X-Snowflake-Warehouse headers). Account host, object names, semantic view name, Snowflake ids, SQL and rows are masked.

Run A, two turns in one ADK session through Runner with SSE streaming:

=== turn 1: 최근 1주 동안 국가별 광고비와 ROAS를 간단히 요약해 주세요.
[tool_call]     <semantic_view_tool> id=<id> args={"pruning_question": "<...>"}
[tool_response] <semantic_view_tool> author=<semantic_view_tool> status=success truncated original_bytes=56397
[final]         status=completed text_chars=956 run_id=<present>
                annotations=0 warnings=4 suggested_queries=0 tables=0 charts=0
[state_delta]   thread_id=<id:8 digits> parent_message_id=<id:10 digits>
[streamed]      text_chars=956 thought_chars=1332 progress_events=13
=== turn 2: 그중 모바일만 다시 보여 주세요.
[final]         status=completed text_chars=765 run_id=<present>   # answer refers back to turn 1
[state_delta]   thread_id=<id:8 digits> parent_message_id=<id:10 digits>   # advanced to the new assistant id
session cursor after 2 turn(s): thread_id=<id:8 digits> parent_message_id=<id:10 digits> schema_version=1

Run B, one turn where the Cortex Agent executed SQL and built a chart:

[tool_call]     <custom_tool> id=<id> args={"pruning_question": "<32 chars>"}
[tool_response] <custom_tool> status=success content_types=['json'] truncated original_bytes=56397
[tool_call]     system_execute_sql id=<id> args={"semantic_model": "<18 chars>", "sql": "<sql masked, 391 chars>"}
[tool_response] system_execute_sql status=success content_types=['json'] rows=2
[tool_call]     <custom_tool> id=<id> args={"skill_name": "<18 chars>"}
[tool_response] <custom_tool> status=success content_types=['json']
[tool_call]     <custom_tool> id=<id> args={"chart": "<1204 chars>", "question": "<91 chars>", "tool_result_id": "<35 chars>"}
[tool_response] <custom_tool> status=success content_types=['json']
[final]         status=completed text_chars=606 run_id=<present>
                annotations=0 warnings=4 suggested_queries=3 tables=1 charts=1
[state_delta]   thread_id=<id:8 digits> parent_message_id=<id:10 digits>
[streamed]      text_chars=606 thought_chars=31141 progress_events=37

Observed: thread creation and two-turn continuity, cursor committed with the assistant message id only, streamed text equal to the final text, four server-side tool calls recorded as FunctionCall/FunctionResponse and paired by tool_use_id, oversized tool result reduced to its key sizes, response.warning, response.table, response.chart and response.suggested_queries on the final event.

Not exercised live: response.text.annotation (the test agent has no Cortex Search).

The cancel path was exercised after the client closed the SSE stream mid-run. Snowflake answered the follow-up cancel with 409 399528 Agent run was already completed, the same answer as for a finished run. A 409 alone does not show whether closing the stream ended the run or it finished on its own, so the adapter treats it as a benign best-effort outcome.

run_id was confirmed to be <thread_id>-<user_message_id>. adk eval over the same agent (2 cases, 3 turns) also completed every turn.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

NO_UNIT_GUIDE=Private transport, parser, and converter modules (_client.py, _sse_parser.py, _event_converter.py) are documented through the SnowflakeCortexAgent unit guide at docs/guides/labs/snowflake/snowflake_cortex_agent/index.md.

@google-cla

google-cla Bot commented Sep 4, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@miewone

This comment was marked as resolved.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant