feat(labs): add SnowflakeCortexAgent (Cortex Agents Run API SSE → ADK events) - #7015
Draft
miewone wants to merge 12 commits into
Draft
feat(labs): add SnowflakeCortexAgent (Cortex Agents Run API SSE → ADK events)#7015miewone wants to merge 12 commits into
miewone wants to merge 12 commits into
Conversation
|
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. |
This comment was marked as resolved.
This comment was marked as resolved.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Please ensure you have read the contribution guide before creating a pull request.
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 +
McpToolsetthe whole run collapses into a singleFunctionResponse, so ADK applications cannot stream progressively, observe tool use, keep citations structured, or continue a Snowflake thread across turns.Solution:
SnowflakeCortexAgent(BaseAgent)undergoogle.adk.labs.snowflakecalls the Cortex Agents Run REST API directly (httpx, no new dependency), parses SSE incrementally, and yields ADK events following theAntigravityAgentprecedent: SSE-gated partial deltas,FunctionCall/FunctionResponseevents for server-side tool trace, one non-partial final event carrying the answer, namespacedcustom_metadata, and the thread cursor asstate_delta. Root-agent only in this PR; composition guards rejectsub_agentsand 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 indocs/guides/README.md)Testing Plan
Unit Tests:
Coverage (mock SSE, no Snowflake access):
data:, non-JSON data,[DONE]/done, terminalerror, unknown events in order, buffer size limitssequence_number, thinking gating bystreaming_mode,tool_use/tool_resultcorrelation bytool_use_id, dedupe, result truncation atmax_tool_result_bytes,client_side_execute=true/permissionrejected, finalresponseaggregation, annotations/warnings/tables/charts/suggested queries metadata (including Cortex Analyst suggestion deltas assembled per index), unknown event passthroughhttpx.MockTransport): thread creation request/response and unusablethread_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 raiseCortexApiErrorwith status, Snowflake code and request id (never the token); connect/read timeouts and dropped connections raiseCortexTransportError; leaving the stream early closes the upstream response; cancel is best effort; a sharedhttpx.AsyncClientis left openparent_message_id=0; second turn reuses the thread with the last assistant id; theRunnerpersists the cursor between turns; user metadata id is never stored; cursor unchanged on terminalerror, on a stream cut before the finalresponse(CortexTransportError), when the finalstatusis notcompleted, and when no assistant id was seen;[DONE]is optional once the finalresponsearrived; per-agent state keys; fingerprint mismatch and malformed cursors fail closed before any request, without quoting ids;sub_agentsandparent_agentrejected, also afterclone();header_providerabsent fromrepr,model_dumpand theadk webagent graph; SSE mode streams partial deltas,NONEyields only persisted events; tool trace recorded asFunctionCall/FunctionResponse; reading stops at[DONE]; disconnect closes the upstream response and cancels the run ({thread_id}-{user_message_id}) unlesscancel_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 openpytest 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, fulltests/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/integrationsVerified on 2026-09-04 against a real Snowflake account (Key Pair JWT auth,
X-Snowflake-Role/X-Snowflake-Warehouseheaders). Account host, object names, semantic view name, Snowflake ids, SQL and rows are masked.Run A, two turns in one ADK session through
Runnerwith SSE streaming:Run B, one turn where the Cortex Agent executed SQL and built a chart:
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/FunctionResponseand paired bytool_use_id, oversized tool result reduced to its key sizes,response.warning,response.table,response.chartandresponse.suggested_querieson 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_idwas confirmed to be<thread_id>-<user_message_id>.adk evalover the same agent (2 cases, 3 turns) also completed every turn.Checklist
Additional context
AntigravityAgent(feat(labs): add experimental Antigravity SDK agent wrapper #6040) for external-runtime tool trace asFunctionCall/FunctionResponse, SSE-gated partials, per-agent state key.AgentToolevent propagation (Feature Request: Support Event Streaming propagation from AgentTool (Sub-agents) to Runner #3984). This adapter is root-agent only untilmode="single_turn"lands in a follow-up.NO_UNIT_GUIDE=Private transport, parser, and converter modules (
_client.py,_sse_parser.py,_event_converter.py) are documented through the SnowflakeCortexAgent unit guide atdocs/guides/labs/snowflake/snowflake_cortex_agent/index.md.