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

[miewone]

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):

• Related: Proposal: Add SnowflakeCortexAgent as a BaseAgent adapter #7014

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. Applications also need ADK after-tool callbacks to postprocess server-side tool results and update session state or save artifacts.

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.

Remote tool result callbacks

after_tool_callback accepts a synchronous or asynchronous function, or an ordered list, using the LlmAgent argument convention: tool, args, tool_context, and tool_response. ADK plugin after-tool callbacks run first in registration order, even when no agent callback is configured. The first non-None result, including {}, replaces the result and stops the applicable chain; a plugin replacement skips agent callbacks. Returning None continues the chain and retains in-place mutations.

Callbacks receive the full normalized result before truncation, paired with the original call name and arguments by tool_use_id. Duplicate results with the same ID invoke callbacks only once per invocation. Results without an identifiable prior call, including results arriving before their calls, remain visible but skip callbacks. ToolContext uses the real session and services, so callback state and artifact deltas accompany the resulting FunctionResponse through Runner persistence. Callback objects are excluded from serialization, repr, and web graph output.

Callback result overrides and ToolContext changes affect ADK-side events, session state, artifacts, and application output. With LlmAgent, an after-tool callback can replace a local tool result before the next ADK-controlled model request consumes it. Here, Snowflake owns the remote loop and may already have consumed the original result. These callbacks cannot retroactively change that result or Snowflake's final answer. Remote tool approval, blocking execution with before_tool_callback, client-side execution, and pause/resume are outside this scope. The adapter does not re-execute remote tools locally.

The size limit is applied after callbacks, including to arbitrary replacement fields. Callback-authored state and artifacts have separate lifetimes and are not bounded by max_tool_result_bytes. A callback exception or cancellation prevents that result and the final cursor update from being emitted; earlier events remain stored and artifacts already saved are not rolled back.

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
• Remote callbacks (real SnowflakeCortexAgent and ADK Runner with httpx.MockTransport): single/list and sync/async agent callbacks, ADK positional fallback, plugin registration order and plugin-before-agent execution, first non-None replacement including {}, unchanged results on None, original name/arguments/ID, duplicate and out-of-order results, missing call metadata, state/artifact readback, two-turn SSE/NONE continuity, callback exceptions/cancellation without a successful result or cursor advancement, size bounds after replacement, serialization/graph exclusion, and no local tool re-execution.
• HTTP API: real /run and /run_sse routes return callback-modified results consistent with session/artifact readback, including the existing SSE split into content and artifact-action frames.

Verified on 2026-09-07 at commit bcaba313. All 187 Snowflake tests passed in each environment. The full tests/unittests suite was also rerun sequentially for every supported Python version:

Python	Passed	Skipped	Xfailed	Xpassed	Failures
3.10.19	14,083	85	27	2	0
3.11.14	14,090	84	27	2	0
3.12.12	14,083	85	27	2	0
3.13.0	14,083	85	27	2	0
3.14.3	14,083	85	27	2	0

Total: 70,422 passed, zero failures. All five tox environments returned exit code 0. Existing skip/xfail/xpass outcomes are included above; no test expectations or eval thresholds were relaxed.

Exact full-suite command (the result and log paths are local verification artifacts):

timeout 3600s .venv/bin/tox run --result-json .local_eval/snowflake_cortex_agent/verification_20260907_full/tox_sequential_results.json -x 'testenv.commands=pytest tests/unittests -o faulthandler_timeout=120' > .local_eval/snowflake_cortex_agent/verification_20260907_full/tox_sequential.log 2>&1

faulthandler_timeout=120 enables diagnostic stack dumps for a long-running test; it does not change test selection or pass/fail expectations.

Checks for the callback changes also passed:

.venv/bin/mypy src/google/adk/labs/snowflake/_event_converter.py src/google/adk/labs/snowflake/_snowflake_cortex_agent.py --follow-imports=silent
pre-commit run --files src/google/adk/labs/snowflake/_event_converter.py src/google/adk/labs/snowflake/_snowflake_cortex_agent.py tests/unittests/labs/snowflake/test_snowflake_cortex_agent.py docs/guides/labs/snowflake/snowflake_cortex_agent/index.md
git diff --check

Mypy reported no issues in the two changed source files, and all applicable pre-commit hooks passed.

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

Remote callback live verification — 2026-09-07

Verified at bcaba313, using the real SnowflakeCortexAgent through ADK Runner with SSE streaming on Python 3.11.14. A uniquely identified user/session created a fresh Snowflake thread and continued it for a second turn. The Cortex Agent already had the server-side request_required_filters custom tool registered; no Snowflake Agent, UDF, or permission configuration was changed.

The local verification harness configured an observing plugin and an async agent after_tool_callback. The callback checked the original tool name, arguments, and tool_use_id, parsed the JSON-string envelope in tool_response['content'][*]['json']['result'], wrote a callback counter through ToolContext.state, saved a receipt artifact, and returned the result with adk_callback_verified=True. After each turn, the harness read the session and artifact back through their services and compared the stored events with Runner output.

Check	Turn 1	Turn 2
Tool trajectory	request_required_filters	request_required_filters
FunctionCall / FunctionResponse / agent callback	1 / 1 / 1	1 / 1 / 1
Plugin ran before agent callback	Verified	Verified
Nested JSON-string envelope received	Yes	Yes
Replacement marker in emitted and stored result	Verified	Verified
State readback: callback count	1	2
Artifact readback	Verified	Verified
Thread/cursor request	Initial parent 0	Same thread, previous assistant message as parent
Remote error events	0	0
Completed response events	1	1

Sanitized completion record:

{"verified": true, "total_callbacks": 2, "target_envelopes": 2, "connections_closed": true}

The harness used a 120-second request timeout, a 150-second per-turn timeout, and a 360-second process limit. HTTP connections and Runner resources were closed. Account details, credentials, keys, SQL, business payloads, answer text, and actual thread/session IDs are omitted from this evidence.

Exact live command (the credential-dependent harness and its output are local evaluation artifacts, not files shipped in this PR):

timeout 360s .venv/bin/python .local_eval/snowflake_cortex_agent/run_live_after_tool_callback.py > .local_eval/snowflake_cortex_agent/verification_20260907_full/live.jsonl

To reproduce the checks with another Cortex Agent, configure an equivalent callback and plugin on the sample agent, run two prompts that invoke a registered server-side tool through Runner, and compare callback counts, FunctionCall/FunctionResponse IDs, replacement output, session/artifact readback, and the second request's thread cursor. The nested envelope assertion above is specific to the tested custom tool; the adapter does not depend on that tool or its schema.

This live run verifies plugin-before-agent ordering and agent result replacement. Multiple-plugin ordering, plugin replacement short-circuiting, exceptions/cancellation, duplicate/out-of-order/missing events, oversized results, and NONE mode are covered by MockTransport tests using the real ADK Runner; these conditions were not forced on the remote server. Response-match scores and the existing eval PASS display were not used as callback-compatibility evidence.

Earlier adapter live verification — 2026-09-04

The following traces cover streaming, server-side tool projection, thread continuity, and cancellation behavior before after-tool callbacks were added.

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. (Not applicable: no dependent downstream changes.)

Additional context

• Precedent: AntigravityAgent (feat(labs): add experimental Antigravity SDK agent wrapper #6040) for external-runtime tool trace as FunctionCall/FunctionResponse, SSE-gated partials, per-agent state key.
• Not addressed here: AgentTool event propagation (Feature Request: Support Event Streaming propagation from AgentTool (Sub-agents) to Runner #3984). The adapter remains root-agent only; composition is outside this PR.
• Snowflake docs: Cortex Agents Run API, Threads.
• This PR was prepared with AI assistance (GPT), including root-cause analysis via source inspection, the fix, the added test, and this description. All changes were verified locally as described above before submission.

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]

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.