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
+208 -1
View File
@@ -526,7 +526,7 @@ workflows/
# LLM Configuration # LLM Configuration
LLM_ENDPOINT="https://api.riotpiao.com/v1/chat/completions" LLM_ENDPOINT="https://api.riotpiao.com/v1/chat/completions"
LLM_MODEL="reasoning" LLM_MODEL="reasoning"
LLM_API_KEY="your-api-key" LOCAL_LLM_BASE_URL="https://api.riotpiao.com" # Override for local dev
# Temporal Configuration # Temporal Configuration
TEMPORAL_HOST_URL="localhost:7233" TEMPORAL_HOST_URL="localhost:7233"
@@ -539,6 +539,122 @@ MEMORY_SERVICE_URL="http://poimen-memory.poimen.svc.cluster.local:8080"
LOG_LEVEL="info" LOG_LEVEL="info"
``` ```
### Authentication & Federation
Poimen supports multiple authentication methods to federate LLM access across customers and tenants:
#### 1. Bearer Token (JWT/OAuth2)
```go
auth := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: "eyJhbGciOiJIUzI1NiIs...", // JWT token from your auth provider
TenantID: "customer-123", // Optional: tenant isolation
Scopes: "llm:read llm:write", // Optional: OAuth2 scopes
}
client := routing.NewLLMClientWithAuth(auth)
```
#### 2. API Key
```go
auth := &routing.LLMAuth{
Type: routing.AuthTypeAPIKey,
APIKey: "sk-xxx-yyy-zzz", // API key from provider
TenantID: "customer-456",
}
client := routing.NewLLMClientWithAuth(auth)
```
#### 3. Custom Header
```go
auth := &routing.LLMAuth{
Type: routing.AuthTypeCustom,
HeaderName: "X-Custom-Auth",
HeaderValue: "custom-token-value",
TenantID: "customer-789",
}
client := routing.NewLLMClientWithAuth(auth)
```
#### 4. Router Configuration with Auth
```go
config := &routing.LLMRouterConfig{
Provider: openaiProvider,
KnowledgeBase: kb,
Auth: &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: jwtToken,
TenantID: customerID, // Tenant isolation in multi-tenant deployments
},
TenantID: customerID, // Additional tenant tracking
}
router := routing.NewLLMRouter(config)
```
#### Multi-Tenant Federated Access
For multi-tenant deployments:
```go
// Per-customer isolated routers
func CreateCustomerRouter(customerID, jwtToken string) (*routing.LLMRouter, error) {
auth := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: jwtToken,
TenantID: customerID, // Passed to LLM API in X-Tenant-ID header
Scopes: "llm:read", // Restrict scopes per customer
}
config := &routing.LLMRouterConfig{
Provider: &routing.LLMClient{},
KnowledgeBase: globalKB,
Auth: auth,
TenantID: customerID,
}
return routing.NewLLMRouter(config)
}
// Usage in activity
input := &routing.LLMRouterInput{
Message: "Analyze code",
MemoryContext: memCtx,
}
output, err := router.Route(ctx, input)
// Headers sent to LLM API:
// Authorization: Bearer <jwtToken>
// X-Tenant-ID: <customerID>
// X-OAuth-Scopes: llm:read
```
#### Token Refresh & Rotation
For long-running workflows, update auth at runtime:
```go
client := routing.NewLLMClientWithAuth(oldAuth)
// Token expires, refresh it
newAuth := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: refreshedToken, // New token from OAuth2 provider
TenantID: customerID,
}
client.UpdateAuth(newAuth)
// Subsequent requests use new token
result, err := client.Chat(ctx, systemPrompt, userMsg)
```
#### Headers Sent to LLM API
| Header | Set When | Value | Purpose |
|--------|----------|-------|---------|
| `Authorization` | Bearer auth | `Bearer {token}` | OAuth2/JWT authentication |
| `X-API-Key` | API Key auth | `{api-key}` | API key authentication |
| `X-Tenant-ID` | Any auth type | `{tenantID}` | Tenant/customer isolation |
| `X-OAuth-Scopes` | Bearer with scopes | Space-separated scopes | OAuth2 scope enforcement |
| Custom header | Custom auth | `{headerValue}` | Custom authentication scheme |
### LLMRouter Configuration (Programmatic) ### LLMRouter Configuration (Programmatic)
```go ```go
config := &routing.LLMRouterConfig{ config := &routing.LLMRouterConfig{
@@ -664,6 +780,97 @@ kubectl apply -f k8s/worker-deployment.yaml
--- ---
## Security & Multi-Tenancy
### Token Management
**Never hardcode tokens!** Use secure secret management:
```go
// ❌ DON'T DO THIS
auth := &routing.LLMAuth{
Token: "eyJhbGciOiJIUzI1NiIs...", // Hardcoded!
}
// ✅ DO THIS
tokenFromVault, _ := vaultClient.GetSecret("llm-token-" + customerID)
auth := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: tokenFromVault,
TenantID: customerID,
}
```
**Recommended secret management:**
- Kubernetes Secrets (development)
- HashiCorp Vault (production)
- AWS Secrets Manager / GCP Secret Manager (cloud)
- Sealed Secrets / Sealed Policies
### Tenant Isolation
The `TenantID` header enforces tenant isolation on the LLM API side:
```go
// Customer A's workflow
authA := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: tokenA,
TenantID: "customer-A", // ← Isolates this customer
}
// Customer B's workflow
authB := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: tokenB,
TenantID: "customer-B", // ← Isolates this customer
}
```
The LLM API server should:
- Validate tenant ownership of tokens
- Enforce data boundaries per tenant
- Log access per tenant ID
- Rate-limit per tenant
### Scope-Based Access Control
Use OAuth2 scopes to limit capabilities:
```go
// Analytics-only customer
analytics := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: token,
Scopes: "llm:read", // Read-only
}
// Full-access customer
admin := &routing.LLMAuth{
Type: routing.AuthTypeBearer,
Token: token,
Scopes: "llm:read llm:write llm:admin", // Full access
}
```
The LLM API server should validate scopes before executing requests.
### Audit & Compliance
All requests include identifying headers for audit trails:
```
POST /v1/chat/completions
Authorization: Bearer eyJhbGc...
X-Tenant-ID: customer-A
X-OAuth-Scopes: llm:read
# Server logs:
# timestamp=2024-01-15T10:30:00Z tenant=customer-A scope=llm:read status=200 tokens=1500
```
---
## Documentation ## Documentation
- **[Routing Workflow Spec](./docs/ROUTING_WORKFLOW_SPEC.md)** — Complete spec format reference - **[Routing Workflow Spec](./docs/ROUTING_WORKFLOW_SPEC.md)** — Complete spec format reference
+110 -1
View File
@@ -8,6 +8,7 @@ import (
"io" "io"
"net/http" "net/http"
"os" "os"
"strings"
) )
var ( 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 // LLMClient is a simple LLM client for routing
type LLMClient struct { type LLMClient struct {
baseURL string baseURL string
httpClient *http.Client httpClient *http.Client
auth *LLMAuth
} }
// NewLLMClient creates a new LLM client // NewLLMClient creates a new LLM client with default (no) auth
func NewLLMClient() *LLMClient { func NewLLMClient() *LLMClient {
return &LLMClient{ return &LLMClient{
baseURL: llmBaseURL, baseURL: llmBaseURL,
httpClient: &http.Client{}, 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,
} }
} }
@@ -107,6 +157,11 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
httpReq.Header.Set("Content-Type", "application/json") 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) resp, err := c.httpClient.Do(httpReq)
if err != nil { if err != nil {
return "", fmt.Errorf("failed to connect to LLM API at %s: %w", c.baseURL, err) return "", fmt.Errorf("failed to connect to LLM API at %s: %w", c.baseURL, err)
@@ -133,3 +188,57 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
return respObj.Choices[0].Message.Content, nil 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 SpecBuilder SpecBuilder
Validators []WorkflowValidator Validators []WorkflowValidator
ParamBinder ParameterBinder 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 // NewLLMRouter creates a new LLM router with custom config
+2 -2
View File
@@ -9,6 +9,6 @@ metadata:
app.kubernetes.io/name: poimen app.kubernetes.io/name: poimen
app.kubernetes.io/component: orchestrator app.kubernetes.io/component: orchestrator
data: data:
GIT_COMMIT: "f3ce019c" # Updated automatically by CI/CD GIT_COMMIT: "30644a8e" # Updated automatically by CI/CD
GIT_BRANCH: "main" GIT_BRANCH: "main"
DEPLOYMENT_DATE: "2026-09-03" DEPLOYMENT_DATE: "2026-09-04"
+2 -2
View File
@@ -13,8 +13,8 @@ spec:
labels: labels:
app: poimen-worker app: poimen-worker
annotations: annotations:
git-commit: "f3ce019c" # ✅ Updated on each push, triggers rolling restart git-commit: "30644a8e" # ✅ Updated on each push, triggers rolling restart
deployment-date: "2026-09-03" deployment-date: "2026-09-04"
spec: spec:
containers: containers:
- name: worker - name: worker