feat: add JWT/OAuth2 authentication & multi-tenant federation
ci / test (push) Failing after 1m49s

- Add LLMAuth struct with support for Bearer, API Key, and Custom auth types
- Implement applyAuth() to inject auth headers into LLM requests
- Add X-Tenant-ID header for multi-tenant isolation
- Add X-OAuth-Scopes header for OAuth2 scope enforcement
- Add UpdateAuth() for runtime token refresh (long-running workflows)
- Update LLMRouterConfig with Auth and TenantID fields
- Document 4 authentication patterns (Bearer, API Key, Custom, Router config)
- Add security best practices: token vault integration, tenant isolation, scopes
- Add audit headers for compliance & logging
- Create multi-tenant router factory pattern

Auth types supported:
- Bearer: JWT/OAuth2 tokens (most secure for federated access)
- API Key: Static keys (X-API-Key header)
- Custom: Any custom header-based scheme
- None: No authentication

Customers can now pass per-tenant JWT tokens with customized scopes
and isolated LLM API access per tenant/customer.
This commit is contained in:
Test
2026-09-04 10:54:22 -07:00
parent 86ad8e7b5e
commit 66c17e821f
5 changed files with 324 additions and 6 deletions
+110 -1
View File
@@ -8,6 +8,7 @@ import (
"io"
"net/http"
"os"
"strings"
)
var (
@@ -22,17 +23,66 @@ func init() {
}
}
// AuthType specifies the authentication mechanism
type AuthType string
const (
// AuthTypeNone - no authentication
AuthTypeNone AuthType = "none"
// AuthTypeBearer - Bearer token (JWT, OAuth2)
AuthTypeBearer AuthType = "bearer"
// AuthTypeAPIKey - API Key authentication
AuthTypeAPIKey AuthType = "api-key"
// AuthTypeCustom - Custom header-based authentication
AuthTypeCustom AuthType = "custom"
)
// LLMAuth configures authentication for LLM API
type LLMAuth struct {
// Type of authentication
Type AuthType `json:"type"`
// Token is the JWT/OAuth2 token for Bearer auth
Token string `json:"token,omitempty"`
// APIKey is the API key for API Key auth
APIKey string `json:"apiKey,omitempty"`
// HeaderName is the custom header name for Custom auth
HeaderName string `json:"headerName,omitempty"`
// HeaderValue is the custom header value for Custom auth
HeaderValue string `json:"headerValue,omitempty"`
// TenantID is the tenant/customer ID for multi-tenant federated access
TenantID string `json:"tenantId,omitempty"`
// Scopes are the OAuth2 scopes (space-separated)
Scopes string `json:"scopes,omitempty"`
}
// LLMClient is a simple LLM client for routing
type LLMClient struct {
baseURL string
httpClient *http.Client
auth *LLMAuth
}
// NewLLMClient creates a new LLM client
// NewLLMClient creates a new LLM client with default (no) auth
func NewLLMClient() *LLMClient {
return &LLMClient{
baseURL: llmBaseURL,
httpClient: &http.Client{},
auth: nil,
}
}
// NewLLMClientWithAuth creates a new LLM client with authentication
func NewLLMClientWithAuth(auth *LLMAuth) *LLMClient {
return &LLMClient{
baseURL: llmBaseURL,
httpClient: &http.Client{},
auth: auth,
}
}
@@ -106,6 +156,11 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
}
httpReq.Header.Set("Content-Type", "application/json")
// Apply authentication headers
if err := c.applyAuth(httpReq); err != nil {
return "", fmt.Errorf("failed to apply authentication: %w", err)
}
resp, err := c.httpClient.Do(httpReq)
if err != nil {
@@ -133,3 +188,57 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
return respObj.Choices[0].Message.Content, nil
}
// applyAuth applies authentication to the HTTP request based on config
func (c *LLMClient) applyAuth(req *http.Request) error {
if c.auth == nil || c.auth.Type == AuthTypeNone {
return nil
}
switch c.auth.Type {
case AuthTypeBearer:
if c.auth.Token == "" {
return fmt.Errorf("bearer token is required but not provided")
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.auth.Token))
case AuthTypeAPIKey:
if c.auth.APIKey == "" {
return fmt.Errorf("API key is required but not provided")
}
// Common API key header names: X-API-Key, api-key, Authorization
req.Header.Set("X-API-Key", c.auth.APIKey)
case AuthTypeCustom:
if c.auth.HeaderName == "" || c.auth.HeaderValue == "" {
return fmt.Errorf("custom header name and value are required but not provided")
}
req.Header.Set(c.auth.HeaderName, c.auth.HeaderValue)
}
// Add tenant ID if specified (for multi-tenant federated access)
if c.auth.TenantID != "" {
req.Header.Set("X-Tenant-ID", c.auth.TenantID)
}
// Add scopes if specified (for OAuth2 flows)
if c.auth.Scopes != "" {
req.Header.Set("X-OAuth-Scopes", c.auth.Scopes)
}
return nil
}
// UpdateAuth updates the authentication config at runtime
func (c *LLMClient) UpdateAuth(auth *LLMAuth) error {
if auth == nil {
return fmt.Errorf("auth config cannot be nil")
}
c.auth = auth
return nil
}
// GetAuth returns the current authentication config
func (c *LLMClient) GetAuth() *LLMAuth {
return c.auth
}
+2
View File
@@ -68,6 +68,8 @@ type LLMRouterConfig struct {
SpecBuilder SpecBuilder
Validators []WorkflowValidator
ParamBinder ParameterBinder
Auth *LLMAuth // Authentication config for LLM API
TenantID string // Tenant ID for multi-tenant isolation
}
// NewLLMRouter creates a new LLM router with custom config