Merge pull request #615 from dgageot/refactor-jwt

Make Docker Desktop JWT available through the env provider

Author: dgageot

pkg/desktop/login.go

@@ -2,7 +2,6 @@ package desktop
 
 import (
 	"context"
-	"os"
 )
 
 type DockerHubInfo struct {
@@ -12,13 +11,6 @@ type DockerHubInfo struct {
 }
 
 func GetToken(ctx context.Context) string {
-	// Allow the user to override the token via an environment variable.
-	// This is e.g. useful when talking to a gateway on staging.
-	manualToken := os.Getenv("DOCKER_TOKEN")
-	if manualToken != "" {
-		return manualToken
-	}
-
 	var token string
 	_ = ClientBackend.Get(ctx, "/registry/token", &token)
 	return token

pkg/environment/default.go

@@ -15,5 +15,8 @@ func NewDefaultProvider() Provider {
 		providers = append(providers, keychainProvider)
 	}
 
+	// Append Docker Desktop provider last
+	providers = append(providers, NewDockerDesktopProvider())
+
 	return NewMultiProvider(providers...)
 }

pkg/environment/docker-desktop.go

@@ -0,0 +1,23 @@
+package environment
+
+import (
+	"context"
+
+	"github.com/docker/cagent/pkg/desktop"
+)
+
+const DockerDesktopTokenEnv = "DOCKER_TOKEN"
+
+type DockerDesktopProvider struct{}
+
+func NewDockerDesktopProvider() *DockerDesktopProvider {
+	return &DockerDesktopProvider{}
+}
+
+func (p *DockerDesktopProvider) Get(ctx context.Context, name string) string {
+	if name != DockerDesktopTokenEnv {
+		return ""
+	}
+
+	return desktop.GetToken(ctx)
+}

pkg/model/provider/anthropic/client.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/docker/cagent/pkg/chat"
 	latest "github.com/docker/cagent/pkg/config/v2"
-	"github.com/docker/cagent/pkg/desktop"
 	"github.com/docker/cagent/pkg/environment"
 	"github.com/docker/cagent/pkg/model/provider/base"
 	"github.com/docker/cagent/pkg/model/provider/options"
@@ -24,11 +23,7 @@ import (
 // It holds the anthropic client and model config
 type Client struct {
 	base.Config
-	client anthropic.Client
-	// When using the Docker AI Gateway, tokens are short-lived. We rebuild
-	// the client per request when in gateway mode.
-	useGateway     bool
-	gatewayBaseURL string
+	clientFn func(context.Context) (anthropic.Client, error)
 }
 
 // interleavedThinkingEnabled returns false unless explicitly enabled via
@@ -76,40 +71,47 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 		opt(&globalOptions)
 	}
 
-	var requestOptions []option.RequestOption
-	useGateway := false
-	gatewayBaseURL := ""
+	var clientFn func(context.Context) (anthropic.Client, error)
 	if gateway := globalOptions.Gateway(); gateway == "" {
 		authToken := env.Get(ctx, "ANTHROPIC_API_KEY")
 		if authToken == "" {
 			return nil, errors.New("ANTHROPIC_API_KEY environment variable is required")
 		}
 
 		slog.Debug("Anthropic API key found, creating client")
-		requestOptions = append(requestOptions,
+		requestOptions := []option.RequestOption{
 			option.WithAPIKey(authToken),
-		)
+		}
 		if cfg.BaseURL != "" {
 			requestOptions = append(requestOptions, option.WithBaseURL(cfg.BaseURL))
 		}
+		client := anthropic.NewClient(requestOptions...)
+		clientFn = func(context.Context) (anthropic.Client, error) {
+			return client, nil
+		}
 	} else {
-		authToken := desktop.GetToken(ctx)
-		if authToken == "" {
+		// Fail fast if Docker Desktop's auth token isn't available
+		if env.Get(ctx, environment.DockerDesktopTokenEnv) == "" {
 			slog.Error("Anthropic client creation failed", "error", "failed to get Docker Desktop's authentication token")
 			return nil, errors.New("sorry, you first need to sign in Docker Desktop to use the Docker AI Gateway")
 		}
 
-		slog.Debug("Docker Desktop's authentication token found, creating client")
-		requestOptions = append(requestOptions,
-			option.WithAuthToken(authToken),
-			option.WithAPIKey(authToken),
-			option.WithBaseURL(gateway),
-		)
-		useGateway = true
-		gatewayBaseURL = gateway
+		// When using a Gateway, tokens are short-lived.
+		clientFn = func(ctx context.Context) (anthropic.Client, error) {
+			// Query a fresh auth token each time the client is used
+			authToken := env.Get(ctx, environment.DockerDesktopTokenEnv)
+			if authToken == "" {
+				return anthropic.Client{}, errors.New("failed to get Docker Desktop token for Gateway")
+			}
+
+			return anthropic.NewClient(
+				option.WithAuthToken(authToken),
+				option.WithAPIKey(authToken),
+				option.WithBaseURL(gateway),
+			), nil
+		}
 	}
 
-	client := anthropic.NewClient(requestOptions...)
 	slog.Debug("Anthropic client created successfully", "model", cfg.Model)
 
 	if globalOptions.StructuredOutput != nil {
@@ -121,25 +123,10 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 			ModelConfig:  cfg,
 			ModelOptions: globalOptions,
 		},
-		client:         client,
-		useGateway:     useGateway,
-		gatewayBaseURL: gatewayBaseURL,
+		clientFn: clientFn,
 	}, nil
 }
 
-// newGatewayClient builds a new Anthropic client using a fresh Docker Desktop token.
-func (c *Client) newGatewayClient(ctx context.Context) anthropic.Client {
-	authToken := desktop.GetToken(ctx)
-	opts := []option.RequestOption{
-		option.WithAuthToken(authToken),
-		option.WithAPIKey(authToken),
-	}
-	if c.gatewayBaseURL != "" {
-		opts = append(opts, option.WithBaseURL(c.gatewayBaseURL))
-	}
-	return anthropic.NewClient(opts...)
-}
-
 // CreateChatCompletionStream creates a streaming chat completion request
 func (c *Client) CreateChatCompletionStream(
 	ctx context.Context,
@@ -155,10 +142,11 @@ func (c *Client) CreateChatCompletionStream(
 	if maxTokens == 0 {
 		maxTokens = 8192 // Default output budget when not specified
 	}
-	// Build a fresh client per request when using the gateway
-	client := c.client
-	if c.useGateway {
-		client = c.newGatewayClient(ctx)
+
+	client, err := c.clientFn(ctx)
+	if err != nil {
+		slog.Error("Failed to create Anthropic client", "error", err)
+		return nil, err
 	}
 
 	// Use Beta API with interleaved thinking only when enabled

pkg/model/provider/gemini/client.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/docker/cagent/pkg/chat"
 	latest "github.com/docker/cagent/pkg/config/v2"
-	"github.com/docker/cagent/pkg/desktop"
 	"github.com/docker/cagent/pkg/environment"
 	"github.com/docker/cagent/pkg/model/provider/base"
 	"github.com/docker/cagent/pkg/model/provider/options"
@@ -24,11 +23,7 @@ import (
 // It implements the provider.Provider interface
 type Client struct {
 	base.Config
-	client *genai.Client
-	// When using the Docker AI Gateway, tokens are short-lived. We rebuild
-	// the client per request when in gateway mode.
-	useGateway     bool
-	gatewayBaseURL string
+	clientFn func(context.Context) (*genai.Client, error)
 }
 
 // NewClient creates a new Gemini client from the provided configuration
@@ -46,66 +41,63 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 		opt(&modelOptions)
 	}
 
-	var apiKey string
-	var httpOptions genai.HTTPOptions
-	useGateway := false
-	gatewayBaseURL := ""
+	var clientFn func(context.Context) (*genai.Client, error)
 	if gateway := modelOptions.Gateway(); gateway == "" {
-		apiKey = env.Get(ctx, "GOOGLE_API_KEY")
+		apiKey := env.Get(ctx, "GOOGLE_API_KEY")
 		if apiKey == "" {
 			return nil, errors.New("GOOGLE_API_KEY environment variable is required")
 		}
+
+		client, err := genai.NewClient(ctx, &genai.ClientConfig{
+			APIKey:  apiKey,
+			Backend: genai.BackendGeminiAPI,
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		clientFn = func(context.Context) (*genai.Client, error) {
+			return client, nil
+		}
 	} else {
-		// genai client requires a non-empty API key
-		apiKey = desktop.GetToken(ctx)
-		if apiKey == "" {
+		// Fail fast if Docker Desktop's auth token isn't available
+		if env.Get(ctx, environment.DockerDesktopTokenEnv) == "" {
+			slog.Error("Gemini client creation failed", "error", "failed to get Docker Desktop's authentication token")
 			return nil, errors.New("sorry, you first need to sign in Docker Desktop to use the Docker AI Gateway")
 		}
-		httpOptions.BaseURL = gateway
-		httpOptions.Headers = make(http.Header)
-		httpOptions.Headers.Set("Authorization", "Bearer "+apiKey)
-		useGateway = true
-		gatewayBaseURL = gateway
-	}
 
-	client, err := genai.NewClient(ctx, &genai.ClientConfig{
-		APIKey:      apiKey,
-		Backend:     genai.BackendGeminiAPI,
-		HTTPOptions: httpOptions,
-	})
-	if err != nil {
-		return nil, err
+		// When using a Gateway, tokens are short-lived.
+		clientFn = func(ctx context.Context) (*genai.Client, error) {
+			// Query a fresh auth token each time the client is used
+			authToken := env.Get(ctx, environment.DockerDesktopTokenEnv)
+			if authToken == "" {
+				return nil, errors.New("failed to get Docker Desktop token for Gateway")
+			}
+
+			return genai.NewClient(ctx, &genai.ClientConfig{
+				APIKey:  authToken,
+				Backend: genai.BackendGeminiAPI,
+				HTTPOptions: genai.HTTPOptions{
+					BaseURL: gateway,
+					Headers: http.Header{
+						"Authorization": []string{"Bearer " + authToken},
+					},
+				},
+			})
+		}
 	}
 
+	slog.Debug("Gemini client created successfully", "model", cfg.Model)
+
 	return &Client{
 		Config: base.Config{
 			ModelConfig:  cfg,
 			ModelOptions: modelOptions,
 		},
-		client:         client,
-		useGateway:     useGateway,
-		gatewayBaseURL: gatewayBaseURL,
+		clientFn: clientFn,
 	}, nil
 }
 
-// newGatewayClient builds a new Gemini client using a fresh Docker Desktop token.
-func (c *Client) newGatewayClient(ctx context.Context) (*genai.Client, error) {
-	token := desktop.GetToken(ctx)
-	if token == "" {
-		return nil, errors.New("failed to get Docker Desktop token for gateway")
-	}
-	httpOptions := genai.HTTPOptions{
-		BaseURL: c.gatewayBaseURL,
-		Headers: make(http.Header),
-	}
-	httpOptions.Headers.Set("Authorization", "Bearer "+token)
-	return genai.NewClient(ctx, &genai.ClientConfig{
-		APIKey:      token,
-		Backend:     genai.BackendGeminiAPI,
-		HTTPOptions: httpOptions,
-	})
-}
-
 // convertMessagesToGemini converts chat.Messages into Gemini Contents
 func convertMessagesToGemini(messages []chat.Message) []*genai.Content {
 	contents := make([]*genai.Content, 0, len(messages))
@@ -338,16 +330,13 @@ func (c *Client) CreateChatCompletionStream(
 		slog.Debug("Message", "index", i, "role", content.Role)
 	}
 
-	// Build a fresh client per request when using the gateway
-	var iter func(func(*genai.GenerateContentResponse, error) bool)
-	if c.useGateway {
-		if gwClient, err := c.newGatewayClient(ctx); err == nil {
-			iter = gwClient.Models.GenerateContentStream(ctx, c.ModelConfig.Model, contents, config)
-		} else {
-			iter = c.client.Models.GenerateContentStream(ctx, c.ModelConfig.Model, contents, config)
-		}
-	} else {
-		iter = c.client.Models.GenerateContentStream(ctx, c.ModelConfig.Model, contents, config)
+	client, err := c.clientFn(ctx)
+	if err != nil {
+		slog.Error("Failed to create Gemini client", "error", err)
+		return nil, err
 	}
+
+	// Build a fresh client per request when using the gateway
+	iter := client.Models.GenerateContentStream(ctx, c.ModelConfig.Model, contents, config)
 	return NewStreamAdapter(iter, c.ModelConfig.Model), nil
 }