Skip to content

fix(agentengine): support Gemini Enterprise AgentSpace streams#777

Open
jankrynauw wants to merge 6 commits into
google:mainfrom
alis-exchange:map-gemini-enterprise-sessions-to-adk
Open

fix(agentengine): support Gemini Enterprise AgentSpace streams#777
jankrynauw wants to merge 6 commits into
google:mainfrom
alis-exchange:map-gemini-enterprise-sessions-to-adk

Conversation

@jankrynauw
Copy link
Copy Markdown

@jankrynauw jankrynauw commented Apr 29, 2026

Problem:

Gemini Enterprise / AgentSpace invokes ADK agents on Agent Engine using the streaming_agent_run_with_events method. ADK Go did not fully match the request and response contract used by the Python ADK Agent Engine implementation.

Observed issues:

  • Gemini Enterprise sends the actual request payload inside input.request_json.
  • The embedded session_id is a Gemini Enterprise / Discovery Engine session resource, not a backend ADK session ID.
  • Passing that external session resource directly to VertexAISessionService fails because user-provided session IDs are not supported.
  • After backend session creation was handled, Gemini Enterprise could invoke the agent, but responses were not rendered correctly because the stream response shape differed from Python ADK.
  • When response rendering began working, Gemini Enterprise displayed duplicate answer text because both partial and final SSE events were emitted.

Solution:

This PR adds Gemini Enterprise / AgentSpace compatibility for streaming_agent_run_with_events in ADK Go:

  • Registers streaming_agent_run_with_events as an Agent Engine stream handler.

  • Decodes the wrapped input.request_json payload.

  • Treats the first Gemini Enterprise session resource as a bootstrap identifier.

  • Checks whether the incoming session_id is already a backend ADK session ID.

  • Creates a generated backend ADK session when the incoming session_id is not already a backend ADK session.

  • Returns the generated backend ADK session ID in the response envelope so Gemini Enterprise can use it on later turns.

  • Reuses the returned backend ADK session ID on subsequent Gemini Enterprise requests.

  • Emits the Python ADK-style response envelope for streaming_agent_run_with_events:

    {
      "events": [],
      "artifacts": [],
      "session_id": "..."
    }
  • Uses non-SSE run mode for streaming_agent_run_with_events to avoid Gemini Enterprise rendering both partial and final text as duplicate answer content.

  • Leaves existing async_stream_query behavior unchanged.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Passed locally:

go test ./server/agentengine/...
go test ./session/vertexai

Coverage added for:

  • decoding input.request_json
  • creating a backend ADK session from the initial Gemini Enterprise session resource
  • reusing a returned backend ADK session ID
  • AgentSpace metadata for streaming_agent_run_with_events
  • Python-style response envelope for AgentSpace streams
  • ensuring streaming_agent_run_with_events uses non-SSE mode to avoid duplicate Gemini Enterprise response text
  • preserving VertexAISessionService behavior that rejects caller-provided session IDs on create

Manual End-to-End (E2E) Tests:

Manually tested with a deployed ADK Go agent invoked from Gemini Enterprise / AgentSpace.

Verified:

  • Gemini Enterprise calls streaming_agent_run_with_events.
  • The request arrives with input.request_json.
  • ADK creates a backend ADK session ID from the first Gemini Enterprise session resource.
  • Gemini Enterprise can continue the conversation using the returned backend ADK session ID.
  • The backend ADK session ID is returned in the AgentSpace response envelope as session_id.
  • The ADK UI shows the expected session, user message, and agent response.
  • Gemini Enterprise renders the agent response after switching to the Python-style stream envelope.
  • Duplicate response text is resolved by using non-SSE mode for streaming_agent_run_with_events.

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

This behavior was compared against the Vertex AI Python SDK AdkApp.streaming_agent_run_with_events implementation. The Go implementation now follows the same high-level contract for Gemini Enterprise / AgentSpace invocation: wrapped request input, generated backend ADK session IDs returned in the response envelope, backend session IDs reused on later turns, and response events wrapped in an AgentSpace-compatible envelope.

@jankrynauw
Copy link
Copy Markdown
Author

Here is an example of a conversation I had on Gemini Enteprise:

Screenshot 2026-04-30 at 00 07 55

And then looking at the corresponding view within Agent Platform:

Screenshot 2026-04-30 at 00 08 39

@kdroste-google kdroste-google self-requested a review April 30, 2026 11:00
Copy link
Copy Markdown
Collaborator

@kdroste-google kdroste-google left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, @jankrynauw,
Thank you for your contribution!!
I think we should not merge those two methods: stream_query and streaming_agent_run_with_events - the separate controllers/method and internal/models should be clearer and more maintainable (no internal "if" and conditionally empty fields).
What do you think?

@jankrynauw
Copy link
Copy Markdown
Author

Hi @kdroste-google, I had that exact same thought this weekend, let me implement the change and get back to you. My initial thought is to keep the logic within the agentengine server and simply split the logic into its own handler so that the current:

// listStreamHandlers returnes a list of handlers for streaming methods
func listStreamHandlers(config *launcher.Config, agentEngineID string) []method.MethodHandler {
	return []method.MethodHandler{
		method.NewStreamQueryHandler(config, agentEngineID, "async_stream_query", "async_stream"),
		method.NewStreamQueryHandler(config, agentEngineID, "streaming_agent_run_with_events", "stream"),
	}
}

effectively becomes

// listStreamHandlers returnes a list of handlers for streaming methods
func listStreamHandlers(config *launcher.Config, agentEngineID string) []method.MethodHandler {
	return []method.MethodHandler{
		method.NewStreamQueryHandler(config, agentEngineID, "async_stream_query", "async_stream"),
		method.NewStreamingAgentRunWithEventsHandler(config, agentEngineID, "streaming_agent_run_with_events", "stream"),
	}
}

@jankrynauw
Copy link
Copy Markdown
Author

just deployed the latest changes via an agent to Agent Runtime on Google Agent Platform, and added it to Gemini Enterprise:
Screenshot 2026-05-04 at 15 27 05
Screenshot 2026-05-04 at 15 28 05
Screenshot 2026-05-04 at 15 30 44

@jankrynauw jankrynauw requested a review from kdroste-google May 4, 2026 13:31
jankrynauw added 4 commits May 5, 2026 14:01
- The Go test check was failing in TestDebugTelemetryGetSpansBySessionID because the test asserted a fixed ordering for returned spans. In practice, the debug telemetry store can return the same set of spans in a different  order depending on span/export timing, especially when spans are created and ended very close together.

- This change keeps the production behavior unchanged and updates the tests to compare spans as a stable sorted set before diffing. The comparison still checks the span names, relevant attributes, and logs, but no longer treats  incidental retrieval order as part of the contract.
Copy link
Copy Markdown
Author

@jankrynauw jankrynauw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kdroste-google I could also remove all references to AgentSpace - and instead use Gemini Enterprise?

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.

Gemini Enterprise integration fails with "unrecognized class method: streaming_agent_run_with_events"

2 participants