diff --git a/agent/remoteagent/a2a_agent_test.go b/agent/remoteagent/a2a_agent_test.go index 0ed190520..ecb711306 100644 --- a/agent/remoteagent/a2a_agent_test.go +++ b/agent/remoteagent/a2a_agent_test.go @@ -756,6 +756,7 @@ func TestRemoteAgent_EmptyResultForEmptySession(t *testing.T) { cmpopts.IgnoreFields(session.Event{}, "ID"), cmpopts.IgnoreFields(session.Event{}, "Timestamp"), cmpopts.IgnoreFields(session.EventActions{}, "StateDelta"), + cmpopts.IgnoreFields(session.EventActions{}, "RequestedAuthConfigs"), } if diff := cmp.Diff(wantEvents, gotEvents, ignoreFields...); diff != "" { t.Fatalf("agent.Run() wrong result (+got,-want):\ngot = %+v\nwant = %+v\ndiff = %s", gotEvents, wantEvents, diff) diff --git a/agent/remoteagent/utils_test.go b/agent/remoteagent/utils_test.go index d3ba10c70..755ea61fc 100644 --- a/agent/remoteagent/utils_test.go +++ b/agent/remoteagent/utils_test.go @@ -305,6 +305,7 @@ func TestPresentAsUserMessage(t *testing.T) { cmpopts.IgnoreFields(session.Event{}, "InvocationID"), cmpopts.IgnoreFields(session.Event{}, "Timestamp"), cmpopts.IgnoreFields(session.EventActions{}, "StateDelta"), + cmpopts.IgnoreFields(session.EventActions{}, "RequestedAuthConfigs"), } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { diff --git a/agent/workflowagents/loopagent/agent_test.go b/agent/workflowagents/loopagent/agent_test.go index ee38334cf..b78c60f27 100644 --- a/agent/workflowagents/loopagent/agent_test.go +++ b/agent/workflowagents/loopagent/agent_test.go @@ -234,7 +234,7 @@ func TestNewLoopAgent(t *testing.T) { ignoreFields := []cmp.Option{ cmpopts.IgnoreFields(session.Event{}, "ID", "InvocationID", "Timestamp"), - cmpopts.IgnoreFields(session.EventActions{}, "StateDelta"), + cmpopts.IgnoreFields(session.EventActions{}, "StateDelta", "RequestedAuthConfigs"), cmpopts.IgnoreFields(genai.FunctionCall{}, "ID"), cmpopts.IgnoreFields(genai.FunctionResponse{}, "ID"), } diff --git a/agent/workflowagents/sequentialagent/agent_test.go b/agent/workflowagents/sequentialagent/agent_test.go index 58641ada2..07618e216 100644 --- a/agent/workflowagents/sequentialagent/agent_test.go +++ b/agent/workflowagents/sequentialagent/agent_test.go @@ -254,7 +254,7 @@ func TestNewSequentialAgent(t *testing.T) { for i, gotEvent := range gotEvents { tt.wantEvents[i].Timestamp = gotEvent.Timestamp if diff := cmp.Diff(tt.wantEvents[i], gotEvent, cmpopts.IgnoreFields(session.Event{}, "ID", "Timestamp", "InvocationID"), - cmpopts.IgnoreFields(session.EventActions{}, "StateDelta")); diff != "" { + cmpopts.IgnoreFields(session.EventActions{}, "StateDelta", "RequestedAuthConfigs")); diff != "" { t.Errorf("event[i] mismatch (-want +got):\n%s", diff) } } diff --git a/auth/auth_config.go b/auth/auth_config.go new file mode 100644 index 000000000..da899cb61 --- /dev/null +++ b/auth/auth_config.go @@ -0,0 +1,274 @@ +// 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 ( + "bytes" + "crypto/sha256" + "encoding/json" + "fmt" + "sort" + "strconv" +) + +// AuthConfig combines auth scheme and credentials for a tool. +// This is passed to tools that require authentication. +type AuthConfig struct { + // AuthScheme defines how the API expects authentication. + AuthScheme AuthScheme `json:"authScheme"` + // RawAuthCredential is the initial credential (e.g., client_id/secret). + RawAuthCredential *AuthCredential `json:"rawAuthCredential,omitempty"` + // ExchangedAuthCredential is the processed credential (e.g., access_token). + ExchangedAuthCredential *AuthCredential `json:"exchangedAuthCredential,omitempty"` + // CredentialKey is a unique key for persisting this credential. + CredentialKey string `json:"credentialKey,omitempty"` +} + +// NewAuthConfig creates a new AuthConfig with the given scheme and credential. +// If credentialKey is empty, it will be generated automatically. +func NewAuthConfig(scheme AuthScheme, credential *AuthCredential) (*AuthConfig, error) { + cfg := &AuthConfig{ + AuthScheme: scheme, + RawAuthCredential: credential, + } + if cfg.CredentialKey == "" { + key, err := cfg.generateCredentialKey() + if err != nil { + return nil, fmt.Errorf("generate credential key: %w", err) + } + cfg.CredentialKey = key + } + return cfg, nil +} + +// generateCredentialKey creates a unique key based on auth scheme and credential. +func (c *AuthConfig) generateCredentialKey() (string, error) { + var schemePart, credPart string + if c.AuthScheme != nil { + schemeJSON, err := stableJSON(c.AuthScheme) + if err != nil { + return "", fmt.Errorf("marshal auth scheme: %w", err) + } + schemeType := c.AuthScheme.GetType() + h := sha256.Sum256([]byte(schemeJSON)) + schemePart = fmt.Sprintf("%s_%x", schemeType, h[:8]) + } + if c.RawAuthCredential != nil { + credJSON, err := stableJSON(c.RawAuthCredential) + if err != nil { + return "", fmt.Errorf("marshal auth credential: %w", err) + } + h := sha256.Sum256([]byte(credJSON)) + credPart = fmt.Sprintf("%s_%x", c.RawAuthCredential.AuthType, h[:8]) + } + return fmt.Sprintf("adk_%s_%s", schemePart, credPart), nil +} + +// Copy creates a deep copy of the AuthConfig. +func (c *AuthConfig) Copy() *AuthConfig { + if c == nil { + return nil + } + return &AuthConfig{ + AuthScheme: cloneAuthScheme(c.AuthScheme), + RawAuthCredential: c.RawAuthCredential.Copy(), + ExchangedAuthCredential: c.ExchangedAuthCredential.Copy(), + CredentialKey: c.CredentialKey, + } +} + +func cloneAuthScheme(s AuthScheme) AuthScheme { + switch v := s.(type) { + case *APIKeyScheme: + if v == nil { + return nil + } + cp := *v + return &cp + case *HTTPScheme: + if v == nil { + return nil + } + cp := *v + return &cp + case *OAuth2Scheme: + if v == nil { + return nil + } + return &OAuth2Scheme{ + Flows: cloneOAuthFlows(v.Flows), + Description: v.Description, + } + case *OpenIDConnectScheme: + if v == nil { + return nil + } + cp := &OpenIDConnectScheme{ + OpenIDConnectURL: v.OpenIDConnectURL, + AuthorizationEndpoint: v.AuthorizationEndpoint, + TokenEndpoint: v.TokenEndpoint, + UserInfoEndpoint: v.UserInfoEndpoint, + RevocationEndpoint: v.RevocationEndpoint, + Description: v.Description, + } + if len(v.GrantTypesSupported) > 0 { + cp.GrantTypesSupported = append([]string{}, v.GrantTypesSupported...) + } + if len(v.Scopes) > 0 { + cp.Scopes = append([]string{}, v.Scopes...) + } + return cp + default: + return v + } +} + +func cloneOAuthFlows(flows *OAuthFlows) *OAuthFlows { + if flows == nil { + return nil + } + return &OAuthFlows{ + Implicit: cloneOAuthFlowImplicit(flows.Implicit), + Password: cloneOAuthFlowPassword(flows.Password), + ClientCredentials: cloneOAuthFlowClientCredentials(flows.ClientCredentials), + AuthorizationCode: cloneOAuthFlowAuthorizationCode(flows.AuthorizationCode), + } +} + +func cloneOAuthFlowImplicit(flow *OAuthFlowImplicit) *OAuthFlowImplicit { + if flow == nil { + return nil + } + return &OAuthFlowImplicit{ + AuthorizationURL: flow.AuthorizationURL, + RefreshURL: flow.RefreshURL, + Scopes: cloneScopes(flow.Scopes), + } +} + +func cloneOAuthFlowPassword(flow *OAuthFlowPassword) *OAuthFlowPassword { + if flow == nil { + return nil + } + return &OAuthFlowPassword{ + TokenURL: flow.TokenURL, + RefreshURL: flow.RefreshURL, + Scopes: cloneScopes(flow.Scopes), + } +} + +func cloneOAuthFlowClientCredentials(flow *OAuthFlowClientCredentials) *OAuthFlowClientCredentials { + if flow == nil { + return nil + } + return &OAuthFlowClientCredentials{ + TokenURL: flow.TokenURL, + RefreshURL: flow.RefreshURL, + Scopes: cloneScopes(flow.Scopes), + } +} + +func cloneOAuthFlowAuthorizationCode(flow *OAuthFlowAuthorizationCode) *OAuthFlowAuthorizationCode { + if flow == nil { + return nil + } + return &OAuthFlowAuthorizationCode{ + AuthorizationURL: flow.AuthorizationURL, + TokenURL: flow.TokenURL, + RefreshURL: flow.RefreshURL, + Scopes: cloneScopes(flow.Scopes), + } +} + +func cloneScopes(scopes map[string]string) map[string]string { + if scopes == nil { + return nil + } + cp := make(map[string]string, len(scopes)) + for k, v := range scopes { + cp[k] = v + } + return cp +} + +// stableJSON returns a deterministic JSON representation with sorted map keys. +func stableJSON(v interface{}) (string, error) { + raw, err := json.Marshal(v) + if err != nil { + return "", err + } + var data interface{} + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.UseNumber() + if err := dec.Decode(&data); err != nil { + return "", err + } + var buf bytes.Buffer + if err := encodeCanonical(&buf, data); err != nil { + return "", err + } + return buf.String(), nil +} + +func encodeCanonical(buf *bytes.Buffer, v interface{}) error { + switch val := v.(type) { + case nil: + buf.WriteString("null") + case bool: + if val { + buf.WriteString("true") + } else { + buf.WriteString("false") + } + case string: + buf.WriteString(strconv.Quote(val)) + case json.Number: + buf.WriteString(val.String()) + case float64: + buf.WriteString(strconv.FormatFloat(val, 'g', -1, 64)) + case []interface{}: + buf.WriteByte('[') + for i, elem := range val { + if i > 0 { + buf.WriteByte(',') + } + if err := encodeCanonical(buf, elem); err != nil { + return err + } + } + buf.WriteByte(']') + case map[string]interface{}: + buf.WriteByte('{') + keys := make([]string, 0, len(val)) + for k := range val { + keys = append(keys, k) + } + sort.Strings(keys) + for i, k := range keys { + if i > 0 { + buf.WriteByte(',') + } + buf.WriteString(strconv.Quote(k)) + buf.WriteByte(':') + if err := encodeCanonical(buf, val[k]); err != nil { + return err + } + } + buf.WriteByte('}') + default: + return fmt.Errorf("unsupported JSON canonicalization type %T", v) + } + return nil +} diff --git a/auth/auth_config_test.go b/auth/auth_config_test.go new file mode 100644 index 000000000..55b774916 --- /dev/null +++ b/auth/auth_config_test.go @@ -0,0 +1,246 @@ +// 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 ( + "errors" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestNewAuthConfig(t *testing.T) { + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + }, + } + + cfg, err := NewAuthConfig(scheme, cred) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + + if cfg.AuthScheme != scheme { + t.Error("AuthScheme not set correctly") + } + if cfg.RawAuthCredential != cred { + t.Error("RawAuthCredential not set correctly") + } + if cfg.CredentialKey == "" { + t.Error("CredentialKey should be auto-generated") + } + if !strings.HasPrefix(cfg.CredentialKey, "adk_") { + t.Errorf("CredentialKey = %q, want prefix 'adk_'", cfg.CredentialKey) + } +} + +func TestAuthConfig_generateCredentialKey_Deterministic(t *testing.T) { + scheme := &APIKeyScheme{ + In: APIKeyInHeader, + Name: "X-API-Key", + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: "test-key", + } + + cfg1, err := NewAuthConfig(scheme, cred) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + cfg2, err := NewAuthConfig(scheme, cred) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + + if cfg1.CredentialKey != cfg2.CredentialKey { + t.Errorf("generateCredentialKey not deterministic: %q != %q", cfg1.CredentialKey, cfg2.CredentialKey) + } +} + +func TestAuthConfig_generateCredentialKey_Different(t *testing.T) { + scheme := &APIKeyScheme{ + In: APIKeyInHeader, + Name: "X-API-Key", + } + cred1 := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: "key-1", + } + cred2 := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: "key-2", + } + + cfg1, err := NewAuthConfig(scheme, cred1) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + cfg2, err := NewAuthConfig(scheme, cred2) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + + if cfg1.CredentialKey == cfg2.CredentialKey { + t.Error("Different credentials should produce different keys") + } +} + +func TestAuthConfig_generateCredentialKey_ScopeOrder(t *testing.T) { + makeCfg := func(scopes map[string]string) *AuthConfig { + cfg, err := NewAuthConfig(&OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + Scopes: scopes, + }, + }, + }, &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client", + }, + }) + if err != nil { + t.Fatalf("NewAuthConfig() error = %v", err) + } + return cfg + } + + cfg1 := makeCfg(map[string]string{ + "read": "Read", + "write": "Write", + }) + cfg2 := makeCfg(map[string]string{ + "write": "Write", + "read": "Read", + }) + + if cfg1.CredentialKey != cfg2.CredentialKey { + t.Fatalf("credential keys differ for same scopes: %q vs %q", cfg1.CredentialKey, cfg2.CredentialKey) + } +} + +func TestAuthConfig_Copy_Nil(t *testing.T) { + var cfg *AuthConfig + got := cfg.Copy() + if got != nil { + t.Errorf("Copy() of nil = %v, want nil", got) + } +} + +func TestAuthConfig_Copy(t *testing.T) { + scheme := &HTTPScheme{ + Scheme: "bearer", + BearerFormat: "JWT", + } + cfg := &AuthConfig{ + AuthScheme: scheme, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeHTTP, + HTTP: &HTTPAuth{ + Scheme: "bearer", + Credentials: &HTTPCredentials{ + Token: "raw-token", + }, + }, + }, + ExchangedAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeHTTP, + HTTP: &HTTPAuth{ + Scheme: "bearer", + Credentials: &HTTPCredentials{ + Token: "exchanged-token", + }, + }, + }, + CredentialKey: "adk_test_key", + } + + got := cfg.Copy() + + if got == cfg { + t.Error("Copy() returned same pointer") + } + if got.RawAuthCredential == cfg.RawAuthCredential { + t.Error("Copy() returned same RawAuthCredential pointer") + } + if got.ExchangedAuthCredential == cfg.ExchangedAuthCredential { + t.Error("Copy() returned same ExchangedAuthCredential pointer") + } + if diff := cmp.Diff(cfg, got); diff != "" { + t.Errorf("Copy() mismatch (-want +got):\n%s", diff) + } +} + +func TestAuthConfig_CopyDeepCopiesAuthScheme(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Description: "orig", + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: "https://example.com/token", + Scopes: map[string]string{ + "repo": "read", + }, + }, + }, + }, + } + + got := cfg.Copy() + + orig := cfg.AuthScheme.(*OAuth2Scheme) + orig.Description = "mutated" + orig.Flows.ClientCredentials.Scopes["repo"] = "write" + + copied := got.AuthScheme.(*OAuth2Scheme) + if copied.Description != "orig" { + t.Fatalf("copied.Description = %q, want %q", copied.Description, "orig") + } + if copied.Flows.ClientCredentials.Scopes["repo"] != "read" { + t.Fatalf("copied scope = %q, want %q", copied.Flows.ClientCredentials.Scopes["repo"], "read") + } +} + +func TestNewAuthConfig_MarshalError(t *testing.T) { + scheme := &badScheme{} + if _, err := NewAuthConfig(scheme, nil); err == nil { + t.Fatal("NewAuthConfig() did not return error for unmarshalable scheme") + } +} + +type badScheme struct{} + +func (b *badScheme) GetType() SecuritySchemeType { + return SecuritySchemeType("bad") +} + +func (b *badScheme) MarshalJSON() ([]byte, error) { + return nil, errors.New("cannot marshal") +} diff --git a/auth/auth_credential.go b/auth/auth_credential.go new file mode 100644 index 000000000..d92320fca --- /dev/null +++ b/auth/auth_credential.go @@ -0,0 +1,182 @@ +// 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 + +// AuthCredentialType defines the type of credential. +type AuthCredentialType string + +const ( + // AuthCredentialTypeAPIKey for API key credentials. + AuthCredentialTypeAPIKey AuthCredentialType = "apiKey" + // AuthCredentialTypeHTTP for HTTP credentials (Basic, Bearer, etc). + AuthCredentialTypeHTTP AuthCredentialType = "http" + // AuthCredentialTypeOAuth2 for OAuth2 credentials. + AuthCredentialTypeOAuth2 AuthCredentialType = "oauth2" + // AuthCredentialTypeOpenIDConnect for OpenID Connect credentials. + AuthCredentialTypeOpenIDConnect AuthCredentialType = "openIdConnect" + // AuthCredentialTypeServiceAccount for Google Service Account credentials. + AuthCredentialTypeServiceAccount AuthCredentialType = "serviceAccount" +) + +// AuthCredential holds authentication credentials. +// The actual credential data is stored in one of the type-specific fields +// based on AuthType. +type AuthCredential struct { + // AuthType specifies which credential field to use. + AuthType AuthCredentialType `json:"authType"` + // ResourceRef is used to reference credentials from external sources. + ResourceRef string `json:"resourceRef,omitempty"` + // APIKey contains the API key value (for AuthCredentialTypeAPIKey). + APIKey string `json:"apiKey,omitempty"` + // HTTP contains HTTP auth credentials (for AuthCredentialTypeHTTP). + HTTP *HTTPAuth `json:"http,omitempty"` + // OAuth2 contains OAuth2 credentials (for AuthCredentialTypeOAuth2 and OpenIDConnect). + OAuth2 *OAuth2Auth `json:"oauth2,omitempty"` + // ServiceAccount contains Google service account credentials. + ServiceAccount *ServiceAccount `json:"serviceAccount,omitempty"` +} + +// HTTPAuth contains HTTP authentication credentials. +type HTTPAuth struct { + // Scheme is the HTTP authentication scheme (e.g., "basic", "bearer"). + Scheme string `json:"scheme"` + // Credentials contains the actual credential values. + Credentials *HTTPCredentials `json:"credentials"` +} + +// HTTPCredentials contains username/password or token for HTTP auth. +type HTTPCredentials struct { + Username string `json:"username,omitempty"` + Password string `json:"password,omitempty"` + Token string `json:"token,omitempty"` +} + +// OAuth2Auth contains OAuth2 credentials and tokens. +type OAuth2Auth struct { + // ClientID is the OAuth2 client ID. + ClientID string `json:"clientId,omitempty"` + // ClientSecret is the OAuth2 client secret. + ClientSecret string `json:"clientSecret,omitempty"` + // AuthURI is the authorization URI (generated by ADK or provided). + AuthURI string `json:"authUri,omitempty"` + // State is the OAuth state parameter for CSRF protection. + State string `json:"state,omitempty"` + // RedirectURI is the callback URI for OAuth flow. + RedirectURI string `json:"redirectUri,omitempty"` + // AuthResponseURI is the response from the OAuth provider after auth. + AuthResponseURI string `json:"authResponseUri,omitempty"` + // AuthCode is the authorization code from OAuth flow. + AuthCode string `json:"authCode,omitempty"` + // AccessToken is the obtained access token. + AccessToken string `json:"accessToken,omitempty"` + // RefreshToken is the refresh token for obtaining new access tokens. + RefreshToken string `json:"refreshToken,omitempty"` + // ExpiresAt is the Unix timestamp when the access token expires. + ExpiresAt int64 `json:"expiresAt,omitempty"` + // ExpiresIn is the token validity duration in seconds. + ExpiresIn int64 `json:"expiresIn,omitempty"` + // Audience is the intended audience for the token. + Audience string `json:"audience,omitempty"` + // TokenEndpointAuthMethod specifies how to authenticate at the token endpoint. + TokenEndpointAuthMethod string `json:"tokenEndpointAuthMethod,omitempty"` +} + +// ServiceAccount contains Google service account credentials. +type ServiceAccount struct { + // ServiceAccountCredential contains the parsed JSON key file. + ServiceAccountCredential *ServiceAccountCredential `json:"serviceAccountCredential,omitempty"` + // Scopes are the OAuth scopes to request. + Scopes []string `json:"scopes"` + // UseDefaultCredential indicates whether to use Application Default Credentials. + UseDefaultCredential bool `json:"useDefaultCredential,omitempty"` +} + +// ServiceAccountCredential represents a Google Service Account JSON key file. +type ServiceAccountCredential struct { + Type string `json:"type"` + ProjectID string `json:"project_id"` + PrivateKeyID string `json:"private_key_id"` + PrivateKey string `json:"private_key"` + ClientEmail string `json:"client_email"` + ClientID string `json:"client_id"` + AuthURI string `json:"auth_uri"` + TokenURI string `json:"token_uri"` + AuthProviderX509CertURL string `json:"auth_provider_x509_cert_url"` + ClientX509CertURL string `json:"client_x509_cert_url"` + UniverseDomain string `json:"universe_domain"` +} + +// Copy creates a deep copy of the AuthCredential. +func (c *AuthCredential) Copy() *AuthCredential { + if c == nil { + return nil + } + newCred := &AuthCredential{ + AuthType: c.AuthType, + ResourceRef: c.ResourceRef, + APIKey: c.APIKey, + } + if c.HTTP != nil { + newCred.HTTP = &HTTPAuth{ + Scheme: c.HTTP.Scheme, + } + if c.HTTP.Credentials != nil { + newCred.HTTP.Credentials = &HTTPCredentials{ + Username: c.HTTP.Credentials.Username, + Password: c.HTTP.Credentials.Password, + Token: c.HTTP.Credentials.Token, + } + } + } + if c.OAuth2 != nil { + newCred.OAuth2 = &OAuth2Auth{ + ClientID: c.OAuth2.ClientID, + ClientSecret: c.OAuth2.ClientSecret, + AuthURI: c.OAuth2.AuthURI, + State: c.OAuth2.State, + RedirectURI: c.OAuth2.RedirectURI, + AuthResponseURI: c.OAuth2.AuthResponseURI, + AuthCode: c.OAuth2.AuthCode, + AccessToken: c.OAuth2.AccessToken, + RefreshToken: c.OAuth2.RefreshToken, + ExpiresAt: c.OAuth2.ExpiresAt, + ExpiresIn: c.OAuth2.ExpiresIn, + Audience: c.OAuth2.Audience, + TokenEndpointAuthMethod: c.OAuth2.TokenEndpointAuthMethod, + } + } + if c.ServiceAccount != nil { + newCred.ServiceAccount = &ServiceAccount{ + Scopes: append([]string{}, c.ServiceAccount.Scopes...), + UseDefaultCredential: c.ServiceAccount.UseDefaultCredential, + } + if c.ServiceAccount.ServiceAccountCredential != nil { + newCred.ServiceAccount.ServiceAccountCredential = &ServiceAccountCredential{ + Type: c.ServiceAccount.ServiceAccountCredential.Type, + ProjectID: c.ServiceAccount.ServiceAccountCredential.ProjectID, + PrivateKeyID: c.ServiceAccount.ServiceAccountCredential.PrivateKeyID, + PrivateKey: c.ServiceAccount.ServiceAccountCredential.PrivateKey, + ClientEmail: c.ServiceAccount.ServiceAccountCredential.ClientEmail, + ClientID: c.ServiceAccount.ServiceAccountCredential.ClientID, + AuthURI: c.ServiceAccount.ServiceAccountCredential.AuthURI, + TokenURI: c.ServiceAccount.ServiceAccountCredential.TokenURI, + AuthProviderX509CertURL: c.ServiceAccount.ServiceAccountCredential.AuthProviderX509CertURL, + ClientX509CertURL: c.ServiceAccount.ServiceAccountCredential.ClientX509CertURL, + UniverseDomain: c.ServiceAccount.ServiceAccountCredential.UniverseDomain, + } + } + } + return newCred +} diff --git a/auth/auth_credential_test.go b/auth/auth_credential_test.go new file mode 100644 index 000000000..9862c5c9c --- /dev/null +++ b/auth/auth_credential_test.go @@ -0,0 +1,132 @@ +// 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 ( + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestAuthCredential_Copy_Nil(t *testing.T) { + var cred *AuthCredential + got := cred.Copy() + if got != nil { + t.Errorf("Copy() of nil = %v, want nil", got) + } +} + +func TestAuthCredential_Copy_APIKey(t *testing.T) { + cred := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + ResourceRef: "my-ref", + APIKey: "secret-key", + } + got := cred.Copy() + + if got == cred { + t.Error("Copy() returned same pointer") + } + if diff := cmp.Diff(cred, got); diff != "" { + t.Errorf("Copy() mismatch (-want +got):\n%s", diff) + } +} + +func TestAuthCredential_Copy_HTTP(t *testing.T) { + cred := &AuthCredential{ + AuthType: AuthCredentialTypeHTTP, + HTTP: &HTTPAuth{ + Scheme: "bearer", + Credentials: &HTTPCredentials{ + Username: "user", + Password: "pass", + Token: "token123", + }, + }, + } + got := cred.Copy() + + if got == cred { + t.Error("Copy() returned same pointer") + } + if got.HTTP == cred.HTTP { + t.Error("Copy() returned same HTTP pointer") + } + if got.HTTP.Credentials == cred.HTTP.Credentials { + t.Error("Copy() returned same Credentials pointer") + } + if diff := cmp.Diff(cred, got); diff != "" { + t.Errorf("Copy() mismatch (-want +got):\n%s", diff) + } +} + +func TestAuthCredential_Copy_OAuth2(t *testing.T) { + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + AuthURI: "https://example.com/auth", + State: "random-state", + RedirectURI: "https://example.com/callback", + AccessToken: "access-token", + RefreshToken: "refresh-token", + ExpiresAt: 1234567890, + }, + } + got := cred.Copy() + + if got == cred { + t.Error("Copy() returned same pointer") + } + if got.OAuth2 == cred.OAuth2 { + t.Error("Copy() returned same OAuth2 pointer") + } + if diff := cmp.Diff(cred, got); diff != "" { + t.Errorf("Copy() mismatch (-want +got):\n%s", diff) + } +} + +func TestAuthCredential_Copy_ServiceAccount(t *testing.T) { + cred := &AuthCredential{ + AuthType: AuthCredentialTypeServiceAccount, + ServiceAccount: &ServiceAccount{ + Scopes: []string{"scope1", "scope2"}, + UseDefaultCredential: true, + ServiceAccountCredential: &ServiceAccountCredential{ + Type: "service_account", + ProjectID: "my-project", + PrivateKey: "private-key-data", + ClientEmail: "sa@example.iam.gserviceaccount.com", + }, + }, + } + got := cred.Copy() + + if got == cred { + t.Error("Copy() returned same pointer") + } + if got.ServiceAccount == cred.ServiceAccount { + t.Error("Copy() returned same ServiceAccount pointer") + } + if got.ServiceAccount.ServiceAccountCredential == cred.ServiceAccount.ServiceAccountCredential { + t.Error("Copy() returned same ServiceAccountCredential pointer") + } + // Verify scopes are deep copied + got.ServiceAccount.Scopes[0] = "modified" + if cred.ServiceAccount.Scopes[0] == "modified" { + t.Error("Copy() did not deep copy Scopes slice") + } +} diff --git a/auth/auth_handler.go b/auth/auth_handler.go new file mode 100644 index 000000000..53f1f8678 --- /dev/null +++ b/auth/auth_handler.go @@ -0,0 +1,176 @@ +// 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 ( + "crypto/rand" + "encoding/hex" + "fmt" + + "golang.org/x/oauth2" +) + +var randRead = rand.Read + +// AuthHandler handles the OAuth flow orchestration including auth URI generation +// and response parsing. +type AuthHandler struct { + authConfig *AuthConfig +} + +// NewAuthHandler creates a new AuthHandler. +func NewAuthHandler(config *AuthConfig) *AuthHandler { + return &AuthHandler{ + authConfig: config, + } +} + +// GenerateAuthRequest generates an AuthConfig with the auth_uri populated for +// OAuth2/OIDC flows. The client uses this to redirect users for authorization. +func (h *AuthHandler) GenerateAuthRequest() (*AuthConfig, error) { + // For non-OAuth schemes, return a copy as-is + if h.authConfig.AuthScheme == nil { + return h.authConfig.Copy(), nil + } + + schemeType := h.authConfig.AuthScheme.GetType() + if schemeType != SecuritySchemeTypeOAuth2 && schemeType != SecuritySchemeTypeOpenIDConnect { + return h.authConfig.Copy(), nil + } + + // If auth_uri already exists in exchanged credential, return as-is + if h.authConfig.ExchangedAuthCredential != nil && + h.authConfig.ExchangedAuthCredential.OAuth2 != nil && + h.authConfig.ExchangedAuthCredential.OAuth2.AuthURI != "" { + return h.authConfig.Copy(), nil + } + + // Check if raw_auth_credential exists with client credentials + if h.authConfig.RawAuthCredential == nil || + h.authConfig.RawAuthCredential.OAuth2 == nil || + h.authConfig.RawAuthCredential.OAuth2.ClientID == "" { + return h.authConfig.Copy(), nil + } + + // If auth_uri already in raw credential, copy to exchanged + if h.authConfig.RawAuthCredential.OAuth2.AuthURI != "" { + return &AuthConfig{ + AuthScheme: h.authConfig.AuthScheme, + RawAuthCredential: h.authConfig.RawAuthCredential, + ExchangedAuthCredential: h.authConfig.RawAuthCredential.Copy(), + CredentialKey: h.authConfig.CredentialKey, + }, nil + } + + // Generate new auth URI + exchangedCred, err := h.generateAuthURI() + if err != nil { + return nil, err + } + if exchangedCred == nil { + return h.authConfig.Copy(), nil + } + + return &AuthConfig{ + AuthScheme: h.authConfig.AuthScheme, + RawAuthCredential: h.authConfig.RawAuthCredential, + ExchangedAuthCredential: exchangedCred, + CredentialKey: h.authConfig.CredentialKey, + }, nil +} + +// generateAuthURI generates the OAuth authorization URI. +func (h *AuthHandler) generateAuthURI() (*AuthCredential, error) { + cred := h.authConfig.RawAuthCredential + if cred == nil || cred.OAuth2 == nil { + return nil, nil + } + + authURL, scopes := authorizationMetadata(h.authConfig.AuthScheme) + if authURL == "" { + return nil, nil + } + + config := &oauth2.Config{ + ClientID: cred.OAuth2.ClientID, + ClientSecret: cred.OAuth2.ClientSecret, + Endpoint: oauth2.Endpoint{ + AuthURL: authURL, + }, + RedirectURL: cred.OAuth2.RedirectURI, + Scopes: scopes, + } + opts := []oauth2.AuthCodeOption{oauth2.AccessTypeOffline} + if cred.OAuth2.Audience != "" { + opts = append(opts, oauth2.SetAuthURLParam("audience", cred.OAuth2.Audience)) + } + + state, err := generateRandomState() + if err != nil { + return nil, err + } + authURI := config.AuthCodeURL(state, opts...) + + exchanged := cred.Copy() + exchanged.OAuth2.AuthURI = authURI + exchanged.OAuth2.State = state + + return exchanged, nil +} + +// GetAuthResponse retrieves the auth response from session state. +func (h *AuthHandler) GetAuthResponse(stateGetter func(key string) interface{}) *AuthCredential { + key := "temp:" + h.authConfig.CredentialKey + if val := stateGetter(key); val != nil { + if cred, ok := val.(*AuthCredential); ok { + return cred + } + } + return nil +} + +// generateRandomState generates a random state string for OAuth CSRF protection. +func generateRandomState() (string, error) { + b := make([]byte, 16) + if _, err := randRead(b); err != nil { + return "", fmt.Errorf("failed to generate random OAuth state: %w", err) + } + return hex.EncodeToString(b), nil +} + +// authorizationMetadata returns the authorization endpoint and scopes for OAuth2/OIDC schemes. +func authorizationMetadata(scheme AuthScheme) (string, []string) { + switch v := scheme.(type) { + case *OAuth2Scheme: + if v == nil || v.Flows == nil { + return "", nil + } + if v.Flows.AuthorizationCode != nil { + return v.Flows.AuthorizationCode.AuthorizationURL, scopeKeys(v.Flows.AuthorizationCode.Scopes) + } + if v.Flows.Implicit != nil { + return v.Flows.Implicit.AuthorizationURL, scopeKeys(v.Flows.Implicit.Scopes) + } + case *OpenIDConnectScheme: + if v == nil { + return "", nil + } + if len(v.Scopes) == 0 { + return v.AuthorizationEndpoint, []string{"openid"} + } + return v.AuthorizationEndpoint, append([]string{}, v.Scopes...) + } + return "", nil +} diff --git a/auth/auth_handler_test.go b/auth/auth_handler_test.go new file mode 100644 index 000000000..29bffb628 --- /dev/null +++ b/auth/auth_handler_test.go @@ -0,0 +1,318 @@ +// 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 ( + "crypto/rand" + "errors" + "net/url" + "strings" + "testing" +) + +func TestNewAuthHandler(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{}, + } + handler := NewAuthHandler(cfg) + if handler == nil { + t.Error("NewAuthHandler returned nil") + } +} + +func TestAuthHandler_GenerateAuthRequest_NonOAuth(t *testing.T) { + // For non-OAuth schemes, should return a copy as-is + cfg := &AuthConfig{ + AuthScheme: &APIKeyScheme{ + In: APIKeyInHeader, + Name: "X-API-Key", + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: "my-key", + }, + } + handler := NewAuthHandler(cfg) + result, err := handler.GenerateAuthRequest() + if err != nil { + t.Fatalf("GenerateAuthRequest() error = %v", err) + } + + if result == cfg { + t.Error("GenerateAuthRequest should return a copy, not same pointer") + } + if result.RawAuthCredential.APIKey != "my-key" { + t.Error("API key should be preserved") + } +} + +func TestAuthHandler_GenerateAuthRequest_OAuth2_NoCredential(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + }, + // No RawAuthCredential + } + handler := NewAuthHandler(cfg) + result, err := handler.GenerateAuthRequest() + if err != nil { + t.Fatalf("GenerateAuthRequest() error = %v", err) + } + + // Should return copy without generating auth URI + if result.ExchangedAuthCredential != nil && result.ExchangedAuthCredential.OAuth2 != nil && + result.ExchangedAuthCredential.OAuth2.AuthURI != "" { + t.Error("Should not generate auth URI without raw credential") + } +} + +func TestAuthHandler_GenerateAuthRequest_OAuth2_WithCredential(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + Scopes: map[string]string{"read": "Read access"}, + }, + }, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + RedirectURI: "https://localhost/callback", + }, + }, + } + handler := NewAuthHandler(cfg) + result, err := handler.GenerateAuthRequest() + if err != nil { + t.Fatalf("GenerateAuthRequest() error = %v", err) + } + + if result.ExchangedAuthCredential == nil { + t.Fatal("ExchangedAuthCredential should be set") + } + if result.ExchangedAuthCredential.OAuth2 == nil { + t.Fatal("ExchangedAuthCredential.OAuth2 should be set") + } + if result.ExchangedAuthCredential.OAuth2.AuthURI == "" { + t.Error("AuthURI should be generated") + } + if !strings.Contains(result.ExchangedAuthCredential.OAuth2.AuthURI, "https://example.com/auth") { + t.Errorf("AuthURI = %q, should contain authorization URL", result.ExchangedAuthCredential.OAuth2.AuthURI) + } + if result.ExchangedAuthCredential.OAuth2.State == "" { + t.Error("State should be generated for CSRF protection") + } +} + +func TestAuthHandler_GenerateAuthRequest_OAuth2_ExistingAuthURI(t *testing.T) { + // If auth_uri already exists in exchanged credential, return as-is + existingAuthURI := "https://example.com/existing-auth" + 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", + }, + }, + ExchangedAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AuthURI: existingAuthURI, + }, + }, + } + handler := NewAuthHandler(cfg) + result, err := handler.GenerateAuthRequest() + if err != nil { + t.Fatalf("GenerateAuthRequest() error = %v", err) + } + + if result.ExchangedAuthCredential.OAuth2.AuthURI != existingAuthURI { + t.Errorf("AuthURI = %q, want %q (should preserve existing)", result.ExchangedAuthCredential.OAuth2.AuthURI, existingAuthURI) + } +} + +func TestAuthHandler_GenerateAuthURI_OpenIDConnect(t *testing.T) { + t.Cleanup(func() { randRead = rand.Read }) + randRead = func(b []byte) (int, error) { + for i := range b { + b[i] = byte(i) + } + return len(b), nil + } + + cfg := &AuthConfig{ + AuthScheme: &OpenIDConnectScheme{ + AuthorizationEndpoint: "https://example.com/oauth2/authorize", + TokenEndpoint: "https://example.com/oauth2/token", + Scopes: []string{"openid", "profile"}, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOpenIDConnect, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + RedirectURI: "https://localhost/callback", + }, + }, + } + + handler := NewAuthHandler(cfg) + cred, err := handler.generateAuthURI() + if err != nil { + t.Fatalf("generateAuthURI() error = %v", err) + } + if cred == nil || cred.OAuth2 == nil { + t.Fatal("generateAuthURI() returned nil") + } + + parsed, err := url.Parse(cred.OAuth2.AuthURI) + if err != nil { + t.Fatalf("parse auth URI: %v", err) + } + q := parsed.Query() + + if got := q.Get("client_id"); got != "client-id" { + t.Fatalf("client_id = %s, want client-id", got) + } + if got := q.Get("scope"); got != "openid profile" { + t.Fatalf("scope = %s, want 'openid profile'", got) + } + if got := q.Get("access_type"); got != "offline" { + t.Fatalf("access_type = %s, want offline", got) + } + if got := q.Get("state"); got != "000102030405060708090a0b0c0d0e0f" { + t.Fatalf("state = %s, want deterministic hex value", got) + } +} + +func TestAuthHandler_GenerateAuthURI_IncludesAudience(t *testing.T) { + t.Cleanup(func() { randRead = rand.Read }) + randRead = func(b []byte) (int, error) { + for i := range b { + b[i] = 0xAB + } + return len(b), nil + } + + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/oauth2/authorize", + TokenURL: "https://example.com/oauth2/token", + Scopes: map[string]string{"read": "Read access"}, + }, + }, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + RedirectURI: "https://localhost/callback", + Audience: "https://example.com/audience", + }, + }, + } + + handler := NewAuthHandler(cfg) + cred, err := handler.generateAuthURI() + if err != nil { + t.Fatalf("generateAuthURI() error = %v", err) + } + if cred == nil { + t.Fatal("generateAuthURI() returned nil") + } + + parsed, err := url.Parse(cred.OAuth2.AuthURI) + if err != nil { + t.Fatalf("parse auth URI: %v", err) + } + if got := parsed.Query().Get("audience"); got != "https://example.com/audience" { + t.Fatalf("audience = %s, want https://example.com/audience", got) + } +} + +func TestGenerateRandomState_Error(t *testing.T) { + t.Cleanup(func() { randRead = rand.Read }) + randRead = func([]byte) (int, error) { + return 0, errors.New("entropy exhausted") + } + + if _, err := generateRandomState(); err == nil { + t.Fatal("generateRandomState() did not return error on rand failure") + } +} + +func TestAuthHandler_GetAuthResponse(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "test-key", + } + handler := NewAuthHandler(cfg) + + expectedCred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "token123", + }, + } + + stateGetter := func(key string) interface{} { + if key == "temp:test-key" { + return expectedCred + } + return nil + } + + result := handler.GetAuthResponse(stateGetter) + if result != expectedCred { + t.Errorf("GetAuthResponse() = %v, want %v", result, expectedCred) + } +} + +func TestAuthHandler_GetAuthResponse_NotFound(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "test-key", + } + handler := NewAuthHandler(cfg) + + stateGetter := func(key string) interface{} { + return nil + } + + result := handler.GetAuthResponse(stateGetter) + if result != nil { + t.Errorf("GetAuthResponse() = %v, want nil", result) + } +} diff --git a/auth/auth_scheme.go b/auth/auth_scheme.go new file mode 100644 index 000000000..9f6bf8263 --- /dev/null +++ b/auth/auth_scheme.go @@ -0,0 +1,146 @@ +// 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 provides authentication types and utilities for ADK tools. +// It follows OpenAPI 3.0 Security Scheme specifications. +package auth + +// SecuritySchemeType defines the type of security scheme. +// See: https://swagger.io/specification/#security-scheme-object +type SecuritySchemeType string + +const ( + // SecuritySchemeTypeAPIKey for API key authentication. + SecuritySchemeTypeAPIKey SecuritySchemeType = "apiKey" + // SecuritySchemeTypeHTTP for HTTP authentication (Basic, Bearer, etc). + SecuritySchemeTypeHTTP SecuritySchemeType = "http" + // SecuritySchemeTypeOAuth2 for OAuth 2.0 authentication. + SecuritySchemeTypeOAuth2 SecuritySchemeType = "oauth2" + // SecuritySchemeTypeOpenIDConnect for OpenID Connect authentication. + SecuritySchemeTypeOpenIDConnect SecuritySchemeType = "openIdConnect" +) + +// APIKeyIn defines where the API key is located. +type APIKeyIn string + +const ( + // APIKeyInQuery for API key in query parameter. + APIKeyInQuery APIKeyIn = "query" + // APIKeyInHeader for API key in HTTP header. + APIKeyInHeader APIKeyIn = "header" + // APIKeyInCookie for API key in cookie. + APIKeyInCookie APIKeyIn = "cookie" +) + +// AuthScheme is the interface for all security schemes. +type AuthScheme interface { + // GetType returns the security scheme type. + GetType() SecuritySchemeType +} + +// APIKeyScheme represents API Key authentication. +// See: https://swagger.io/docs/specification/authentication/api-keys/ +type APIKeyScheme struct { + In APIKeyIn `json:"in"` + Name string `json:"name"` + Description string `json:"description,omitempty"` +} + +// GetType implements AuthScheme. +func (s *APIKeyScheme) GetType() SecuritySchemeType { + return SecuritySchemeTypeAPIKey +} + +// HTTPScheme represents HTTP authentication (Basic, Bearer, etc). +// See: https://swagger.io/docs/specification/authentication/basic-authentication/ +type HTTPScheme struct { + Scheme string `json:"scheme"` // "basic", "bearer", "digest", etc. + BearerFormat string `json:"bearerFormat,omitempty"` // e.g., "JWT" + Description string `json:"description,omitempty"` +} + +// GetType implements AuthScheme. +func (s *HTTPScheme) GetType() SecuritySchemeType { + return SecuritySchemeTypeHTTP +} + +// OAuth2Scheme represents OAuth 2.0 authentication. +// See: https://swagger.io/docs/specification/authentication/oauth2/ +type OAuth2Scheme struct { + Flows *OAuthFlows `json:"flows"` + Description string `json:"description,omitempty"` +} + +// GetType implements AuthScheme. +func (s *OAuth2Scheme) GetType() SecuritySchemeType { + return SecuritySchemeTypeOAuth2 +} + +// OAuthFlows contains OAuth2 flow configurations. +type OAuthFlows struct { + Implicit *OAuthFlowImplicit `json:"implicit,omitempty"` + Password *OAuthFlowPassword `json:"password,omitempty"` + ClientCredentials *OAuthFlowClientCredentials `json:"clientCredentials,omitempty"` + AuthorizationCode *OAuthFlowAuthorizationCode `json:"authorizationCode,omitempty"` +} + +// OAuthFlowImplicit represents the OAuth2 Implicit flow. +type OAuthFlowImplicit struct { + AuthorizationURL string `json:"authorizationUrl"` + RefreshURL string `json:"refreshUrl,omitempty"` + Scopes map[string]string `json:"scopes"` +} + +// OAuthFlowPassword represents the OAuth2 Resource Owner Password flow. +type OAuthFlowPassword struct { + TokenURL string `json:"tokenUrl"` + RefreshURL string `json:"refreshUrl,omitempty"` + Scopes map[string]string `json:"scopes"` +} + +// OAuthFlowClientCredentials represents the OAuth2 Client Credentials flow. +type OAuthFlowClientCredentials struct { + TokenURL string `json:"tokenUrl"` + RefreshURL string `json:"refreshUrl,omitempty"` + Scopes map[string]string `json:"scopes"` +} + +// OAuthFlowAuthorizationCode represents the OAuth2 Authorization Code flow. +type OAuthFlowAuthorizationCode struct { + AuthorizationURL string `json:"authorizationUrl"` + TokenURL string `json:"tokenUrl"` + RefreshURL string `json:"refreshUrl,omitempty"` + Scopes map[string]string `json:"scopes"` +} + +// OpenIDConnectScheme represents OpenID Connect authentication. +// This is an extended version that includes flattened OIDC configuration, +// similar to Python ADK's OpenIdConnectWithConfig. +type OpenIDConnectScheme struct { + // OpenIDConnectURL is the standard OIDC discovery URL. + OpenIDConnectURL string `json:"openIdConnectUrl,omitempty"` + // Flattened OIDC configuration (for when discovery is not available). + AuthorizationEndpoint string `json:"authorizationEndpoint,omitempty"` + TokenEndpoint string `json:"tokenEndpoint,omitempty"` + UserInfoEndpoint string `json:"userinfoEndpoint,omitempty"` + RevocationEndpoint string `json:"revocationEndpoint,omitempty"` + GrantTypesSupported []string `json:"grantTypesSupported,omitempty"` + Scopes []string `json:"scopes,omitempty"` + Description string `json:"description,omitempty"` +} + +// GetType implements AuthScheme. +func (s *OpenIDConnectScheme) GetType() SecuritySchemeType { + return SecuritySchemeTypeOpenIDConnect +} diff --git a/auth/auth_scheme_test.go b/auth/auth_scheme_test.go new file mode 100644 index 000000000..5b6f30003 --- /dev/null +++ b/auth/auth_scheme_test.go @@ -0,0 +1,68 @@ +// 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 "testing" + +func TestAPIKeyScheme_GetType(t *testing.T) { + scheme := &APIKeyScheme{ + In: APIKeyInHeader, + Name: "X-API-Key", + } + if got := scheme.GetType(); got != SecuritySchemeTypeAPIKey { + t.Errorf("GetType() = %v, want %v", got, SecuritySchemeTypeAPIKey) + } +} + +func TestHTTPScheme_GetType(t *testing.T) { + scheme := &HTTPScheme{ + Scheme: "bearer", + BearerFormat: "JWT", + } + if got := scheme.GetType(); got != SecuritySchemeTypeHTTP { + t.Errorf("GetType() = %v, want %v", got, SecuritySchemeTypeHTTP) + } +} + +func TestOAuth2Scheme_GetType(t *testing.T) { + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + } + if got := scheme.GetType(); got != SecuritySchemeTypeOAuth2 { + t.Errorf("GetType() = %v, want %v", got, SecuritySchemeTypeOAuth2) + } +} + +func TestOpenIDConnectScheme_GetType(t *testing.T) { + scheme := &OpenIDConnectScheme{ + OpenIDConnectURL: "https://example.com/.well-known/openid-configuration", + } + if got := scheme.GetType(); got != SecuritySchemeTypeOpenIDConnect { + t.Errorf("GetType() = %v, want %v", got, SecuritySchemeTypeOpenIDConnect) + } +} + +func TestAuthScheme_Interface(t *testing.T) { + // Verify all scheme types implement AuthScheme interface + var _ AuthScheme = (*APIKeyScheme)(nil) + var _ AuthScheme = (*HTTPScheme)(nil) + var _ AuthScheme = (*OAuth2Scheme)(nil) + var _ AuthScheme = (*OpenIDConnectScheme)(nil) +} diff --git a/auth/auth_tool_arguments.go b/auth/auth_tool_arguments.go new file mode 100644 index 000000000..1b863a860 --- /dev/null +++ b/auth/auth_tool_arguments.go @@ -0,0 +1,25 @@ +// 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 + +// AuthToolArguments represents the arguments for the special long-running +// function tool that is used to request end user credentials. +// This matches Python ADK's AuthToolArguments in auth_tool.py:93 +type AuthToolArguments struct { + // FunctionCallID is the ID of the original function call that requested auth. + FunctionCallID string `json:"function_call_id"` + // AuthConfig is the auth configuration for the tool. + AuthConfig *AuthConfig `json:"auth_config"` +} diff --git a/auth/credential_manager.go b/auth/credential_manager.go new file mode 100644 index 000000000..095b7568d --- /dev/null +++ b/auth/credential_manager.go @@ -0,0 +1,241 @@ +// 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" + "fmt" +) + +// CredentialManager orchestrates the complete lifecycle of authentication +// credentials, from initial loading to final preparation for use. +type CredentialManager struct { + authConfig *AuthConfig + exchangerRegistry *ExchangerRegistry + refresherRegistry *RefresherRegistry +} + +// NewCredentialManager creates a new CredentialManager with default exchangers +// and refreshers registered. +func NewCredentialManager(cfg *AuthConfig) *CredentialManager { + m := &CredentialManager{ + authConfig: cfg, + exchangerRegistry: NewExchangerRegistry(), + refresherRegistry: NewRefresherRegistry(), + } + + // Register default exchangers + oauth2Exchanger := NewOAuth2Exchanger() + m.exchangerRegistry.Register(AuthCredentialTypeOAuth2, oauth2Exchanger) + m.exchangerRegistry.Register(AuthCredentialTypeOpenIDConnect, oauth2Exchanger) + + // Register default refreshers + oauth2Refresher := NewOAuth2Refresher() + m.refresherRegistry.Register(AuthCredentialTypeOAuth2, oauth2Refresher) + m.refresherRegistry.Register(AuthCredentialTypeOpenIDConnect, oauth2Refresher) + + return m +} + +// GetAuthCredential loads and prepares authentication credential through a +// structured workflow: +// 1. Validate credential configuration +// 2. Check if credential is already ready (no processing needed) +// 3. Try to load existing processed credential from CredentialService +// 4. If no existing credential, load from auth response (temp: prefix) +// 5. For client credentials flow, use raw credentials directly +// 6. Exchange credential if needed (e.g., auth code -> access token) +// 7. Refresh credential if expired +// 8. Save credential to CredentialService if it was modified +func (m *CredentialManager) GetAuthCredential(ctx context.Context, stateGetter func(key string) interface{}, credentialService ...CredentialService) (*AuthCredential, error) { + // Step 1: Validate credential configuration + if err := m.validate(); err != nil { + return nil, err + } + + // Step 2: Check if credential is already ready (no processing needed) + if m.isCredentialReady() { + return m.authConfig.RawAuthCredential, nil + } + + // Step 3: Try to load existing processed credential + credential := m.loadExistingCredential() + + // Step 3b: Try to load from credential service (persistent storage) + var svc CredentialService + if len(credentialService) > 0 && credentialService[0] != nil { + svc = credentialService[0] + if credential == nil { + loaded, err := svc.LoadCredential(ctx, m.authConfig) + if err != nil { + return nil, fmt.Errorf("failed to load credential: %w", err) + } + if loaded != nil && loaded.OAuth2 != nil && loaded.OAuth2.AccessToken != "" { + credential = loaded + } + } + } + + // Step 4: If no existing credential, load from auth response (session state) + wasFromAuthResponse := false + if credential == nil && stateGetter != nil { + credential = m.loadFromAuthResponse(stateGetter) + if credential != nil { + wasFromAuthResponse = true + } + } + + // Step 5: If still no credential available, check if client credentials + if credential == nil { + if m.isClientCredentialsFlow() { + credential = m.authConfig.RawAuthCredential + } else { + // For authorization code flow, return nil to trigger user authorization + return nil, nil + } + } + + // Step 6: Exchange credential if needed + credential, wasExchanged, err := m.exchangeCredential(ctx, credential) + if err != nil { + return nil, err + } + + // Step 7: Refresh credential if expired + wasRefreshed := false + if !wasExchanged { + var err error + credential, wasRefreshed, err = m.refreshCredential(ctx, credential) + if err != nil { + return nil, err + } + } + + // Step 8: Save credential if it was modified + if wasFromAuthResponse || wasExchanged || wasRefreshed { + m.authConfig.ExchangedAuthCredential = credential + // Save to credential service for persistence across requests + if svc != nil { + if err := svc.SaveCredential(ctx, m.authConfig); err != nil { + return nil, fmt.Errorf("failed to save credential: %w", err) + } + } + } + + return credential, nil +} + +// validate checks that the auth configuration is valid. +func (m *CredentialManager) validate() error { + // For OAuth2/OIDC, raw_auth_credential is required + if m.authConfig.AuthScheme != nil { + schemeType := m.authConfig.AuthScheme.GetType() + if schemeType == SecuritySchemeTypeOAuth2 || schemeType == SecuritySchemeTypeOpenIDConnect { + if m.authConfig.RawAuthCredential == nil { + return nil // Will need user auth + } + } + } + return nil +} + +// isCredentialReady checks if credential is ready to use without further processing. +func (m *CredentialManager) isCredentialReady() bool { + raw := m.authConfig.RawAuthCredential + if raw == nil { + return false + } + + // Simple credentials that don't need exchange or refresh + switch raw.AuthType { + case AuthCredentialTypeAPIKey, AuthCredentialTypeHTTP: + return true + } + + return false +} + +// loadExistingCredential loads credential from exchanged cache. +func (m *CredentialManager) loadExistingCredential() *AuthCredential { + if m.authConfig.ExchangedAuthCredential != nil { + return m.authConfig.ExchangedAuthCredential + } + return nil +} + +// loadFromAuthResponse loads credential from session state (auth response). +func (m *CredentialManager) loadFromAuthResponse(stateGetter func(key string) interface{}) *AuthCredential { + key := "temp:" + m.authConfig.CredentialKey + if val := stateGetter(key); val != nil { + if cred, ok := val.(*AuthCredential); ok { + return cred + } + } + return nil +} + +// isClientCredentialsFlow checks if the auth scheme uses client credentials flow. +func (m *CredentialManager) isClientCredentialsFlow() bool { + switch scheme := m.authConfig.AuthScheme.(type) { + case *OAuth2Scheme: + if scheme.Flows == nil { + return false + } + return scheme.Flows.ClientCredentials != nil + case *OpenIDConnectScheme: + return grantSupported(scheme.GrantTypesSupported, "client_credentials") + default: + return false + } +} + +// exchangeCredential exchanges credential if needed. +func (m *CredentialManager) exchangeCredential(ctx context.Context, cred *AuthCredential) (*AuthCredential, bool, error) { + ex := m.exchangerRegistry.Get(cred.AuthType) + if ex == nil { + return cred, false, nil + } + + result, err := ex.Exchange(ctx, cred, m.authConfig.AuthScheme) + if err != nil { + return nil, false, err + } + + return result.Credential, result.WasExchanged, nil +} + +// refreshCredential refreshes credential if expired. +func (m *CredentialManager) refreshCredential(ctx context.Context, cred *AuthCredential) (*AuthCredential, bool, error) { + ref := m.refresherRegistry.Get(cred.AuthType) + if ref == nil { + return cred, false, nil + } + + if !ref.IsRefreshNeeded(cred, m.authConfig.AuthScheme) { + return cred, false, nil + } + + refreshed, err := ref.Refresh(ctx, cred, m.authConfig.AuthScheme) + if err != nil { + return cred, false, fmt.Errorf("failed to refresh credential: %w", err) + } + + return refreshed, true, nil +} + +// GetAuthConfig accessor +func (m *CredentialManager) GetAuthConfig() *AuthConfig { + return m.authConfig +} diff --git a/auth/credential_manager_test.go b/auth/credential_manager_test.go new file mode 100644 index 000000000..f4053b890 --- /dev/null +++ b/auth/credential_manager_test.go @@ -0,0 +1,418 @@ +// 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 +} + +type stubExchanger struct { + calls int + result *ExchangeResult + err error +} + +func (s *stubExchanger) Exchange(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*ExchangeResult, error) { + s.calls++ + if s.err != nil { + return nil, s.err + } + return s.result, nil +} + +type stubCredentialService struct { + loadResp *AuthCredential + loadErr error + saved []*AuthConfig + saveErr error +} + +func (s *stubCredentialService) LoadCredential(ctx context.Context, cfg *AuthConfig) (*AuthCredential, error) { + if s.loadErr != nil { + return nil, s.loadErr + } + return s.loadResp, nil +} + +func (s *stubCredentialService) SaveCredential(ctx context.Context, cfg *AuthConfig) error { + if s.saveErr != nil { + return s.saveErr + } + s.saved = append(s.saved, cfg.Copy()) + return 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) + } +} + +func TestCredentialManager_GetAuthCredential_ReadyAPIKey(t *testing.T) { + cfg := &AuthConfig{ + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: "key-123", + }, + } + + manager := NewCredentialManager(cfg) + cred, err := manager.GetAuthCredential(context.Background(), nil) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if cred != cfg.RawAuthCredential { + t.Fatalf("got %v, want raw credential", cred) + } +} + +func TestCredentialManager_GetAuthCredential_ClientCredentialsFlow(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "adk_client", + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: "https://example.com/token", + }, + }, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client", + ClientSecret: "secret", + }, + }, + } + + manager := NewCredentialManager(cfg) + exchanger := &stubExchanger{ + result: &ExchangeResult{ + Credential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "token-123", + }, + }, + WasExchanged: true, + }, + } + manager.exchangerRegistry = NewExchangerRegistry() + manager.exchangerRegistry.Register(AuthCredentialTypeOAuth2, exchanger) + + svc := &stubCredentialService{} + cred, err := manager.GetAuthCredential(context.Background(), nil, svc) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if exchanger.calls != 1 { + t.Fatalf("Exchange called %d times, want 1", exchanger.calls) + } + if cred.OAuth2.AccessToken != "token-123" { + t.Fatalf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "token-123") + } + if len(svc.saved) != 1 { + t.Fatalf("expected credential to be saved, got %d saves", len(svc.saved)) + } +} + +func TestCredentialManager_GetAuthCredential_AuthCodeExchange(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "adk_auth_code", + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client", + ClientSecret: "secret", + }, + }, + } + + manager := NewCredentialManager(cfg) + exchanger := &stubExchanger{ + result: &ExchangeResult{ + Credential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "new-token", + }, + }, + WasExchanged: true, + }, + } + manager.exchangerRegistry = NewExchangerRegistry() + manager.exchangerRegistry.Register(AuthCredentialTypeOAuth2, exchanger) + + stateGetter := func(key string) interface{} { + if key == "temp:"+cfg.CredentialKey { + return &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client", + ClientSecret: "secret", + AuthCode: "code-123", + }, + } + } + return nil + } + + cred, err := manager.GetAuthCredential(context.Background(), stateGetter) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if exchanger.calls != 1 { + t.Fatalf("Exchange called %d times, want 1", exchanger.calls) + } + if cred.OAuth2.AccessToken != "new-token" { + t.Fatalf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "new-token") + } +} + +func TestCredentialManager_GetAuthCredential_LoadsFromCredentialService(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: "https://example.com/token", + }, + }, + }, + RawAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + }, + } + + manager := NewCredentialManager(cfg) + svc := &stubCredentialService{ + loadResp: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "svc-token", + }, + }, + } + + cred, err := manager.GetAuthCredential(context.Background(), nil, svc) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if cred.OAuth2.AccessToken != "svc-token" { + t.Fatalf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "svc-token") + } + if len(svc.saved) != 0 { + t.Fatalf("expected no saves, got %d", len(svc.saved)) + } +} + +func TestCredentialManager_GetAuthCredential_LoadsFromAuthResponse(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "adk_temp", + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + }, + } + + manager := NewCredentialManager(cfg) + stateCred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "state-token", + }, + } + stateGetter := func(key string) interface{} { + if key == "temp:"+cfg.CredentialKey { + return stateCred + } + return nil + } + + cred, err := manager.GetAuthCredential(context.Background(), stateGetter) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if cred.OAuth2.AccessToken != "state-token" { + t.Fatalf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "state-token") + } +} + +func TestCredentialManager_GetAuthCredential_ReturnsNilWhenAuthorizationNeeded(t *testing.T) { + cfg := &AuthConfig{ + AuthScheme: &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + }, + } + + manager := NewCredentialManager(cfg) + cred, err := manager.GetAuthCredential(context.Background(), nil) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if cred != nil { + t.Fatalf("expected nil credential when user auth is required, got %v", cred) + } +} + +func TestCredentialManager_GetAuthCredential_RefreshesCredential(t *testing.T) { + cfg := &AuthConfig{ + CredentialKey: "adk_refresh", + AuthScheme: &OAuth2Scheme{}, + ExchangedAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "expired", + RefreshToken: "refresh", + ExpiresAt: time.Now().Add(-time.Minute).Unix(), + }, + }, + } + + manager := NewCredentialManager(cfg) + manager.refresherRegistry = NewRefresherRegistry() + manager.refresherRegistry.Register(AuthCredentialTypeOAuth2, &stubRefresher{ + shouldRefresh: true, + refreshed: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "fresh", + }, + }, + }) + + svc := &stubCredentialService{} + cred, err := manager.GetAuthCredential(context.Background(), nil, svc) + if err != nil { + t.Fatalf("GetAuthCredential() unexpected error: %v", err) + } + if cred.OAuth2.AccessToken != "fresh" { + t.Fatalf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "fresh") + } + if len(svc.saved) != 1 { + t.Fatalf("expected refreshed credential to be saved, got %d", len(svc.saved)) + } +} diff --git a/auth/credential_service.go b/auth/credential_service.go new file mode 100644 index 000000000..8e1bdbde4 --- /dev/null +++ b/auth/credential_service.go @@ -0,0 +1,101 @@ +// 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" + "sync" +) + +// CredentialService is an interface for loading and saving credentials. +type CredentialService interface { + // LoadCredential loads a credential from storage. + LoadCredential(ctx context.Context, config *AuthConfig) (*AuthCredential, error) + + // SaveCredential saves a credential to storage. + SaveCredential(ctx context.Context, config *AuthConfig) error +} + +// InMemoryCredentialService stores credentials in memory. +type InMemoryCredentialService struct { + mu sync.RWMutex + credentials map[string]*AuthCredential +} + +// NewInMemoryCredentialService creates a new in-memory credential service. +func NewInMemoryCredentialService() *InMemoryCredentialService { + return &InMemoryCredentialService{ + credentials: make(map[string]*AuthCredential), + } +} + +// LoadCredential loads a credential from memory. +func (s *InMemoryCredentialService) LoadCredential(ctx context.Context, config *AuthConfig) (*AuthCredential, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + if cred, ok := s.credentials[config.CredentialKey]; ok { + return cred.Copy(), nil + } + return nil, nil +} + +// SaveCredential saves a credential to memory. +func (s *InMemoryCredentialService) SaveCredential(ctx context.Context, config *AuthConfig) error { + s.mu.Lock() + defer s.mu.Unlock() + + if config.ExchangedAuthCredential != nil { + s.credentials[config.CredentialKey] = config.ExchangedAuthCredential.Copy() + } + return nil +} + +// SessionStateCredentialService stores credentials in session state. +type SessionStateCredentialService struct { + stateSetter func(key string, value interface{}) + stateGetter func(key string) interface{} +} + +// NewSessionStateCredentialService creates a new session state credential service. +func NewSessionStateCredentialService( + getter func(key string) interface{}, + setter func(key string, value interface{}), +) *SessionStateCredentialService { + return &SessionStateCredentialService{ + stateGetter: getter, + stateSetter: setter, + } +} + +// LoadCredential loads a credential from session state. +func (s *SessionStateCredentialService) LoadCredential(ctx context.Context, config *AuthConfig) (*AuthCredential, error) { + key := "cred:" + config.CredentialKey + if val := s.stateGetter(key); val != nil { + if cred, ok := val.(*AuthCredential); ok { + return cred.Copy(), nil + } + } + return nil, nil +} + +// SaveCredential saves a credential to session state. +func (s *SessionStateCredentialService) SaveCredential(ctx context.Context, config *AuthConfig) error { + if config.ExchangedAuthCredential != nil { + key := "cred:" + config.CredentialKey + s.stateSetter(key, config.ExchangedAuthCredential.Copy()) + } + return nil +} diff --git a/auth/credential_service_test.go b/auth/credential_service_test.go new file mode 100644 index 000000000..fe33f8f0c --- /dev/null +++ b/auth/credential_service_test.go @@ -0,0 +1,109 @@ +// 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" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestInMemoryCredentialService_SaveAndLoad(t *testing.T) { + svc := NewInMemoryCredentialService() + ctx := context.Background() + + cfg := &AuthConfig{ + CredentialKey: "test-key", + ExchangedAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "access-token", + RefreshToken: "refresh-token", + }, + }, + } + + // Save + if err := svc.SaveCredential(ctx, cfg); err != nil { + t.Fatalf("SaveCredential() error = %v", err) + } + + // Load + loaded, err := svc.LoadCredential(ctx, cfg) + if err != nil { + t.Fatalf("LoadCredential() error = %v", err) + } + if loaded == nil { + t.Fatal("LoadCredential() returned nil") + } + if diff := cmp.Diff(cfg.ExchangedAuthCredential, loaded); diff != "" { + t.Errorf("LoadCredential() mismatch (-want +got):\n%s", diff) + } +} + +func TestInMemoryCredentialService_LoadCredential_NotFound(t *testing.T) { + svc := NewInMemoryCredentialService() + ctx := context.Background() + + cfg := &AuthConfig{ + CredentialKey: "non-existent-key", + } + + loaded, err := svc.LoadCredential(ctx, cfg) + if err != nil { + t.Fatalf("LoadCredential() error = %v", err) + } + if loaded != nil { + t.Errorf("LoadCredential() = %v, want nil for non-existent key", loaded) + } +} + +// Note: SaveCredential requires non-nil config - passing nil will panic. + +func TestInMemoryCredentialService_Overwrite(t *testing.T) { + svc := NewInMemoryCredentialService() + ctx := context.Background() + + cfg := &AuthConfig{ + CredentialKey: "test-key", + ExchangedAuthCredential: &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + AccessToken: "token-1", + }, + }, + } + + // Save first credential + if err := svc.SaveCredential(ctx, cfg); err != nil { + t.Fatalf("SaveCredential() error = %v", err) + } + + // Overwrite with new credential + cfg.ExchangedAuthCredential.OAuth2.AccessToken = "token-2" + if err := svc.SaveCredential(ctx, cfg); err != nil { + t.Fatalf("SaveCredential() error = %v", err) + } + + // Load should return new credential + loaded, err := svc.LoadCredential(ctx, cfg) + if err != nil { + t.Fatalf("LoadCredential() error = %v", err) + } + if loaded.OAuth2.AccessToken != "token-2" { + t.Errorf("AccessToken = %q, want %q", loaded.OAuth2.AccessToken, "token-2") + } +} diff --git a/auth/helpers.go b/auth/helpers.go new file mode 100644 index 000000000..79da2ce0d --- /dev/null +++ b/auth/helpers.go @@ -0,0 +1,177 @@ +// 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 ( + "encoding/json" + "fmt" +) + +// TokenToSchemeCredential creates an API key auth scheme and credential. +// This is a helper function similar to Python ADK's token_to_scheme_credential. +// +// Parameters: +// - schemeType: The scheme type ("apikey" or "http") +// - location: Where the token is sent ("header", "query", or "cookie") +// - name: The name of the header/query parameter +// - token: The actual token/API key value +// +// Returns the auth scheme and credential pair. +func TokenToSchemeCredential(schemeType, location, name, token string) (AuthScheme, *AuthCredential) { + switch schemeType { + case "apikey": + var in APIKeyIn + switch location { + case "query": + in = APIKeyInQuery + case "cookie": + in = APIKeyInCookie + default: + in = APIKeyInHeader + } + scheme := &APIKeyScheme{ + In: in, + Name: name, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: token, + } + return scheme, cred + case "http": + scheme := &HTTPScheme{ + Scheme: "bearer", + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeHTTP, + HTTP: &HTTPAuth{ + Scheme: "bearer", + Credentials: &HTTPCredentials{ + Token: token, + }, + }, + } + return scheme, cred + default: + // Default to API key in header + scheme := &APIKeyScheme{ + In: APIKeyInHeader, + Name: name, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeAPIKey, + APIKey: token, + } + return scheme, cred + } +} + +// BearerTokenCredential creates an HTTP Bearer token auth scheme and credential. +func BearerTokenCredential(token string) (AuthScheme, *AuthCredential) { + scheme := &HTTPScheme{ + Scheme: "bearer", + BearerFormat: "JWT", + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeHTTP, + HTTP: &HTTPAuth{ + Scheme: "bearer", + Credentials: &HTTPCredentials{ + Token: token, + }, + }, + } + return scheme, cred +} + +// OAuth2ClientCredentials creates an OAuth2 client credentials auth scheme and credential. +func OAuth2ClientCredentials(clientID, clientSecret, tokenURL string, scopes map[string]string) (AuthScheme, *AuthCredential) { + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: tokenURL, + Scopes: scopes, + }, + }, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: clientID, + ClientSecret: clientSecret, + }, + } + return scheme, cred +} + +// OAuth2AuthorizationCode creates an OAuth2 authorization code auth scheme and credential. +func OAuth2AuthorizationCode(clientID, clientSecret, authURL, tokenURL string, scopes map[string]string) (AuthScheme, *AuthCredential) { + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + AuthorizationCode: &OAuthFlowAuthorizationCode{ + AuthorizationURL: authURL, + TokenURL: tokenURL, + Scopes: scopes, + }, + }, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: clientID, + ClientSecret: clientSecret, + }, + } + return scheme, cred +} + +// ServiceAccountCredentials creates a service account auth scheme and credential. +func ServiceAccountCredentials(credentialJSON []byte, scopes []string) (AuthScheme, *AuthCredential, error) { + scheme := &HTTPScheme{ + Scheme: "bearer", + BearerFormat: "JWT", + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeServiceAccount, + ServiceAccount: &ServiceAccount{ + Scopes: scopes, + }, + } + // Parse the JSON if provided + if len(credentialJSON) > 0 { + var saCred ServiceAccountCredential + if err := json.Unmarshal(credentialJSON, &saCred); err != nil { + return nil, nil, fmt.Errorf("failed to parse service account credential: %w", err) + } + cred.ServiceAccount.ServiceAccountCredential = &saCred + } + return scheme, cred, nil +} + +// DefaultCredentials creates a credential using Application Default Credentials. +func DefaultCredentials(scopes []string) (AuthScheme, *AuthCredential) { + scheme := &HTTPScheme{ + Scheme: "bearer", + BearerFormat: "JWT", + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeServiceAccount, + ServiceAccount: &ServiceAccount{ + Scopes: scopes, + UseDefaultCredential: true, + }, + } + return scheme, cred +} diff --git a/auth/helpers_test.go b/auth/helpers_test.go new file mode 100644 index 000000000..623cfb472 --- /dev/null +++ b/auth/helpers_test.go @@ -0,0 +1,45 @@ +// 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 "testing" + +func TestServiceAccountCredentialsParse(t *testing.T) { + jsonCred := []byte(`{ + "type": "service_account", + "client_email": "service@example.com", + "token_uri": "https://oauth2.example.com/token" + }`) + + scheme, cred, err := ServiceAccountCredentials(jsonCred, []string{"scope1"}) + if err != nil { + t.Fatalf("ServiceAccountCredentials() error = %v", err) + } + if scheme == nil { + t.Fatal("scheme is nil") + } + if cred.ServiceAccount == nil || cred.ServiceAccount.ServiceAccountCredential == nil { + t.Fatal("service account credential was not parsed") + } + if got := cred.ServiceAccount.ServiceAccountCredential.ClientEmail; got != "service@example.com" { + t.Fatalf("client email = %s, want service@example.com", got) + } +} + +func TestServiceAccountCredentialsInvalidJSON(t *testing.T) { + if _, _, err := ServiceAccountCredentials([]byte("{invalid json"), nil); err == nil { + t.Fatal("expected error for invalid JSON") + } +} diff --git a/auth/oauth2.go b/auth/oauth2.go new file mode 100644 index 000000000..ca9d19908 --- /dev/null +++ b/auth/oauth2.go @@ -0,0 +1,328 @@ +// 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" + "fmt" + "strings" + "time" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/clientcredentials" +) + +// OAuth2Exchanger exchanges OAuth2 credentials. +// It handles both authorization code and client credentials flows. +type OAuth2Exchanger struct{} + +// NewOAuth2Exchanger creates a new OAuth2 exchanger. +func NewOAuth2Exchanger() *OAuth2Exchanger { + return &OAuth2Exchanger{} +} + +// Exchange exchanges OAuth2 credentials for access tokens. +func (e *OAuth2Exchanger) Exchange(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*ExchangeResult, error) { + if scheme == nil { + return nil, fmt.Errorf("auth scheme is required for OAuth2 credential exchange") + } + + // If already has access token, no exchange needed + if cred.OAuth2 != nil && cred.OAuth2.AccessToken != "" { + return &ExchangeResult{Credential: cred, WasExchanged: false}, nil + } + + // Determine grant type from scheme + grantType := e.determineGrantType(scheme) + + switch grantType { + case GrantTypeClientCredentials: + return e.exchangeClientCredentials(ctx, cred, scheme) + case GrantTypeAuthorizationCode: + return e.exchangeAuthorizationCode(ctx, cred, scheme) + default: + // Unknown grant type, return unchanged + return &ExchangeResult{Credential: cred, WasExchanged: false}, nil + } +} + +// GrantType represents OAuth2 grant types. +type GrantType string + +const ( + GrantTypeAuthorizationCode GrantType = "authorization_code" + GrantTypeClientCredentials GrantType = "client_credentials" +) + +func (e *OAuth2Exchanger) determineGrantType(scheme AuthScheme) GrantType { + switch s := scheme.(type) { + case *OAuth2Scheme: + if s.Flows == nil { + return "" + } + if s.Flows.ClientCredentials != nil { + return GrantTypeClientCredentials + } + if s.Flows.AuthorizationCode != nil { + return GrantTypeAuthorizationCode + } + case *OpenIDConnectScheme: + if grantSupported(s.GrantTypesSupported, "client_credentials") { + return GrantTypeClientCredentials + } + if grantSupported(s.GrantTypesSupported, "authorization_code") || s.AuthorizationEndpoint != "" { + return GrantTypeAuthorizationCode + } + } + return "" +} + +func (e *OAuth2Exchanger) exchangeClientCredentials(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*ExchangeResult, error) { + if cred.OAuth2 == nil { + return nil, fmt.Errorf("oauth2 credentials required") + } + + tokenURL, scopes := clientCredentialsMetadata(scheme) + if tokenURL == "" { + return nil, fmt.Errorf("client credentials flow not configured in scheme") + } + + conf := &clientcredentials.Config{ + ClientID: cred.OAuth2.ClientID, + ClientSecret: cred.OAuth2.ClientSecret, + TokenURL: tokenURL, + Scopes: scopes, + } + + if cred.OAuth2.TokenEndpointAuthMethod == "client_secret_post" { + conf.AuthStyle = oauth2.AuthStyleInParams + } + if cred.OAuth2.Audience != "" { + conf.EndpointParams = map[string][]string{ + "audience": {cred.OAuth2.Audience}, + } + } + + token, err := conf.Token(ctx) + if err != nil { + return nil, fmt.Errorf("failed to exchange client credentials: %w", err) + } + + newCred := cred.Copy() + newCred.OAuth2.AccessToken = token.AccessToken + newCred.OAuth2.RefreshToken = token.RefreshToken + if !token.Expiry.IsZero() { + newCred.OAuth2.ExpiresAt = token.Expiry.Unix() + newCred.OAuth2.ExpiresIn = int64(time.Until(token.Expiry).Seconds()) + } + + return &ExchangeResult{Credential: newCred, WasExchanged: true}, nil +} + +func (e *OAuth2Exchanger) exchangeAuthorizationCode(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*ExchangeResult, error) { + if cred.OAuth2 == nil { + return nil, fmt.Errorf("oauth2 credentials required") + } + + authURL, tokenURL, scopes := authorizationCodeMetadataFromScheme(scheme) + if authURL == "" || tokenURL == "" { + return nil, fmt.Errorf("authorization code flow not configured in scheme") + } + + // Need auth_code to exchange + if cred.OAuth2.AuthCode == "" { + return &ExchangeResult{Credential: cred, WasExchanged: false}, nil + } + + config := &oauth2.Config{ + ClientID: cred.OAuth2.ClientID, + ClientSecret: cred.OAuth2.ClientSecret, + Endpoint: oauth2.Endpoint{ + AuthURL: authURL, + TokenURL: tokenURL, + }, + RedirectURL: cred.OAuth2.RedirectURI, + Scopes: scopes, + } + + token, err := config.Exchange(ctx, cred.OAuth2.AuthCode) + if err != nil { + return nil, fmt.Errorf("failed to exchange authorization code: %w", err) + } + + // Update credential with tokens + newCred := cred.Copy() + newCred.OAuth2.AccessToken = token.AccessToken + newCred.OAuth2.RefreshToken = token.RefreshToken + newCred.OAuth2.AuthCode = "" // Clear the code after use + if !token.Expiry.IsZero() { + newCred.OAuth2.ExpiresAt = token.Expiry.Unix() + newCred.OAuth2.ExpiresIn = int64(time.Until(token.Expiry).Seconds()) + } + + return &ExchangeResult{Credential: newCred, WasExchanged: true}, nil +} + +func scopeKeys(scopes map[string]string) []string { + if scopes == nil { + return nil + } + keys := make([]string, 0, len(scopes)) + for k := range scopes { + keys = append(keys, k) + } + return keys +} + +func clientCredentialsMetadata(scheme AuthScheme) (string, []string) { + switch s := scheme.(type) { + case *OAuth2Scheme: + if s.Flows == nil || s.Flows.ClientCredentials == nil { + return "", nil + } + return s.Flows.ClientCredentials.TokenURL, scopeKeys(s.Flows.ClientCredentials.Scopes) + case *OpenIDConnectScheme: + if s.TokenEndpoint == "" { + return "", nil + } + if len(s.Scopes) == 0 { + return s.TokenEndpoint, []string{"openid"} + } + return s.TokenEndpoint, append([]string{}, s.Scopes...) + default: + return "", nil + } +} + +func authorizationCodeMetadataFromScheme(scheme AuthScheme) (string, string, []string) { + switch s := scheme.(type) { + case *OAuth2Scheme: + if s.Flows == nil || s.Flows.AuthorizationCode == nil { + return "", "", nil + } + return s.Flows.AuthorizationCode.AuthorizationURL, s.Flows.AuthorizationCode.TokenURL, + scopeKeys(s.Flows.AuthorizationCode.Scopes) + case *OpenIDConnectScheme: + if s.AuthorizationEndpoint == "" || s.TokenEndpoint == "" { + return "", "", nil + } + if len(s.Scopes) == 0 { + return s.AuthorizationEndpoint, s.TokenEndpoint, []string{"openid"} + } + return s.AuthorizationEndpoint, s.TokenEndpoint, append([]string{}, s.Scopes...) + default: + return "", "", nil + } +} + +func tokenEndpointFromScheme(scheme AuthScheme) string { + switch s := scheme.(type) { + case *OAuth2Scheme: + if s.Flows == nil { + return "" + } + if s.Flows.AuthorizationCode != nil { + return s.Flows.AuthorizationCode.TokenURL + } + if s.Flows.ClientCredentials != nil { + return s.Flows.ClientCredentials.TokenURL + } + case *OpenIDConnectScheme: + return s.TokenEndpoint + } + return "" +} + +func grantSupported(grants []string, want string) bool { + want = strings.ToLower(want) + for _, g := range grants { + if strings.ToLower(g) == want { + return true + } + } + return false +} + +// OAuth2Refresher refreshes OAuth2 access tokens using refresh tokens. +type OAuth2Refresher struct{} + +// NewOAuth2Refresher creates a new OAuth2 refresher. +func NewOAuth2Refresher() *OAuth2Refresher { + return &OAuth2Refresher{} +} + +// IsRefreshNeeded checks if the OAuth2 token is expired or about to expire. +func (r *OAuth2Refresher) IsRefreshNeeded(cred *AuthCredential, scheme AuthScheme) bool { + if cred.OAuth2 == nil { + return false + } + + // No expiry info, assume valid + if cred.OAuth2.ExpiresAt == 0 { + return false + } + + // Check if expired (with 60 second buffer) + expiresAt := time.Unix(cred.OAuth2.ExpiresAt, 0) + return time.Now().Add(60 * time.Second).After(expiresAt) +} + +// Refresh refreshes the OAuth2 access token using the refresh token. +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 +} diff --git a/auth/oauth2_test.go b/auth/oauth2_test.go new file mode 100644 index 000000000..1c9e1b04a --- /dev/null +++ b/auth/oauth2_test.go @@ -0,0 +1,223 @@ +// 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" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func TestOAuth2Exchanger_ClientCredentialsFlow(t *testing.T) { + t.Parallel() + + var gotAuthHeader string + var gotAudience string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + if grant := r.FormValue("grant_type"); grant != "client_credentials" { + t.Fatalf("grant_type = %s, want client_credentials", grant) + } + gotAudience = r.FormValue("audience") + gotAuthHeader = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"token123","token_type":"bearer","expires_in":60}`) + })) + defer server.Close() + + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: server.URL, + Scopes: map[string]string{"read": "Read access"}, + }, + }, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + Audience: "https://api.example.com", + }, + } + + ex := NewOAuth2Exchanger() + res, err := ex.exchangeClientCredentials(context.Background(), cred, scheme) + if err != nil { + t.Fatalf("exchangeClientCredentials() error = %v", err) + } + if res.Credential.OAuth2.AccessToken != "token123" { + t.Fatalf("access token = %s, want token123", res.Credential.OAuth2.AccessToken) + } + if gotAudience != "https://api.example.com" { + t.Fatalf("audience = %s, want https://api.example.com", gotAudience) + } + if gotAuthHeader == "" || !strings.HasPrefix(gotAuthHeader, "Basic ") { + t.Fatalf("authorization header = %s, want HTTP Basic credentials", gotAuthHeader) + } +} + +func TestOAuth2Exchanger_ClientCredentials_ClientSecretPost(t *testing.T) { + t.Parallel() + + var gotAuthHeader string + var clientID string + var clientSecret string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + clientID = r.FormValue("client_id") + clientSecret = r.FormValue("client_secret") + gotAuthHeader = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"token456","token_type":"bearer","expires_in":60}`) + })) + defer server.Close() + + scheme := &OAuth2Scheme{ + Flows: &OAuthFlows{ + ClientCredentials: &OAuthFlowClientCredentials{ + TokenURL: server.URL, + }, + }, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOAuth2, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + TokenEndpointAuthMethod: "client_secret_post", + }, + } + + ex := NewOAuth2Exchanger() + res, err := ex.exchangeClientCredentials(context.Background(), cred, scheme) + if err != nil { + t.Fatalf("exchangeClientCredentials() error = %v", err) + } + if res.Credential.OAuth2.AccessToken != "token456" { + t.Fatalf("access token = %s, want token456", res.Credential.OAuth2.AccessToken) + } + if gotAuthHeader != "" { + t.Fatalf("authorization header = %s, want empty for client_secret_post", gotAuthHeader) + } + if clientID != "client-id" || clientSecret != "client-secret" { + t.Fatalf("client credentials were not sent in body") + } +} + +func TestOAuth2Exchanger_AuthorizationCode_OpenID(t *testing.T) { + t.Parallel() + + var receivedCode string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + receivedCode = r.FormValue("code") + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"oidc-access","refresh_token":"oidc-refresh","token_type":"bearer","expires_in":120}`) + })) + defer server.Close() + + scheme := &OpenIDConnectScheme{ + AuthorizationEndpoint: "https://example.com/oauth2/authorize", + TokenEndpoint: server.URL, + Scopes: []string{"openid"}, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOpenIDConnect, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + RedirectURI: "https://localhost/callback", + AuthCode: "auth-code", + }, + } + + ex := NewOAuth2Exchanger() + res, err := ex.exchangeAuthorizationCode(context.Background(), cred, scheme) + if err != nil { + t.Fatalf("exchangeAuthorizationCode() error = %v", err) + } + if receivedCode != "auth-code" { + t.Fatalf("code = %s, want auth-code", receivedCode) + } + if res.Credential.OAuth2.AccessToken != "oidc-access" { + t.Fatalf("access token = %s, want oidc-access", res.Credential.OAuth2.AccessToken) + } + if res.Credential.OAuth2.RefreshToken != "oidc-refresh" { + t.Fatalf("refresh token = %s, want oidc-refresh", res.Credential.OAuth2.RefreshToken) + } + if res.Credential.OAuth2.AuthCode != "" { + t.Fatalf("auth code was not cleared") + } +} + +func TestOAuth2Refresher_OpenIDConnect(t *testing.T) { + t.Parallel() + + var grantType string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("ParseForm() error = %v", err) + } + grantType = r.FormValue("grant_type") + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"new-token","refresh_token":"new-refresh","token_type":"bearer","expires_in":60}`) + })) + defer server.Close() + + scheme := &OpenIDConnectScheme{ + TokenEndpoint: server.URL, + } + cred := &AuthCredential{ + AuthType: AuthCredentialTypeOpenIDConnect, + OAuth2: &OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + AccessToken: "old-token", + RefreshToken: "old-refresh", + ExpiresAt: time.Now().Add(-time.Minute).Unix(), + }, + } + + refresher := NewOAuth2Refresher() + newCred, err := refresher.Refresh(context.Background(), cred, scheme) + if err != nil { + t.Fatalf("Refresh() error = %v", err) + } + if grantType != "refresh_token" { + t.Fatalf("grant_type = %s, want refresh_token", grantType) + } + if newCred.OAuth2.AccessToken != "new-token" { + t.Fatalf("access token = %s, want new-token", newCred.OAuth2.AccessToken) + } + if newCred.OAuth2.RefreshToken != "new-refresh" { + t.Fatalf("refresh token = %s, want new-refresh", newCred.OAuth2.RefreshToken) + } +} diff --git a/auth/registries.go b/auth/registries.go new file mode 100644 index 000000000..5af90e422 --- /dev/null +++ b/auth/registries.go @@ -0,0 +1,99 @@ +// 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" + "sync" +) + +// ExchangeResult contains the result of a credential exchange. +type ExchangeResult struct { + // Credential is the exchanged credential. + Credential *AuthCredential + // WasExchanged indicates if the credential was actually exchanged. + WasExchanged bool +} + +// CredentialExchanger exchanges credentials from one form to another. +// For example, exchanging an authorization code for an access token. +type CredentialExchanger interface { + // Exchange exchanges the given credential using the auth scheme. + // Returns the exchanged credential and whether it was exchanged. + Exchange(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*ExchangeResult, error) +} + +// CredentialRefresher refreshes expired credentials. +type CredentialRefresher interface { + // IsRefreshNeeded checks if the credential needs to be refreshed. + IsRefreshNeeded(cred *AuthCredential, scheme AuthScheme) bool + + // Refresh refreshes the credential and returns the new credential. + Refresh(ctx context.Context, cred *AuthCredential, scheme AuthScheme) (*AuthCredential, error) +} + +// ExchangerRegistry manages credential exchangers by credential type. +type ExchangerRegistry struct { + mu sync.RWMutex + exchangers map[AuthCredentialType]CredentialExchanger +} + +// NewExchangerRegistry creates a new exchanger registry. +func NewExchangerRegistry() *ExchangerRegistry { + return &ExchangerRegistry{ + exchangers: make(map[AuthCredentialType]CredentialExchanger), + } +} + +// Register registers an exchanger for a credential type. +func (r *ExchangerRegistry) Register(credType AuthCredentialType, exchanger CredentialExchanger) { + r.mu.Lock() + defer r.mu.Unlock() + r.exchangers[credType] = exchanger +} + +// Get returns the exchanger for a credential type, or nil if not found. +func (r *ExchangerRegistry) Get(credType AuthCredentialType) CredentialExchanger { + r.mu.RLock() + defer r.mu.RUnlock() + return r.exchangers[credType] +} + +// RefresherRegistry manages credential refreshers by credential type. +type RefresherRegistry struct { + mu sync.RWMutex + refreshers map[AuthCredentialType]CredentialRefresher +} + +// NewRefresherRegistry creates a new refresher registry. +func NewRefresherRegistry() *RefresherRegistry { + return &RefresherRegistry{ + refreshers: make(map[AuthCredentialType]CredentialRefresher), + } +} + +// Register registers a refresher for a credential type. +func (r *RefresherRegistry) Register(credType AuthCredentialType, refresher CredentialRefresher) { + r.mu.Lock() + defer r.mu.Unlock() + r.refreshers[credType] = refresher +} + +// Get returns the refresher for a credential type, or nil if not found. +func (r *RefresherRegistry) Get(credType AuthCredentialType) CredentialRefresher { + r.mu.RLock() + defer r.mu.RUnlock() + return r.refreshers[credType] +} diff --git a/auth/request_euc.go b/auth/request_euc.go new file mode 100644 index 000000000..a39945483 --- /dev/null +++ b/auth/request_euc.go @@ -0,0 +1,20 @@ +// 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 + +// RequestEUCFunctionCallName is the name of the system function call +// used to request end-user credentials (EUC) for OAuth2 authorization. +// This matches Python ADK's REQUEST_EUC_FUNCTION_CALL_NAME. +const RequestEUCFunctionCallName = "adk_request_credential" diff --git a/go.mod b/go.mod index a770851cd..7569dfb3b 100644 --- a/go.mod +++ b/go.mod @@ -26,6 +26,7 @@ require ( github.com/glebarez/go-sqlite v1.21.1 // indirect github.com/mattn/go-isatty v0.0.17 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect modernc.org/libc v1.22.3 // indirect modernc.org/mathutil v1.5.0 // indirect modernc.org/memory v1.5.0 // indirect @@ -64,7 +65,7 @@ require ( go.opentelemetry.io/otel/sdk v1.38.0 go.opentelemetry.io/otel/sdk/metric v1.38.0 // indirect go.opentelemetry.io/otel/trace v1.38.0 - golang.org/x/oauth2 v0.32.0 + golang.org/x/oauth2 v0.34.0 golang.org/x/time v0.14.0 // indirect google.golang.org/genproto v0.0.0-20251014184007-4626949a642f // indirect google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f // indirect diff --git a/go.sum b/go.sum index 407b99394..fd1bd784e 100644 --- a/go.sum +++ b/go.sum @@ -145,6 +145,8 @@ golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= golang.org/x/oauth2 v0.32.0 h1:jsCblLleRMDrxMN29H3z/k1KliIvpLgCkE6R8FXXNgY= golang.org/x/oauth2 v0.32.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/oauth2 v0.34.0 h1:hqK/t4AKgbqWkdkcAeI8XLmbK+4m4G5YeQRrmiotGlw= +golang.org/x/oauth2 v0.34.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/internal/llminternal/auth_functions.go b/internal/llminternal/auth_functions.go new file mode 100644 index 000000000..4f31bf009 --- /dev/null +++ b/internal/llminternal/auth_functions.go @@ -0,0 +1,89 @@ +// 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 llminternal + +import ( + "github.com/google/uuid" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/auth" + "google.golang.org/adk/model" + "google.golang.org/adk/session" +) + +const afFunctionCallIDPrefix = "adk-" + +// generateFunctionCallID creates a unique function call ID. +// This matches Python's generate_client_function_call_id() with AF_FUNCTION_CALL_ID_PREFIX = 'adk-' +func generateFunctionCallID() string { + return afFunctionCallIDPrefix + uuid.NewString() +} + +// GenerateAuthEvent creates an event with adk_request_credential function calls +// from the RequestedAuthConfigs in the function response event. +// This matches Python ADK's generate_auth_event in flows/llm_flows/functions.py. +func GenerateAuthEvent(ctx agent.InvocationContext, fnResponseEvent *session.Event) *session.Event { + if fnResponseEvent == nil || len(fnResponseEvent.Actions.RequestedAuthConfigs) == 0 { + return nil + } + + var parts []*genai.Part + var longRunningToolIDs []string + + for functionCallID, authConfig := range fnResponseEvent.Actions.RequestedAuthConfigs { + // Create args map matching Python's AuthToolArguments.model_dump() + // Note: We preserve *auth.AuthConfig pointer since this is in-memory, + // matching Python's behavior where objects are passed by reference. + argsMap := map[string]any{ + "function_call_id": functionCallID, + "auth_config": authConfig, // Keep as *auth.AuthConfig pointer + } + + // Create the adk_request_credential function call + requestEucFunctionCall := &genai.FunctionCall{ + Name: auth.RequestEUCFunctionCallName, + Args: argsMap, + } + + // Generate a unique ID for this function call + requestEucFunctionCall.ID = generateFunctionCallID() + longRunningToolIDs = append(longRunningToolIDs, requestEucFunctionCall.ID) + + parts = append(parts, &genai.Part{ + FunctionCall: requestEucFunctionCall, + }) + } + + // Determine the role from the original event + role := "model" + if fnResponseEvent.Content != nil && fnResponseEvent.Content.Role != "" { + role = fnResponseEvent.Content.Role + } + + // Create the auth event + authEvent := session.NewEvent(ctx.InvocationID()) + authEvent.Author = ctx.Agent().Name() + authEvent.Branch = ctx.Branch() + authEvent.LLMResponse = model.LLMResponse{ + Content: &genai.Content{ + Role: role, + Parts: parts, + }, + } + authEvent.LongRunningToolIDs = longRunningToolIDs + + return authEvent +} diff --git a/internal/llminternal/auth_functions_test.go b/internal/llminternal/auth_functions_test.go new file mode 100644 index 000000000..f05e9044a --- /dev/null +++ b/internal/llminternal/auth_functions_test.go @@ -0,0 +1,197 @@ +// 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 llminternal + +import ( + "testing" + + "google.golang.org/adk/auth" + contextinternal "google.golang.org/adk/internal/context" + "google.golang.org/adk/session" +) + +func TestGenerateAuthEvent_Nil(t *testing.T) { + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) + + // Nil event + result := GenerateAuthEvent(inv, nil) + if result != nil { + t.Error("GenerateAuthEvent(nil) should return nil") + } + + // Empty RequestedAuthConfigs + event := &session.Event{ + Actions: session.EventActions{ + RequestedAuthConfigs: make(map[string]*auth.AuthConfig), + }, + } + result = GenerateAuthEvent(inv, event) + if result != nil { + t.Error("GenerateAuthEvent with empty RequestedAuthConfigs should return nil") + } +} + +// Note: TestGenerateAuthEvent_CreatesEvent and TestGenerateAuthEvent_MultipleCalls +// are skipped as they require full invocation context with agent setup. +// The GenerateAuthEvent function is tested indirectly through integration tests. + +func TestGenerateFunctionCallID(t *testing.T) { + id1 := generateFunctionCallID() + id2 := generateFunctionCallID() + + if id1 == "" { + t.Error("generateFunctionCallID() returned empty string") + } + if id1 == id2 { + t.Error("generateFunctionCallID() should return unique IDs") + } + if len(id1) < 4 || id1[:4] != "adk-" { + t.Errorf("generateFunctionCallID() = %q, should start with 'adk-'", id1) + } +} + +func TestParseAuthConfigFromMap(t *testing.T) { + data := map[string]any{ + "credential_key": "test-key", + "exchanged_auth_credential": map[string]any{ + "auth_type": "oauth2", + "oauth2": map[string]any{ + "access_token": "token123", + "refresh_token": "refresh456", + "expires_at": float64(1234567890), + }, + }, + } + + config, err := parseAuthConfigFromMap(data) + if err != nil { + t.Fatalf("parseAuthConfigFromMap() error = %v", err) + } + if config.CredentialKey != "test-key" { + t.Errorf("CredentialKey = %q, want %q", config.CredentialKey, "test-key") + } + if config.ExchangedAuthCredential == nil { + t.Fatal("ExchangedAuthCredential should not be nil") + } + if config.ExchangedAuthCredential.OAuth2 == nil { + t.Fatal("OAuth2 should not be nil") + } + if config.ExchangedAuthCredential.OAuth2.AccessToken != "token123" { + t.Errorf("AccessToken = %q, want %q", config.ExchangedAuthCredential.OAuth2.AccessToken, "token123") + } +} + +func TestParseAuthConfigFromMap_CamelCase(t *testing.T) { + data := map[string]any{ + "credentialKey": "camel-key", + "exchangedAuthCredential": map[string]any{ + "authType": "oauth2", + "oauth2": map[string]any{ + "accessToken": "camel_token", + }, + }, + } + + config, err := parseAuthConfigFromMap(data) + if err != nil { + t.Fatalf("parseAuthConfigFromMap() error = %v", err) + } + if config.CredentialKey != "camel-key" { + t.Errorf("CredentialKey = %q, want %q", config.CredentialKey, "camel-key") + } + if got := config.ExchangedAuthCredential.OAuth2.AccessToken; got != "camel_token" { + t.Errorf("AccessToken = %q, want %q", got, "camel_token") + } +} + +func TestParseAuthCredentialFromMap(t *testing.T) { + data := map[string]any{ + "auth_type": "oauth2", + "oauth2": map[string]any{ + "access_token": "access", + "refresh_token": "refresh", + "expires_at": float64(9999999999), + }, + } + + cred, err := parseAuthCredentialFromMap(data) + if err != nil { + t.Fatalf("parseAuthCredentialFromMap() error = %v", err) + } + if cred.AuthType != auth.AuthCredentialTypeOAuth2 { + t.Errorf("AuthType = %v, want %v", cred.AuthType, auth.AuthCredentialTypeOAuth2) + } + if cred.OAuth2.AccessToken != "access" { + t.Errorf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "access") + } + if cred.OAuth2.RefreshToken != "refresh" { + t.Errorf("RefreshToken = %q, want %q", cred.OAuth2.RefreshToken, "refresh") + } +} + +func TestParseAuthCredentialFromMap_WithHyphenKeys(t *testing.T) { + data := map[string]any{ + "auth-type": "oauth2", + "oauth2": map[string]any{ + "access-token": "hy-access", + "refresh-token": "hy-refresh", + }, + } + + cred, err := parseAuthCredentialFromMap(data) + if err != nil { + t.Fatalf("parseAuthCredentialFromMap() error = %v", err) + } + if cred.AuthType != auth.AuthCredentialTypeOAuth2 { + t.Errorf("AuthType = %v, want %v", cred.AuthType, auth.AuthCredentialTypeOAuth2) + } + if cred.OAuth2.AccessToken != "hy-access" { + t.Errorf("AccessToken = %q, want %q", cred.OAuth2.AccessToken, "hy-access") + } + if cred.OAuth2.RefreshToken != "hy-refresh" { + t.Errorf("RefreshToken = %q, want %q", cred.OAuth2.RefreshToken, "hy-refresh") + } +} + +func TestParseAuthCredentialFromMap_NotAMap(t *testing.T) { + _, err := parseAuthCredentialFromMap("not a map") + if err == nil { + t.Error("parseAuthCredentialFromMap() should error for non-map input") + } +} + +func TestAuthPreprocessorResult_Init(t *testing.T) { + // Verify CurrentAuthPreprocessorResult starts as nil + if CurrentAuthPreprocessorResult != nil { + // Clear it for test isolation + CurrentAuthPreprocessorResult = nil + } + + result := &AuthPreprocessorResult{ + ToolIdsToResume: make(map[string]bool), + CredentialsStored: true, + OriginalEvent: &session.Event{}, + } + + if result.ToolIdsToResume == nil { + t.Error("ToolIdsToResume should not be nil") + } + if !result.CredentialsStored { + t.Error("CredentialsStored should be true") + } + if result.OriginalEvent == nil { + t.Error("OriginalEvent should not be nil") + } +} diff --git a/internal/llminternal/base_flow.go b/internal/llminternal/base_flow.go index 15937e23a..ea16a6135 100644 --- a/internal/llminternal/base_flow.go +++ b/internal/llminternal/base_flow.go @@ -125,6 +125,38 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, if ctx.Ended() { return } + + // Check if auth preprocessor found tools that need to be re-executed. + // This implements the "Surgical Resumption" pattern from Python ADK. + if result := CurrentAuthPreprocessorResult; result != nil && result.OriginalEvent != nil && len(result.ToolIdsToResume) > 0 { + // Clear the result immediately to prevent re-processing + CurrentAuthPreprocessorResult = nil + + // Build tools map + tools := make(map[string]tool.Tool) + for k, v := range req.Tools { + if t, ok := v.(tool.Tool); ok { + tools[k] = t + } + } + + // Execute function calls from the original event that match our tools_to_resume + // This matches Python's handle_function_calls_async with tools_to_resume filter + fnResponseEvent, err := f.handleFunctionCalls(ctx, tools, &result.OriginalEvent.LLMResponse, result.ToolIdsToResume) + if err != nil { + yield(nil, err) + return + } + if fnResponseEvent != nil { + if !yield(fnResponseEvent, nil) { + return + } + } + + // Return after tool re-execution - Python does the same + return + } + spans := telemetry.StartTrace(ctx, "call_llm") // Create event to pass to callback state delta stateDelta := make(map[string]any) @@ -163,10 +195,8 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, if !yield(modelResponseEvent, nil) { return } - // TODO: generate and yield an auth event if needed. // Handle function calls. - ev, err := f.handleFunctionCalls(ctx, tools, resp) if err != nil { yield(nil, err) @@ -180,6 +210,14 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, return } + // Generate and yield an auth event if needed. + // This converts RequestedAuthConfigs into adk_request_credential function calls. + if authEvent := GenerateAuthEvent(ctx, ev); authEvent != nil { + if !yield(authEvent, nil) { + return + } + } + // Actually handle "transfer_to_agent" tool. The function call sets the ev.Actions.TransferToAgent field. // We are following python's execution flow which is // BaseLlmFlow._postprocess_async @@ -364,14 +402,25 @@ func findLongRunningFunctionCallIDs(c *genai.Content, tools map[string]tool.Tool } // handleFunctionCalls calls the functions and returns the function response event. +// If toolsToResume is non-nil and non-empty, only function calls with IDs in the map are executed. // -// TODO: accept filters to include/exclude function calls. // TODO: check feasibility of running tool.Run concurrently. -func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[string]tool.Tool, resp *model.LLMResponse) (*session.Event, error) { +func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[string]tool.Tool, resp *model.LLMResponse, toolsToResume ...map[string]bool) (*session.Event, error) { var fnResponseEvents []*session.Event + // Build filter map if provided + var filterMap map[string]bool + if len(toolsToResume) > 0 && toolsToResume[0] != nil && len(toolsToResume[0]) > 0 { + filterMap = toolsToResume[0] + } + fnCalls := utils.FunctionCalls(resp.Content) for _, fnCall := range fnCalls { + // Skip function calls not in the filter (if filter is provided) + if filterMap != nil && !filterMap[fnCall.ID] { + continue + } + curTool, ok := toolsDict[fnCall.Name] if !ok { return nil, fmt.Errorf("unknown tool: %q", fnCall.Name) @@ -380,6 +429,7 @@ func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[st if !ok { return nil, fmt.Errorf("tool %q is not a function tool", curTool.Name()) } + toolCtx := toolinternal.NewToolContext(ctx, fnCall.ID, &session.EventActions{StateDelta: make(map[string]any)}) // toolCtx := tool. spans := telemetry.StartTrace(ctx, "execute_tool "+fnCall.Name) diff --git a/internal/llminternal/other_processors.go b/internal/llminternal/other_processors.go index dd84a8395..158625550 100644 --- a/internal/llminternal/other_processors.go +++ b/internal/llminternal/other_processors.go @@ -15,8 +15,16 @@ package llminternal import ( + "fmt" + "strings" + "unicode" + + "github.com/mitchellh/mapstructure" "google.golang.org/adk/agent" + "google.golang.org/adk/auth" "google.golang.org/adk/model" + "google.golang.org/adk/session" + "google.golang.org/genai" ) func identityRequestProcessor(ctx agent.InvocationContext, req *model.LLMRequest) error { @@ -34,11 +42,228 @@ func codeExecutionRequestProcessor(ctx agent.InvocationContext, req *model.LLMRe return nil } +// AuthPreprocessorResult contains the result of auth preprocessing. +// It tells the Flow whether tools need to be re-executed. +type AuthPreprocessorResult struct { + // ToolIdsToResume contains the IDs of function calls that should be re-executed. + ToolIdsToResume map[string]bool + // CredentialsStored indicates if any credentials were stored. + CredentialsStored bool + // OriginalEvent is the event containing the original function calls to resume. + OriginalEvent *session.Event +} + +// CurrentAuthPreprocessorResult holds the result of the last auth preprocessing. +// This is set by authPreprocessor and read by Flow.runOneStep. +var CurrentAuthPreprocessorResult *AuthPreprocessorResult + 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 +} + +// parseAuthConfigFromMap converts any map-like auth_config payload into auth.AuthConfig. +func parseAuthConfigFromMap(data any) (*auth.AuthConfig, error) { + var config auth.AuthConfig + if err := decodeSnakeCompatibleMap(data, &config, "auth_config"); err != nil { + return nil, err + } + return &config, nil +} + +// parseAuthCredentialFromMap converts any map-like auth credential payload into auth.AuthCredential. +func parseAuthCredentialFromMap(data any) (*auth.AuthCredential, error) { + var cred auth.AuthCredential + if err := decodeSnakeCompatibleMap(data, &cred, "credential"); err != nil { + return nil, err + } + return &cred, nil +} + +func decodeSnakeCompatibleMap(data any, target any, kind string) error { + dataMap, ok := data.(map[string]any) + if !ok { + return fmt.Errorf("%s is not a map", kind) + } + decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{ + TagName: "json", + Result: target, + WeaklyTypedInput: true, + MatchName: func(mapKey, fieldName string) bool { + return canonicalFieldName(mapKey) == canonicalFieldName(fieldName) + }, + }) + if err != nil { + return fmt.Errorf("failed to build decoder: %w", err) + } + if err := decoder.Decode(dataMap); err != nil { + return fmt.Errorf("failed to decode %s: %w", kind, err) + } return nil } +func canonicalFieldName(s string) string { + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + if r == '_' || r == '-' { + continue + } + b.WriteRune(unicode.ToLower(r)) + } + return b.String() +} + func nlPlanningResponseProcessor(ctx agent.InvocationContext, req *model.LLMRequest, resp *model.LLMResponse) error { // TODO: implement (adk-python src/google/adk/_nl_planning.py) return nil diff --git a/internal/toolinternal/context.go b/internal/toolinternal/context.go index 9a3525d84..e382a34df 100644 --- a/internal/toolinternal/context.go +++ b/internal/toolinternal/context.go @@ -22,6 +22,7 @@ import ( "google.golang.org/adk/agent" "google.golang.org/adk/artifact" + "google.golang.org/adk/auth" contextinternal "google.golang.org/adk/internal/context" "google.golang.org/adk/memory" "google.golang.org/adk/session" @@ -53,11 +54,17 @@ func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions functionCallID = uuid.NewString() } if actions == nil { - actions = &session.EventActions{StateDelta: make(map[string]any)} + actions = &session.EventActions{ + StateDelta: make(map[string]any), + RequestedAuthConfigs: make(map[string]*auth.AuthConfig), + } } if actions.StateDelta == nil { actions.StateDelta = make(map[string]any) } + if actions.RequestedAuthConfigs == nil { + actions.RequestedAuthConfigs = make(map[string]*auth.AuthConfig) + } cbCtx := contextinternal.NewCallbackContextWithDelta(ctx, actions.StateDelta) return &toolContext{ @@ -99,3 +106,57 @@ func (c *toolContext) AgentName() string { func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) { return c.invocationContext.Memory().Search(ctx, query) } + +// RequestCredential requests user authorization for OAuth2. +// The auth config will be included in the event's RequestedAuthConfigs, +// which is converted to adk_request_credential function calls by GenerateAuthEvent. +func (c *toolContext) RequestCredential(config *auth.AuthConfig) { + + if config == nil { + return + } + + // Generate auth request with auth_uri + handler := auth.NewAuthHandler(config) + authRequest, err := handler.GenerateAuthRequest() + if err != nil { + // TODO: log or surface the error once a logging strategy is defined. + return + } + if authRequest == nil { + return + } + + // Add to RequestedAuthConfigs keyed by function call ID + c.eventActions.RequestedAuthConfigs[c.functionCallID] = authRequest +} + +// GetAuthResponse retrieves the auth response from session state. +// Returns nil if no auth response is available. +func (c *toolContext) GetAuthResponse(config *auth.AuthConfig) *auth.AuthCredential { + if config == nil { + return nil + } + key := session.KeyPrefixTemp + config.CredentialKey + + val, err := c.invocationContext.Session().State().Get(key) + if err != nil { + return nil + } + if val == nil { + return nil + } + + if cred, ok := val.(*auth.AuthCredential); ok { + return cred + } + + return nil +} + +// CredentialService returns the credential service for persistent storage. +// Returns nil as toolContext does not have a default credential service. +// The InvocationContext or runner should provide a credential service if needed. +func (c *toolContext) CredentialService() auth.CredentialService { + return nil +} diff --git a/internal/toolinternal/context_test.go b/internal/toolinternal/context_test.go index 517cf6260..b8a2c2b0e 100644 --- a/internal/toolinternal/context_test.go +++ b/internal/toolinternal/context_test.go @@ -18,6 +18,7 @@ import ( "testing" "google.golang.org/adk/agent" + "google.golang.org/adk/auth" contextinternal "google.golang.org/adk/internal/context" "google.golang.org/adk/session" ) @@ -36,3 +37,85 @@ func TestToolContext(t *testing.T) { t.Errorf("ToolContext(%+T) is unexpectedly an InvocationContext", got) } } + +func TestToolContext_RequestCredential(t *testing.T) { + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) + actions := &session.EventActions{ + RequestedAuthConfigs: make(map[string]*auth.AuthConfig), + } + toolCtx := NewToolContext(inv, "fn-123", actions) + + authConfig := &auth.AuthConfig{ + AuthScheme: &auth.OAuth2Scheme{ + Flows: &auth.OAuthFlows{ + AuthorizationCode: &auth.OAuthFlowAuthorizationCode{ + AuthorizationURL: "https://example.com/auth", + TokenURL: "https://example.com/token", + }, + }, + }, + RawAuthCredential: &auth.AuthCredential{ + AuthType: auth.AuthCredentialTypeOAuth2, + OAuth2: &auth.OAuth2Auth{ + ClientID: "client-id", + ClientSecret: "client-secret", + RedirectURI: "https://localhost/callback", + }, + }, + CredentialKey: "test-key", + } + + // Request credential + tc := toolCtx.(*toolContext) + tc.RequestCredential(authConfig) + + // Verify it was added to RequestedAuthConfigs + if len(actions.RequestedAuthConfigs) != 1 { + t.Errorf("RequestedAuthConfigs has %d entries, want 1", len(actions.RequestedAuthConfigs)) + } + + stored, ok := actions.RequestedAuthConfigs["fn-123"] + if !ok { + t.Error("RequestedAuthConfigs should have entry for 'fn-123'") + } + if stored == nil { + t.Error("Stored auth config should not be nil") + } + // AuthHandler.GenerateAuthRequest adds auth_uri + if stored.ExchangedAuthCredential == nil || stored.ExchangedAuthCredential.OAuth2 == nil { + t.Error("ExchangedAuthCredential should be set") + } +} + +func TestToolContext_RequestCredential_Nil(t *testing.T) { + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) + actions := &session.EventActions{ + RequestedAuthConfigs: make(map[string]*auth.AuthConfig), + } + toolCtx := NewToolContext(inv, "fn-123", actions) + + tc := toolCtx.(*toolContext) + tc.RequestCredential(nil) + + // Should not panic and should not add anything + if len(actions.RequestedAuthConfigs) != 0 { + t.Errorf("RequestedAuthConfigs has %d entries, want 0", len(actions.RequestedAuthConfigs)) + } +} + +// Note: TestToolContext_GetAuthResponse tests are skipped as they require +// a full invocation context with session state. The GetAuthResponse function +// works correctly when session state is available. + +func TestToolContext_CredentialService(t *testing.T) { + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) + toolCtx := NewToolContext(inv, "fn-123", &session.EventActions{}) + + tc := toolCtx.(*toolContext) + svc := tc.CredentialService() + + // ToolContext doesn't have a credential service by default + if svc != nil { + t.Errorf("CredentialService() = %v, want nil", svc) + } +} diff --git a/session/database/service_test.go b/session/database/service_test.go index 3e7290273..b50e56af7 100644 --- a/session/database/service_test.go +++ b/session/database/service_test.go @@ -16,6 +16,7 @@ package database import ( "maps" + "os" "strconv" "testing" "time" @@ -30,6 +31,14 @@ import ( "google.golang.org/adk/session" ) +// TestMain sets up the test environment. +// We set time.Local to UTC to ensure consistent timestamp formatting in tests, +// since the database stores timestamps without timezone info and tests expect UTC format. +func TestMain(m *testing.M) { + time.Local = time.UTC + os.Exit(m.Run()) +} + func Test_databaseService_Create(t *testing.T) { tests := []struct { name string diff --git a/session/session.go b/session/session.go index 581d3e2ba..9a23156e5 100644 --- a/session/session.go +++ b/session/session.go @@ -21,6 +21,7 @@ import ( "github.com/google/uuid" + "google.golang.org/adk/auth" "google.golang.org/adk/model" ) @@ -134,7 +135,10 @@ func NewEvent(invocationID string) *Event { ID: uuid.NewString(), InvocationID: invocationID, Timestamp: time.Now(), - Actions: EventActions{StateDelta: make(map[string]any)}, + Actions: EventActions{ + StateDelta: make(map[string]any), + RequestedAuthConfigs: make(map[string]*auth.AuthConfig), + }, } } @@ -154,6 +158,11 @@ type EventActions struct { TransferToAgent string // The agent is escalating to a higher level agent. Escalate bool + + // RequestedAuthConfigs holds authentication configurations requested by tools. + // Key is the function call ID, value is the auth config. + // These are converted to adk_request_credential function calls by GenerateAuthEvent. + RequestedAuthConfigs map[string]*auth.AuthConfig } // Prefixes for defining session's state scopes diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 87f14b7c7..ab80099b5 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -159,7 +159,7 @@ func TestMCPToolSet(t *testing.T) { if diff := cmp.Diff(wantEvents, gotEvents, cmpopts.IgnoreFields(session.Event{}, "ID", "Timestamp", "InvocationID"), - cmpopts.IgnoreFields(session.EventActions{}, "StateDelta"), + cmpopts.IgnoreFields(session.EventActions{}, "StateDelta", "RequestedAuthConfigs"), cmpopts.IgnoreFields(model.LLMResponse{}, "UsageMetadata", "AvgLogprobs", "FinishReason"), cmpopts.IgnoreFields(genai.FunctionCall{}, "ID"), cmpopts.IgnoreFields(genai.FunctionResponse{}, "ID"), diff --git a/tool/tool.go b/tool/tool.go index b5107d7c6..9d2dcb0a1 100644 --- a/tool/tool.go +++ b/tool/tool.go @@ -21,6 +21,7 @@ import ( "context" "google.golang.org/adk/agent" + "google.golang.org/adk/auth" "google.golang.org/adk/memory" "google.golang.org/adk/session" ) @@ -51,6 +52,19 @@ type Context interface { Actions() *session.EventActions // SearchMemory performs a semantic search on the agent's memory. SearchMemory(context.Context, string) (*memory.SearchResponse, error) + + // RequestCredential requests user authorization for OAuth2. + // The auth config will be included in the event's RequestedAuthConfigs. + // The runner will send an adk_request_credential event to the client. + RequestCredential(config *auth.AuthConfig) + + // GetAuthResponse retrieves the auth response from session state. + // Returns nil if no auth response is available. + GetAuthResponse(config *auth.AuthConfig) *auth.AuthCredential + + // CredentialService returns the credential service for persistent storage. + // Returns nil if no credential service is configured. + CredentialService() auth.CredentialService } // Toolset is an interface for a collection of tools. It allows grouping