Feature/auth flow - #456
Conversation
- Add AuthConfig, AuthCredential, AuthScheme types - Add AuthHandler for OAuth2 flow orchestration - Add CredentialManager and CredentialService for token management - Add OAuth2 token exchange helpers - Add unit tests for auth package
Add support for tools to request authentication credentials during execution: - Add RequestedAuthConfigs field to EventActions for storing auth requests - Implement GenerateAuthEvent to convert auth requests to adk_request_credential function calls - Add RequestAuthConfig method to ToolContext for tools to request authentication - Add AuthConfig field to Tool interface for tools to declare auth requirements - Integrate auth event generation into the LLM base flow Also fix test comparisons to ignore RequestedAuthConfigs field initialization and add UTC timezone setup in database tests for consistent timestamp formatting.
Summary of ChangesHello @cpunion, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Agent Development Kit by integrating a complete and robust authentication system. It provides the foundational components and logic necessary for agents to securely interact with external services that require authentication, particularly focusing on OAuth2 and OpenID Connect standards. The system is designed to manage the entire credential lifecycle, from initial request and exchange to refreshing expired tokens, and includes a novel mechanism to gracefully pause and resume agent workflows for interactive user authentication. This capability is crucial for building more powerful and versatile agents that can operate effectively in secure, real-world environments. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive authentication framework, including support for OAuth2/OIDC flows, credential management, and integration with the agent and tool execution lifecycle. The implementation is extensive, adding numerous new files and modifying core components to handle authentication seamlessly. My review focuses on ensuring the robustness and correctness of this new auth flow. I've identified a critical issue that could lead to an infinite loop, a high-severity issue regarding unhandled errors, and a medium-severity issue related to maintainability and robustness of data parsing. Addressing these points will significantly improve the stability of the new authentication feature.
| // This is set by authPreprocessor and read by Flow.runOneStep. | ||
| var CurrentAuthPreprocessorResult *AuthPreprocessorResult | ||
|
|
||
| func authPreprocessor(ctx agent.InvocationContext, req *model.LLMRequest) error { |
There was a problem hiding this comment.
This preprocessor logic has a high risk of causing an infinite loop. The function processes the last user event containing an auth response, but it doesn't mark the event as processed. In the next iteration of the agent's run loop, authPreprocessor will be called again, find the same user event, and re-process it, causing the same tool resumption logic to trigger repeatedly.
To fix this, you should mark the user event as processed after handling it, preventing it from being processed again. A good way to do this is by storing the event's ID in the session's temporary state.
Here's how you can implement this:
- At the beginning of the function, after finding
lastEventWithContent, check if it has already been processed. If so, return early. - After successfully processing the event and setting
CurrentAuthPreprocessorResult, store the event's ID in the session state to mark it as processed.
Here is a snippet for the first part of the fix:
func authPreprocessor(ctx agent.InvocationContext, req *model.LLMRequest) error {
// Reset the result
CurrentAuthPreprocessorResult = nil
// This implements Python ADK's auth_preprocessor logic exactly.
// It checks SESSION EVENTS (not userContent) for auth responses.
// This is crucial - checking session events means we won't re-process
// the same auth response on every runOneStep iteration.
events := ctx.Session().Events()
if events.Len() == 0 {
return nil
}
// Find the last event with non-None content (Python lines 54-60)
var lastEventWithContent *session.Event
for i := events.Len() - 1; i >= 0; i-- {
event := events.At(i)
if event.Content != nil {
lastEventWithContent = event
break
}
}
// Check if the last event with content is authored by user (Python lines 62-64)
if lastEventWithContent == nil || lastEventWithContent.Author != "user" {
return nil
}
processedKey := "processed_auth_event:" + lastEventWithContent.ID
if val, _ := ctx.Session().State().Get(processedKey); val != nil {
// This event has already been processed for auth, skip.
return nil
}| func parseAuthConfigFromMap(data any) (*auth.AuthConfig, error) { | ||
| // The response from OAuth flow should contain the credential | ||
| dataMap, ok := data.(map[string]any) | ||
| if !ok { | ||
| return nil, fmt.Errorf("auth_config is not a map") | ||
| } | ||
|
|
||
| config := &auth.AuthConfig{} | ||
|
|
||
| if credKey, ok := dataMap["credential_key"].(string); ok { | ||
| config.CredentialKey = credKey | ||
| } | ||
|
|
||
| // Try to extract the exchanged credential | ||
| if credData, ok := dataMap["exchanged_auth_credential"]; ok { | ||
| cred, err := parseAuthCredentialFromMap(credData) | ||
| if err == nil { | ||
| config.ExchangedAuthCredential = cred | ||
| } | ||
| } | ||
|
|
||
| return config, nil | ||
| } |
There was a problem hiding this comment.
The manual parsing in parseAuthConfigFromMap (and parseAuthCredentialFromMap) is brittle and incomplete. It only handles a small subset of fields and relies on specific type assertions that might fail if the input map structure changes slightly.
A more robust and maintainable approach is to leverage json.Marshal and json.Unmarshal. This will correctly handle all fields defined with json tags in your structs, including nested structures and different number types.
By refactoring parseAuthConfigFromMap this way, parseAuthCredentialFromMap becomes redundant and can be removed, as json.Unmarshal will handle the nested ExchangedAuthCredential field automatically.
func parseAuthConfigFromMap(data any) (*auth.AuthConfig, error) {
jsonData, err := json.Marshal(data)
if err != nil {
return nil, fmt.Errorf("failed to marshal auth_config data: %w", err)
}
var config auth.AuthConfig
if err := json.Unmarshal(jsonData, &config); err != nil {
return nil, fmt.Errorf("failed to unmarshal auth_config: %w", err)
}
return &config, nil
}
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive authentication package to support various auth flows, particularly OAuth2 and OIDC. The changes are well-structured across a new auth package, with clear separation of concerns for configuration, credentials, handlers, and lifecycle management. The integration into the core agent flow to handle credential requests from tools is also well-implemented. The code is generally of high quality with good test coverage for many of the new components.
My review includes two main suggestions for improvement:
- Ensuring deep copies of
AuthConfigto prevent subtle bugs from shared mutable state. - Expanding test coverage for the
CredentialManagerto better validate its complex logic across different authentication scenarios.
| return nil | ||
| } | ||
| return &AuthConfig{ | ||
| AuthScheme: c.AuthScheme, // AuthScheme is typically immutable |
There was a problem hiding this comment.
The Copy() method performs a shallow copy of AuthScheme. The comment // AuthScheme is typically immutable suggests an assumption, but the AuthScheme interface is implemented by pointer types (e.g., *OAuth2Scheme) which can contain mutable fields like maps and slices (e.g., Scopes). If the original AuthScheme is modified after a copy is made, the changes will be reflected in the copy, which can lead to unexpected behavior and subtle bugs.
To ensure true immutability of the copied AuthConfig, I recommend performing a deep copy of the AuthScheme. This could be done by adding a Copy() AuthScheme method to the AuthScheme interface and implementing it for each concrete scheme type.
For example, OAuth2Scheme contains *OAuthFlows, which in turn contains maps that should be copied.
| // Copyright 2025 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. | ||
|
|
||
| package auth | ||
|
|
||
| import ( | ||
| "context" | ||
| "errors" | ||
| "strings" | ||
| "testing" | ||
| "time" | ||
| ) | ||
|
|
||
| type failingCredentialService struct { | ||
| loadErr error | ||
| } | ||
|
|
||
| func (f *failingCredentialService) LoadCredential(context.Context, *AuthConfig) (*AuthCredential, error) { | ||
| return nil, f.loadErr | ||
| } | ||
|
|
||
| func (f *failingCredentialService) SaveCredential(context.Context, *AuthConfig) error { | ||
| return nil | ||
| } | ||
|
|
||
| type stubRefresher struct { | ||
| shouldRefresh bool | ||
| err error | ||
| refreshed *AuthCredential | ||
| } | ||
|
|
||
| func (s *stubRefresher) IsRefreshNeeded(*AuthCredential, AuthScheme) bool { | ||
| return s.shouldRefresh | ||
| } | ||
|
|
||
| func (s *stubRefresher) Refresh(context.Context, *AuthCredential, AuthScheme) (*AuthCredential, error) { | ||
| if s.err != nil { | ||
| return nil, s.err | ||
| } | ||
| return s.refreshed, nil | ||
| } | ||
|
|
||
| func TestCredentialManager_GetAuthCredential_LoadCredentialError(t *testing.T) { | ||
| cfg := &AuthConfig{ | ||
| AuthScheme: &OAuth2Scheme{ | ||
| Flows: &OAuthFlows{ | ||
| ClientCredentials: &OAuthFlowClientCredentials{ | ||
| TokenURL: "https://example.com/token", | ||
| }, | ||
| }, | ||
| }, | ||
| RawAuthCredential: &AuthCredential{ | ||
| AuthType: AuthCredentialTypeOAuth2, | ||
| OAuth2: &OAuth2Auth{ | ||
| ClientID: "client-id", | ||
| ClientSecret: "client-secret", | ||
| AccessToken: "existing-token", | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| manager := NewCredentialManager(cfg) | ||
|
|
||
| svc := &failingCredentialService{loadErr: errors.New("database offline")} | ||
|
|
||
| if _, err := manager.GetAuthCredential(context.Background(), nil, svc); err == nil || !strings.Contains(err.Error(), "failed to load credential") { | ||
| t.Fatalf("GetAuthCredential() error = %v, want load credential error", err) | ||
| } | ||
| } | ||
|
|
||
| func TestCredentialManager_GetAuthCredential_RefreshError(t *testing.T) { | ||
| cfg := &AuthConfig{ | ||
| AuthScheme: &OAuth2Scheme{ | ||
| Flows: &OAuthFlows{ | ||
| AuthorizationCode: &OAuthFlowAuthorizationCode{ | ||
| AuthorizationURL: "https://example.com/auth", | ||
| TokenURL: "https://example.com/token", | ||
| }, | ||
| }, | ||
| }, | ||
| RawAuthCredential: &AuthCredential{ | ||
| AuthType: AuthCredentialTypeOAuth2, | ||
| OAuth2: &OAuth2Auth{ | ||
| ClientID: "client-id", | ||
| ClientSecret: "client-secret", | ||
| }, | ||
| }, | ||
| ExchangedAuthCredential: &AuthCredential{ | ||
| AuthType: AuthCredentialTypeOAuth2, | ||
| OAuth2: &OAuth2Auth{ | ||
| ClientID: "client-id", | ||
| ClientSecret: "client-secret", | ||
| AccessToken: "expired-token", | ||
| RefreshToken: "refresh", | ||
| ExpiresAt: time.Now().Add(-time.Minute).Unix(), | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| manager := NewCredentialManager(cfg) | ||
| manager.refresherRegistry.Register(AuthCredentialTypeOAuth2, &stubRefresher{ | ||
| shouldRefresh: true, | ||
| err: errors.New("refresh failed"), | ||
| }) | ||
|
|
||
| if _, err := manager.GetAuthCredential(context.Background(), nil); err == nil || !strings.Contains(err.Error(), "failed to refresh credential") { | ||
| t.Fatalf("GetAuthCredential() error = %v, want refresh error", err) | ||
| } | ||
| } |
There was a problem hiding this comment.
The CredentialManager.GetAuthCredential method contains complex logic with multiple branches for handling different authentication scenarios (e.g., ready credentials, client credentials flow, auth code exchange, refresh). The current tests only cover two error paths (LoadCredentialError and RefreshError).
To ensure the correctness and robustness of this critical component, I recommend adding more test cases to cover the various success paths and edge cases, including:
- A simple credential type that is considered "ready" (e.g., APIKey).
- The OAuth2 client credentials flow.
- The OAuth2 authorization code exchange flow.
- The token refresh flow.
- Loading a credential from the
ExchangedAuthCredentialcache. - Loading a credential from a
CredentialService. - Loading a credential from the
stateGetter(simulating an auth response from the user). - The case where no credential is available and user authorization is required (returns
nil, nil).
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive authentication package (/auth) to handle various authentication schemes, with a strong focus on OAuth2/OIDC. This is a significant feature addition that appears to be a port of the Python ADK's auth functionality. The new package is well-structured, with clear separation of concerns for schemes, credentials, handlers, and persistence. The integration into the core agent flow (base_flow.go) and the exposure of auth capabilities to tools (tool.go, context.go) are well-thought-out, enabling tools to request credentials and the agent to handle the multi-step auth process. The code is of high quality and includes extensive testing. My review includes a few suggestions to improve error handling and maintainability in some of the new complex logic.
| func (r *OAuth2Refresher) Refresh(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*AuthCredential, error) { | ||
| if cred.OAuth2 == nil || cred.OAuth2.RefreshToken == "" { | ||
| // No refresh token, return original | ||
| return cred, nil | ||
| } | ||
|
|
||
| tokenURL := tokenEndpointFromScheme(scheme) | ||
|
|
||
| if tokenURL == "" { | ||
| return cred, nil | ||
| } | ||
|
|
||
| config := &oauth2.Config{ | ||
| ClientID: cred.OAuth2.ClientID, | ||
| ClientSecret: cred.OAuth2.ClientSecret, | ||
| Endpoint: oauth2.Endpoint{ | ||
| TokenURL: tokenURL, | ||
| }, | ||
| } | ||
|
|
||
| // Create token source from existing token | ||
| oldToken := &oauth2.Token{ | ||
| AccessToken: cred.OAuth2.AccessToken, | ||
| RefreshToken: cred.OAuth2.RefreshToken, | ||
| Expiry: time.Unix(cred.OAuth2.ExpiresAt, 0), | ||
| } | ||
|
|
||
| tokenSource := config.TokenSource(ctx, oldToken) | ||
| newToken, err := tokenSource.Token() | ||
| if err != nil { | ||
| return cred, err | ||
| } | ||
|
|
||
| // Update credential with new tokens | ||
| newCred := cred.Copy() | ||
| newCred.OAuth2.AccessToken = newToken.AccessToken | ||
| if newToken.RefreshToken != "" { | ||
| newCred.OAuth2.RefreshToken = newToken.RefreshToken | ||
| } | ||
| if !newToken.Expiry.IsZero() { | ||
| newCred.OAuth2.ExpiresAt = newToken.Expiry.Unix() | ||
| newCred.OAuth2.ExpiresIn = int64(time.Until(newToken.Expiry).Seconds()) | ||
| } | ||
|
|
||
| return newCred, nil | ||
| } |
There was a problem hiding this comment.
The Refresh method's error handling could be improved. In several failure scenarios, it returns the original, potentially expired, credential instead of a clear error. This can mask failures and lead to the use of stale credentials.
It is better to return (nil, error) on failure to make it clear that a valid refreshed credential could not be obtained.
Specifically:
- If there is no refresh token (line 284), it should return an error instead of
cred, nil. - If the token URL cannot be determined (line 291), it should also return an error.
- If the token exchange fails (line 312), it should return
nilfor the credential along with the error, not the originalcred.
Here are the suggested changes:
// In place of lines 284-287
if cred.OAuth2 == nil {
return nil, fmt.Errorf("oauth2 credential cannot be nil for refresh")
}
if cred.OAuth2.RefreshToken == "" {
return nil, fmt.Errorf("no refresh token available to refresh credential")
}// In place of lines 291-293
if tokenURL == "" {
return nil, fmt.Errorf("no token endpoint URL found in auth scheme to refresh token")
}// In place of lines 312-314
if err != nil {
return nil, err
}| func authPreprocessor(ctx agent.InvocationContext, req *model.LLMRequest) error { | ||
| // TODO: implement (adk-python src/google/adk/auth/auth_preprocessor.py) | ||
| // Reset the result | ||
| CurrentAuthPreprocessorResult = nil | ||
|
|
||
| // This implements Python ADK's auth_preprocessor logic exactly. | ||
| // It checks SESSION EVENTS (not userContent) for auth responses. | ||
| // This is crucial - checking session events means we won't re-process | ||
| // the same auth response on every runOneStep iteration. | ||
|
|
||
| events := ctx.Session().Events() | ||
| if events.Len() == 0 { | ||
| return nil | ||
| } | ||
|
|
||
| // Find the last event with non-None content (Python lines 54-60) | ||
| var lastEventWithContent *session.Event | ||
| for i := events.Len() - 1; i >= 0; i-- { | ||
| event := events.At(i) | ||
| if event.Content != nil { | ||
| lastEventWithContent = event | ||
| break | ||
| } | ||
| } | ||
|
|
||
| // Check if the last event with content is authored by user (Python lines 62-64) | ||
| if lastEventWithContent == nil || lastEventWithContent.Author != "user" { | ||
| return nil | ||
| } | ||
|
|
||
| // Get function responses from the event (Python lines 66-68) | ||
| var functionResponses []*genai.FunctionResponse | ||
| for _, part := range lastEventWithContent.Content.Parts { | ||
| if part.FunctionResponse != nil { | ||
| functionResponses = append(functionResponses, part.FunctionResponse) | ||
| } | ||
| } | ||
| if len(functionResponses) == 0 { | ||
| return nil | ||
| } | ||
|
|
||
| // Collect request_euc function call IDs and store credentials (Python lines 70-80) | ||
| requestEucFunctionCallIDs := make(map[string]bool) | ||
| for _, funcResponse := range functionResponses { | ||
| if funcResponse.Name != auth.RequestEUCFunctionCallName { | ||
| continue | ||
| } | ||
| // Found the function call response for the system long running request euc function call | ||
| requestEucFunctionCallIDs[funcResponse.ID] = true | ||
|
|
||
| // Parse and store the credential | ||
| if funcResponse.Response != nil { | ||
| if authConfigData, ok := funcResponse.Response["auth_config"]; ok { | ||
| authConfig, err := parseAuthConfigFromMap(authConfigData) | ||
| if err != nil { | ||
| continue | ||
| } | ||
| // Store the credential in session state | ||
| if authConfig.CredentialKey != "" && authConfig.ExchangedAuthCredential != nil { | ||
| key := session.KeyPrefixTemp + authConfig.CredentialKey | ||
| if err := ctx.Session().State().Set(key, authConfig.ExchangedAuthCredential); err != nil { | ||
| return fmt.Errorf("failed to store auth credential: %w", err) | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| if len(requestEucFunctionCallIDs) == 0 { | ||
| return nil | ||
| } | ||
|
|
||
| // Now find the original tool calls that need to be resumed. | ||
| // Python lines 85-130: Search backwards for adk_request_credential function calls, | ||
| // then find the original tool calls that triggered them. | ||
|
|
||
| result := &AuthPreprocessorResult{ | ||
| ToolIdsToResume: make(map[string]bool), | ||
| } | ||
|
|
||
| for i := events.Len() - 2; i >= 0; i-- { | ||
| event := events.At(i) | ||
| if event.Content == nil { | ||
| continue | ||
| } | ||
|
|
||
| // Look for adk_request_credential function calls in this event (Python lines 87-101) | ||
| var functionCalls []*genai.FunctionCall | ||
| for _, part := range event.Content.Parts { | ||
| if part.FunctionCall != nil { | ||
| functionCalls = append(functionCalls, part.FunctionCall) | ||
| } | ||
| } | ||
| if len(functionCalls) == 0 { | ||
| continue | ||
| } | ||
|
|
||
| toolsToResume := make(map[string]bool) | ||
| for _, fc := range functionCalls { | ||
| if !requestEucFunctionCallIDs[fc.ID] { | ||
| continue | ||
| } | ||
| // Extract function_call_id from args (the original tool that requested auth) | ||
| if args := fc.Args; args != nil { | ||
| if fcID, ok := args["function_call_id"].(string); ok { | ||
| toolsToResume[fcID] = true | ||
| } | ||
| } | ||
| } | ||
|
|
||
| if len(toolsToResume) == 0 { | ||
| continue | ||
| } | ||
|
|
||
| // Found the system long running request euc function call | ||
| // Now looking for original function call that requests euc (Python lines 103-129) | ||
| for j := i - 1; j >= 0; j-- { | ||
| originalEvent := events.At(j) | ||
| if originalEvent.Content == nil { | ||
| continue | ||
| } | ||
|
|
||
| var originalFunctionCalls []*genai.FunctionCall | ||
| for _, part := range originalEvent.Content.Parts { | ||
| if part.FunctionCall != nil { | ||
| originalFunctionCalls = append(originalFunctionCalls, part.FunctionCall) | ||
| } | ||
| } | ||
| if len(originalFunctionCalls) == 0 { | ||
| continue | ||
| } | ||
|
|
||
| // Check if any function call matches our tools_to_resume | ||
| hasMatch := false | ||
| for _, fc := range originalFunctionCalls { | ||
| if toolsToResume[fc.ID] { | ||
| hasMatch = true | ||
| break | ||
| } | ||
| } | ||
|
|
||
| if hasMatch { | ||
| // Found the original event containing function calls to resume | ||
| result.ToolIdsToResume = toolsToResume | ||
| result.OriginalEvent = originalEvent | ||
| result.CredentialsStored = true | ||
| CurrentAuthPreprocessorResult = result | ||
| return nil | ||
| } | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| return nil | ||
| } |
There was a problem hiding this comment.
The authPreprocessor function is very long (over 150 lines) and has a high cyclomatic complexity due to multiple nested loops and backward searches through the session's event history. This makes the function difficult to understand, maintain, and test. Consider refactoring it into several smaller, well-named helper functions, each with a single responsibility. For example:
- A function to find the last user event containing function responses.
- A function to parse and store credentials from function responses.
- A function to find the original tool call event that needs to be resumed.
This will improve readability and make the logic easier to follow.
| if err != nil { | ||
| continue | ||
| } |
There was a problem hiding this comment.
The error from parseAuthConfigFromMap is silently ignored by the continue statement. If parsing fails, the credential from the user response will be skipped without any notification. This can make debugging difficult. This error should be logged to provide visibility into why a credential was not stored.
| if err != nil { | ||
| // TODO: log or surface the error once a logging strategy is defined. | ||
| return | ||
| } |
There was a problem hiding this comment.
The error returned by handler.GenerateAuthRequest() is silently ignored. If an error occurs during auth request generation (e.g., due to an invalid auth scheme configuration), the request for credentials will be silently dropped. This can lead to the authentication flow stalling without a clear reason. The error should be handled, for instance by logging it, to aid in debugging. Since this function doesn't return an error, consider adding logging here.
1 participant
Base on:
#455 feat(auth): add OAuth2/OIDC authentication package