From 4a34c8e67200d7b7260be8f6cc2c0303b5e396eb Mon Sep 17 00:00:00 2001 From: Test Date: Fri, 4 Sep 2026 14:11:54 -0700 Subject: [PATCH] feat: add AssumeRoleActivity for temporary LLM API token grants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- README.md | 114 +++++++++++ action/assume_role.go | 178 ++++++++++++++++++ cmd/worker/main.go | 3 + internal/routing/activity_knowledge_base.json | 71 ++++++- k8s/git-commit.yaml | 2 +- k8s/worker-deployment.yaml | 2 +- 6 files changed, 366 insertions(+), 4 deletions(-) create mode 100644 action/assume_role.go diff --git a/README.md b/README.md index de93506..4e9a9f2 100644 --- a/README.md +++ b/README.md @@ -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: "user@company.com", // 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: "user@company.com", + 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 diff --git a/action/assume_role.go b/action/assume_role.go new file mode 100644 index 0000000..7828d12 --- /dev/null +++ b/action/assume_role.go @@ -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 +} diff --git a/cmd/worker/main.go b/cmd/worker/main.go index 17da652..af18955 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -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) diff --git a/internal/routing/activity_knowledge_base.json b/internal/routing/activity_knowledge_base.json index 884d473..8d0697c 100644 --- a/internal/routing/activity_knowledge_base.json +++ b/internal/routing/activity_knowledge_base.json @@ -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": ["user@company.com", "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 } } } diff --git a/k8s/git-commit.yaml b/k8s/git-commit.yaml index 066ae05..0028857 100644 --- a/k8s/git-commit.yaml +++ b/k8s/git-commit.yaml @@ -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" diff --git a/k8s/worker-deployment.yaml b/k8s/worker-deployment.yaml index f3d8bc7..2b3987e 100644 --- a/k8s/worker-deployment.yaml +++ b/k8s/worker-deployment.yaml @@ -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: