feat(client): add support for short-lived tokens (#799)

Co-authored-by: stainless-app[bot] <142633134+stainless-app[bot]@users.noreply.github.com>

Author: karpetrosyan

.stats.yml

@@ -1,4 +1,4 @@
 configured_endpoints: 139
-openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/openai%2Fopenai-dd99495ad509338e6de862802839360dfe394d5cd6d6ba6d13fec8fca92328b8.yml
+openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/openai%2Fopenai-a6eca1bd01e0c434af356fe5275c206057216a4e626d1051d294c27016cd6d05.yml
 openapi_spec_hash: 68abda9122013a9ae3f084cfdbe8e8c1
-config_hash: 5635033cdc8c930255f8b529a78de722
+config_hash: 4975e16a94e8f9901428022044131888

README.md

@@ -953,6 +953,104 @@ You may also replace the default `http.Client` with
 accepted (this overwrites any previous client) and receives requests after any
 middleware has been applied.
 
+## Workload Identity Authentication
+
+For cloud workloads (Kubernetes, Azure, Google Cloud Platform), you can use workload identity authentication instead of API keys. This provides short-lived tokens that are automatically refreshed.
+
+### Kubernetes
+
+```go
+import (
+	"github.com/openai/openai-go/v3"
+	"github.com/openai/openai-go/v3/auth"
+	"github.com/openai/openai-go/v3/option"
+)
+
+client := openai.NewClient(
+	option.WithWorkloadIdentity(auth.WorkloadIdentity{
+		ClientID:           "your-client-id",
+		IdentityProviderID: "idp-123",
+		ServiceAccountID:   "sa-456",
+		Provider:           auth.K8sServiceAccountTokenProvider(""),
+	}),
+)
+```
+
+### Azure Managed Identity
+
+```go
+client := openai.NewClient(
+	option.WithWorkloadIdentity(auth.WorkloadIdentity{
+		ClientID:           "your-client-id",
+		IdentityProviderID: "idp-123",
+		ServiceAccountID:   "sa-456",
+		Provider:           auth.AzureManagedIdentityTokenProvider(nil),
+	}),
+)
+```
+
+### Google Cloud Compute Engine
+
+```go
+client := openai.NewClient(
+	option.WithWorkloadIdentity(auth.WorkloadIdentity{
+		ClientID:           "your-client-id",
+		IdentityProviderID: "idp-123",
+		ServiceAccountID:   "sa-456",
+		Provider:           auth.GCPIDTokenProvider(nil),
+	}),
+)
+```
+
+### Custom Subject Token Provider
+
+You can implement your own subject token provider:
+
+```go
+import (
+	"context"
+
+	"github.com/openai/openai-go/v3"
+	"github.com/openai/openai-go/v3/auth"
+	"github.com/openai/openai-go/v3/option"
+)
+
+type customTokenProvider struct{}
+
+func (p *customTokenProvider) TokenType() auth.SubjectTokenType {
+	return auth.SubjectTokenTypeJWT
+}
+
+func (p *customTokenProvider) GetToken(ctx context.Context, httpClient auth.HTTPDoer) (string, error) {
+	return "your-token", nil
+}
+
+client := openai.NewClient(
+	option.WithWorkloadIdentity(auth.WorkloadIdentity{
+		ClientID:           "your-client-id",
+		IdentityProviderID: "idp-123",
+		ServiceAccountID:   "sa-456",
+		Provider:           &customTokenProvider{},
+	}),
+)
+```
+
+### Customizing Refresh Buffer
+
+By default, tokens are refreshed 20 minutes (1200 seconds) before expiry. You can customize this:
+
+```go
+client := openai.NewClient(
+	option.WithWorkloadIdentity(auth.WorkloadIdentity{
+		ClientID:             "your-client-id",
+		IdentityProviderID:   "idp-123",
+		ServiceAccountID:     "sa-456",
+		Provider:             auth.K8sServiceAccountTokenProvider(""),
+		RefreshBufferSeconds: 600,
+	}),
+)
+```
+
 ## Microsoft Azure OpenAI
 
 To use this library with [Azure OpenAI]https://learn.microsoft.com/azure/ai-services/openai/overview),

aliases.go

@@ -448,6 +448,15 @@ type FunctionParameters = shared.FunctionParameters
 // This is an alias to an internal type.
 type Metadata = shared.Metadata
 
+// This is an alias to an internal type.
+type OAuthErrorCode = shared.OAuthErrorCode
+
+// Equals "invalid_grant"
+const OAuthErrorCodeInvalidGrant = shared.OAuthErrorCodeInvalidGrant
+
+// Equals "invalid_subject_token"
+const OAuthErrorCodeInvalidSubjectToken = shared.OAuthErrorCodeInvalidSubjectToken
+
 // **gpt-5 and o-series models only**
 //
 // Configuration options for

api.md

@@ -7,6 +7,7 @@
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#FunctionDefinitionParam">FunctionDefinitionParam</a>
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#FunctionParameters">FunctionParameters</a>
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#Metadata">Metadata</a>
+- <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#OAuthErrorCode">OAuthErrorCode</a>
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#ReasoningParam">ReasoningParam</a>
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#ReasoningEffort">ReasoningEffort</a>
 - <a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared">shared</a>.<a href="https://pkg.go.dev/github.com/openai/openai-go/v3/shared#ResponseFormatJSONObjectParam">ResponseFormatJSONObjectParam</a>

auth/error.go

@@ -0,0 +1,37 @@
+package auth
+
+import (
+	"fmt"
+
+	"github.com/openai/openai-go/v3/shared"
+)
+
+// SubjectTokenProviderError is raised when failing to get the subject token from the cloud environment.
+type SubjectTokenProviderError struct {
+	Provider string
+	Message  string
+	Cause    error
+}
+
+func (e *SubjectTokenProviderError) Error() string {
+	if e.Cause != nil {
+		return fmt.Sprintf("%s provider error: %s: %v", e.Provider, e.Message, e.Cause)
+	}
+	return fmt.Sprintf("%s provider error: %s", e.Provider, e.Message)
+}
+
+func (e *SubjectTokenProviderError) Unwrap() error {
+	return e.Cause
+}
+
+// OAuthError is raised when there is an error in the RFC-defined OAuth flow, such as failing to exchange the token.
+// See https://datatracker.ietf.org/doc/html/rfc8693 for the OAuth 2.0 Token Exchange specification.
+type OAuthError struct {
+	StatusCode       int
+	ErrorCode        shared.OAuthErrorCode
+	ErrorDescription string
+}
+
+func (e *OAuthError) Error() string {
+	return fmt.Sprintf("OAuth error (status %d): %s - %s", e.StatusCode, e.ErrorCode, e.ErrorDescription)
+}

auth/middleware.go

@@ -0,0 +1,50 @@
+package auth
+
+import (
+	"net/http"
+)
+
+func WorkloadIdentityMiddleware(
+	wia *WorkloadIdentityAuth,
+	httpClient HTTPDoer,
+	req *http.Request,
+	next func(*http.Request) (*http.Response, error),
+) (*http.Response, error) {
+	token, err := wia.GetToken(req.Context(), httpClient)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Authorization", "Bearer "+token)
+
+	resp, err := next(req)
+	if err != nil || resp == nil || resp.StatusCode != http.StatusUnauthorized {
+		return resp, err
+	}
+
+	wia.invalidateToken()
+
+	if req.Body != nil && req.GetBody == nil {
+		return resp, nil
+	}
+
+	retryReq := req.Clone(req.Context())
+
+	token, err = wia.GetToken(req.Context(), httpClient)
+	if err != nil {
+		resp.Body.Close()
+		return nil, err
+	}
+	retryReq.Header.Set("Authorization", "Bearer "+token)
+
+	if req.GetBody != nil {
+		retryReq.Body, err = req.GetBody()
+		if err != nil {
+			resp.Body.Close()
+			return nil, err
+		}
+	}
+
+	resp.Body.Close()
+	return next(retryReq)
+}