feat: add AssumeRoleActivity for temporary LLM API token grants
Implements AWS AssumeRole-like pattern for Poimen: - User/service requests temporary access with identity + scope - AssumeRoleActivity exchanges credentials with OAuth2 auth server - Returns JWT token valid for limited time (default: 1hr, max: 24hrs) - Token used in all subsequent LLM API calls to api.riotpiao.com Key features: - Credentials from vault/K8s secrets (never hardcoded) - Scope-based access control (llm:read, llm:read llm:write, llm:admin) - Automatic token expiration tracking - Retry support for transient auth failures (2x, 1.5s backoff) - Configurable auth server endpoint Usage pattern: 1. AssumeRoleActivity(identity, scope) → JWT token 2. LLMRouter uses token in LLMAuth config 3. All activity calls validated against token + scopes 4. Workflow optionally refreshes token before expiry Security: - No credentials in code/logs (env or vault only) - Short-lived tokens (1hr default, 24hr max) - Server-enforced scope validation - Token revocation support Activity registered: #10 (authentication category) Knowledge base updated with full activity spec New file: action/assume_role.go (5.2 KB)
This commit is contained in:
Test
@@ -120,6 +120,7 @@ Poimen embodies three core principles:
|
||||
|
||||
| Activity | Purpose | Timeout | Retry |
|
||||
|----------|---------|---------|-------|
|
||||
| **AssumeRoleActivity** | Request temporary JWT token (like AWS AssumeRole) | 30s | 2x |
|
||||
| **CloneRepo** | Clone git repository | 30s | 3x |
|
||||
| **AnalyzeCode** | Static analysis (SAST) | 120s | 2x |
|
||||
| **SecurityScan** | Dependency & vulnerability scan | 60s | 2x |
|
||||
@@ -820,6 +821,119 @@ The LLM provider (api.riotpiao.com) should:
|
||||
|
||||
---
|
||||
|
||||
## AssumeRoleActivity: Temporary LLM Token Grants
|
||||
|
||||
**Like AWS AssumeRole**, AssumeRoleActivity requests temporary credentials for accessing LLM APIs:
|
||||
|
||||
```go
|
||||
// 1. User requests temporary token
|
||||
assumeRoleInput := &routing.AssumeRoleInput{
|
||||
Identity: "[email protected]", // Who is accessing
|
||||
Scope: "llm:read llm:write", // What permissions
|
||||
DurationSeconds: 1800, // 30 minutes
|
||||
}
|
||||
|
||||
// 2. Activity exchanges with auth server → returns JWT
|
||||
output, err := temporalClient.ExecuteActivity(ctx,
|
||||
routing.AssumeRoleActivity,
|
||||
assumeRoleInput)
|
||||
|
||||
// 3. Extract token from result
|
||||
var tokenOutput *routing.AssumeRoleOutput
|
||||
output.Get(&tokenOutput)
|
||||
|
||||
// 4. Use token in LLM Router
|
||||
auth := &routing.LLMAuth{
|
||||
Type: routing.AuthTypeBearer,
|
||||
Token: tokenOutput.Token, // ← JWT valid for 30 minutes
|
||||
}
|
||||
router := routing.NewLLMRouter(config)
|
||||
```
|
||||
|
||||
### Workflow Pattern: AssumeRole → LLM Router → Activities
|
||||
|
||||
```go
|
||||
// Step 1: Get temporary credentials
|
||||
assumeRoleResult := workflow.ExecuteActivity(ctx, routing.AssumeRoleActivity, &routing.AssumeRoleInput{
|
||||
Identity: workflowInput.UserID,
|
||||
Scope: "llm:read llm:write",
|
||||
DurationSeconds: 1800,
|
||||
})
|
||||
var token *routing.AssumeRoleOutput
|
||||
assumeRoleResult.Get(&token)
|
||||
|
||||
// Step 2: Use token for all LLM router calls
|
||||
routerInput := &routing.LLMRouterInput{
|
||||
Message: "Analyze code for security",
|
||||
Context: map[string]interface{}{"repo": "myrepo"},
|
||||
}
|
||||
|
||||
routerOutput := workflow.ExecuteActivity(ctx, routing.LLMRouterActivity, routerInput)
|
||||
// LLMRouter automatically uses the token from LLMRouterConfig
|
||||
|
||||
// Step 3: Execute generated workflow with same token
|
||||
// (token baked into all activity calls)
|
||||
```
|
||||
|
||||
### Configuration: Credentials from Vault
|
||||
|
||||
Never hardcode credentials. Use Kubernetes Secrets or Hashicorp Vault:
|
||||
|
||||
```bash
|
||||
# In K8s secret
|
||||
kubectl create secret generic llm-oauth-creds \
|
||||
--from-literal=OAUTH_CLIENT_ID="client-xxx" \
|
||||
--from-literal=OAUTH_CLIENT_SECRET="secret-yyy" \
|
||||
--from-literal=AUTH_SERVER_URL="https://auth.company.com"
|
||||
|
||||
# Pod reads from secret
|
||||
env:
|
||||
- name: OAUTH_CLIENT_ID
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: llm-oauth-creds
|
||||
key: OAUTH_CLIENT_ID
|
||||
```
|
||||
|
||||
In code, AssumeRoleActivity reads from environment:
|
||||
```go
|
||||
input := &routing.AssumeRoleInput{
|
||||
Identity: "[email protected]",
|
||||
Scope: "llm:read",
|
||||
// clientId, clientSecret, authServerUrl read from env automatically
|
||||
}
|
||||
result, _ := AssumeRoleActivity(ctx, input)
|
||||
```
|
||||
|
||||
### Token Lifecycle
|
||||
|
||||
| Stage | Duration | Action |
|
||||
|-------|----------|--------|
|
||||
| **Request** | T+0s | User calls AssumeRoleActivity with identity + scope |
|
||||
| **Grant** | T+1s | Auth server validates, issues JWT (default: 1hr validity) |
|
||||
| **Use** | T+1s to T+3600s | LLMRouter uses token for all api.riotpiao.com calls |
|
||||
| **Refresh** | Before expiry | If workflow > 1hr, request new token via AssumeRole again |
|
||||
| **Revoke** | On demand | Auth server can immediately revoke token if needed |
|
||||
|
||||
### Scopes & Access Control
|
||||
|
||||
Scopes define granular permissions:
|
||||
|
||||
```go
|
||||
// Read-only access (safe for analytics)
|
||||
asScope: "llm:read"
|
||||
|
||||
// Full access (for agent workflows)
|
||||
scope: "llm:read llm:write"
|
||||
|
||||
// Admin access (for operator/setup)
|
||||
scope: "llm:admin"
|
||||
```
|
||||
|
||||
The LLM API validates scopes on every request. AssumeRoleActivity can't escalate privileges—scopes returned by auth server are trusted.
|
||||
|
||||
---
|
||||
|
||||
## Documentation
|
||||
|
||||
- **[Routing Workflow Spec](./docs/ROUTING_WORKFLOW_SPEC.md)** — Complete spec format reference
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
package action
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AssumeRoleInput is the input to AssumeRoleActivity
|
||||
type AssumeRoleInput struct {
|
||||
// Identity is the user/service identity requesting access
|
||||
Identity string `json:"identity"`
|
||||
|
||||
// ClientID is the OAuth2/OIDC client ID (from vault or env)
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
|
||||
// ClientSecret is the OAuth2/OIDC client secret (from vault or env)
|
||||
ClientSecret string `json:"clientSecret,omitempty"`
|
||||
|
||||
// Scope defines what APIs this token can access (e.g., "llm:read llm:write")
|
||||
Scope string `json:"scope"`
|
||||
|
||||
// DurationSeconds is how long the token is valid (default: 3600 = 1 hour)
|
||||
DurationSeconds int `json:"durationSeconds,omitempty"`
|
||||
|
||||
// AuthServerURL is the auth server endpoint (from env if not provided)
|
||||
AuthServerURL string `json:"authServerUrl,omitempty"`
|
||||
}
|
||||
|
||||
// AssumeRoleOutput is the output from AssumeRoleActivity
|
||||
type AssumeRoleOutput struct {
|
||||
// Token is the JWT token for calling api.riotpiao.com
|
||||
Token string `json:"token"`
|
||||
|
||||
// ExpiresAt is when the token expires (Unix timestamp)
|
||||
ExpiresAt int64 `json:"expiresAt"`
|
||||
|
||||
// ExpiresIn is the duration in seconds until expiration
|
||||
ExpiresIn int `json:"expiresIn"`
|
||||
|
||||
// TokenType is typically "Bearer"
|
||||
TokenType string `json:"tokenType"`
|
||||
}
|
||||
|
||||
// oauthTokenRequest is sent to the auth server
|
||||
type oauthTokenRequest struct {
|
||||
GrantType string `json:"grant_type"`
|
||||
ClientID string `json:"client_id"`
|
||||
ClientSecret string `json:"client_secret"`
|
||||
Scope string `json:"scope"`
|
||||
Subject string `json:"subject,omitempty"` // The identity being assumed
|
||||
}
|
||||
|
||||
// oauthTokenResponse is returned from the auth server
|
||||
type oauthTokenResponse struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
TokenType string `json:"token_type"`
|
||||
ExpiresIn int `json:"expires_in"`
|
||||
Scope string `json:"scope"`
|
||||
}
|
||||
|
||||
// AssumeRoleActivity requests a temporary JWT token for accessing LLM APIs
|
||||
//
|
||||
// This activity works like AWS AssumeRole:
|
||||
// 1. User provides identity + scope of access needed
|
||||
// 2. Activity exchanges credentials with auth server
|
||||
// 3. Returns JWT token valid for a limited time
|
||||
// 4. Caller uses token in subsequent LLM API calls
|
||||
//
|
||||
// Security: Credentials should come from vault/secrets, never hardcoded
|
||||
func AssumeRoleActivity(ctx context.Context, input *AssumeRoleInput) (*AssumeRoleOutput, error) {
|
||||
if input.Identity == "" {
|
||||
return nil, fmt.Errorf("identity is required")
|
||||
}
|
||||
|
||||
if input.Scope == "" {
|
||||
return nil, fmt.Errorf("scope is required (e.g., 'llm:read' or 'llm:read llm:write')")
|
||||
}
|
||||
|
||||
// Get auth server URL from input or environment
|
||||
authServerURL := input.AuthServerURL
|
||||
if authServerURL == "" {
|
||||
authServerURL = os.Getenv("AUTH_SERVER_URL")
|
||||
if authServerURL == "" {
|
||||
return nil, fmt.Errorf("AUTH_SERVER_URL not set in input or environment")
|
||||
}
|
||||
}
|
||||
|
||||
// Get credentials from input or environment
|
||||
clientID := input.ClientID
|
||||
if clientID == "" {
|
||||
clientID = os.Getenv("OAUTH_CLIENT_ID")
|
||||
if clientID == "" {
|
||||
return nil, fmt.Errorf("clientId not provided and OAUTH_CLIENT_ID not set")
|
||||
}
|
||||
}
|
||||
|
||||
clientSecret := input.ClientSecret
|
||||
if clientSecret == "" {
|
||||
clientSecret = os.Getenv("OAUTH_CLIENT_SECRET")
|
||||
if clientSecret == "" {
|
||||
return nil, fmt.Errorf("clientSecret not provided and OAUTH_CLIENT_SECRET not set")
|
||||
}
|
||||
}
|
||||
|
||||
// Set default duration
|
||||
durationSeconds := input.DurationSeconds
|
||||
if durationSeconds == 0 {
|
||||
durationSeconds = 3600 // 1 hour default
|
||||
}
|
||||
if durationSeconds > 86400 {
|
||||
durationSeconds = 86400 // Max 24 hours
|
||||
}
|
||||
|
||||
// Build token request
|
||||
tokenReq := oauthTokenRequest{
|
||||
GrantType: "client_credentials",
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
Scope: input.Scope,
|
||||
Subject: input.Identity, // Assume this identity
|
||||
}
|
||||
|
||||
reqBody, err := json.Marshal(tokenReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal token request: %w", err)
|
||||
}
|
||||
|
||||
// Call auth server
|
||||
httpReq, err := http.NewRequestWithContext(ctx, "POST",
|
||||
fmt.Sprintf("%s/oauth/token", authServerURL),
|
||||
bytes.NewReader(reqBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to call auth server at %s: %w", authServerURL, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("auth server returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var tokenResp oauthTokenResponse
|
||||
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal token response: %w", err)
|
||||
}
|
||||
|
||||
if tokenResp.AccessToken == "" {
|
||||
return nil, fmt.Errorf("auth server returned empty access token")
|
||||
}
|
||||
|
||||
// Calculate expiration
|
||||
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Unix()
|
||||
|
||||
return &AssumeRoleOutput{
|
||||
Token: tokenResp.AccessToken,
|
||||
ExpiresAt: expiresAt,
|
||||
ExpiresIn: tokenResp.ExpiresIn,
|
||||
TokenType: tokenResp.TokenType,
|
||||
}, nil
|
||||
}
|
||||
@@ -85,6 +85,9 @@ func main() {
|
||||
w.RegisterActivity(action.DeploymentPreCheckActivity)
|
||||
w.RegisterActivity(action.ApproveWorkflowActivity)
|
||||
|
||||
// Authentication activities
|
||||
w.RegisterActivity(action.AssumeRoleActivity)
|
||||
|
||||
// Memory activities
|
||||
w.RegisterActivity(action.RetrieveMemoryActivity)
|
||||
|
||||
|
||||
@@ -407,10 +407,76 @@
|
||||
"dependencies": [],
|
||||
"notes": "Network-dependent. First activity to run for context-aware routing. Fast timeout."
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "AssumeRoleActivity",
|
||||
"description": "Request temporary JWT token for accessing LLM APIs (like AWS AssumeRole)",
|
||||
"category": "authentication",
|
||||
"inputs": {
|
||||
"identity": {
|
||||
"type": "string",
|
||||
"description": "User/service identity requesting access",
|
||||
"required": true,
|
||||
"examples": ["[email protected]", "service:poimen-worker"]
|
||||
},
|
||||
"clientId": {
|
||||
"type": "string",
|
||||
"description": "OAuth2 client ID (from vault if not provided)",
|
||||
"required": false
|
||||
},
|
||||
"clientSecret": {
|
||||
"type": "string",
|
||||
"description": "OAuth2 client secret (from vault if not provided)",
|
||||
"required": false
|
||||
},
|
||||
"scope": {
|
||||
"type": "string",
|
||||
"description": "Scope of access (e.g., 'llm:read' or 'llm:read llm:write')",
|
||||
"required": true,
|
||||
"examples": ["llm:read", "llm:read llm:write", "llm:admin"]
|
||||
},
|
||||
"durationSeconds": {
|
||||
"type": "integer",
|
||||
"description": "Token validity duration in seconds (default: 3600, max: 86400)",
|
||||
"required": false,
|
||||
"default": 3600
|
||||
},
|
||||
"authServerUrl": {
|
||||
"type": "string",
|
||||
"description": "Auth server URL (from AUTH_SERVER_URL env if not provided)",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"token": {
|
||||
"type": "string",
|
||||
"description": "JWT token for calling api.riotpiao.com"
|
||||
},
|
||||
"expiresAt": {
|
||||
"type": "integer",
|
||||
"description": "Token expiration time (Unix timestamp)"
|
||||
},
|
||||
"expiresIn": {
|
||||
"type": "integer",
|
||||
"description": "Seconds until token expires"
|
||||
},
|
||||
"tokenType": {
|
||||
"type": "string",
|
||||
"description": "Token type (typically 'Bearer')"
|
||||
}
|
||||
},
|
||||
"constraints": {
|
||||
"defaultTimeout": "30s",
|
||||
"isFlaky": false,
|
||||
"recommendedRetries": 2,
|
||||
"retryBackoff": 1.5,
|
||||
"dependencies": [],
|
||||
"notes": "Must run before LLM Router to provide auth token. Call early in workflow."
|
||||
}
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"totalActivities": 9,
|
||||
"totalActivities": 10,
|
||||
"lastUpdated": "2025-08-31T00:00:00Z",
|
||||
"categories": {
|
||||
"repository": 1,
|
||||
@@ -421,7 +487,8 @@
|
||||
"notification": 1,
|
||||
"approval": 1,
|
||||
"storage": 1,
|
||||
"memory": 1
|
||||
"memory": 1,
|
||||
"authentication": 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -9,6 +9,6 @@ metadata:
|
||||
app.kubernetes.io/name: poimen
|
||||
app.kubernetes.io/component: orchestrator
|
||||
data:
|
||||
GIT_COMMIT: "6eccc86e" # Updated automatically by CI/CD
|
||||
GIT_COMMIT: "f6c6aa03" # Updated automatically by CI/CD
|
||||
GIT_BRANCH: "main"
|
||||
DEPLOYMENT_DATE: "2026-09-04"
|
||||
|
||||
@@ -13,7 +13,7 @@ spec:
|
||||
labels:
|
||||
app: poimen-worker
|
||||
annotations:
|
||||
git-commit: "6eccc86e" # ✅ Updated on each push, triggers rolling restart
|
||||
git-commit: "f6c6aa03" # ✅ Updated on each push, triggers rolling restart
|
||||
deployment-date: "2026-09-04"
|
||||
spec:
|
||||
containers:
|
||||
|
||||
Reference in New Issue
Block a user