Skip to content

Migrate from shared secret to client credential for publisher api#620

Merged
hanzjk merged 6 commits intowso2:mainfrom
nadheesh:migrate-publisher-key-v2
Mar 24, 2026
Merged

Migrate from shared secret to client credential for publisher api#620
hanzjk merged 6 commits intowso2:mainfrom
nadheesh:migrate-publisher-key-v2

Conversation

@nadheesh
Copy link
Copy Markdown
Contributor

@nadheesh nadheesh commented Mar 23, 2026

Purpose

Resolves: #436

Summary by CodeRabbit

  • New Features

    • Switched publisher auth to OAuth2 client-credentials with runtime token management and added bootstrapping for a publisher OAuth client.
  • Chores

    • Removed static publisher API-key configuration and related env/secret entries across deployments, charts and CI.
    • Updated workflows and evaluation invocation to use IdP credentials and secrets.
    • Publisher endpoints now require JWT auth and JWT validation requires a JWKS URL (fail-closed except local dev).
  • Tests

    • Updated tests to use IDP credentials and a mocked token manager.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented Mar 23, 2026

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Replaces static publisher API-key auth with OAuth2 client-credentials: evaluation job obtains Bearer tokens, agent-manager enforces JWT validation via JWKS (dev-only unsigned fallback removed), Thunder bootstraps a publisher OAuth client, and Helm/CI/manifests/configs were updated to remove static API-key usage and wire IdP credentials and secrets.

Changes

Cohort / File(s) Summary
Agent Manager — config & controller
agent-manager-service/config/config.go, agent-manager-service/config/config_loader.go, agent-manager-service/controllers/monitor_scores_publisher_controller.go, agent-manager-service/wiring/wire_gen.go
Removed static APIKey and DefaultIssuer fields and their loading/validation; publisher controller no longer checks x-api-key and now authorizes via JWT claims; wiring updated to stop injecting config into publisher controller.
Agent Manager — routing & JWT validation
agent-manager-service/api/app.go, agent-manager-service/middleware/jwtassertion/auth.go
Publisher routes now use AuthMiddleware. Fallback path that accepted unsigned tokens without JWKS except for a default issuer was removed; missing KEY_MANAGER_JWKS_URL now fails outside local-dev.
Helm / Kubernetes — agent-manager chart & env examples
deployments/helm-charts/wso2-agent-manager/values.yaml, .../templates/agent-manager-service/deployment.yaml, .../templates/agent-manager-service/secret.yaml, .../templates/agent-manager-service/configmap.yaml, agent-manager-service/.env.example
Removed publisherApiKey Helm values and PUBLISHER_API_KEY env/secret rendering; removed KEY_MANAGER_DEFAULT_ISSUER from ConfigMap and example; added amp-publisher-client to key manager audience.
Evaluation extension & workflows
deployments/helm-charts/wso2-amp-evaluation-extension/templates/.../cluster-workflow-monitor-evaluation.yaml, .../workflow-monitor-evaluation.yaml, .../values.yaml
Replaced publisher-api-key parameter with idp-token-url, idp-client-id, publisher-secret-name; workflows now fetch client secret into a Kubernetes Secret (ExternalSecret) and pass IdP params to containers.
Thunder extension & bootstrap
deployments/helm-charts/wso2-amp-thunder-extension/templates/amp-thunder-bootstrap.yaml, deployments/helm-charts/wso2-amp-thunder-extension/values.yaml
Added bootstrap script and values to create/update an ampPublisherClient OAuth2 application (client_credentials) in Thunder and expose clientId/clientSecret values.
Evaluation job (runtime) & tests
evaluation-job/main.py, evaluation-job/test_main.py
Added OAuth2TokenManager to obtain/cache client-credentials tokens; publish_scores switched to Bearer tokens; tests updated to mock token manager and assert Authorization: Bearer ....
CI, scripts, docker, quick-start
.github/workflows/agent-manager-service-pr-checks.yaml, deployments/docker-compose.yml, deployments/scripts/setup-openchoreo.sh, deployments/quick-start/uninstall.sh, deployments/single-cluster/values-openbao.yaml
Removed PUBLISHER_API_KEY usage in CI/helm flags and Helm chart values; updated docker-compose env for key-manager issuer/audience; changed uninstall namespace default; added BAO bootstrap step for publisher secret.
Misc (templates & values)
deployments/helm-charts/wso2-amp-evaluation-extension/..., other templates
Various template/value adjustments to support IdP-based publisher auth and secret retrieval across charts and workflow templates (ExternalSecret, secret name parameters, audience updates).

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant EvalJob as Evaluation Job
  participant IdP as Identity Provider (Thunder)
  participant Agent as Agent Manager Service
  participant JWKS as JWKS/KeyManager

  EvalJob->>IdP: POST /token (client_id, client_secret, grant_type=client_credentials)
  IdP-->>EvalJob: 200 OK (access_token, expires_in)
  EvalJob->>Agent: POST /api/v1/publisher/scores (Authorization: Bearer <token>)
  Agent->>JWKS: Fetch JWKS and validate JWT signature / audience / expiry
  JWKS-->>Agent: JWKS response
  Agent-->>EvalJob: 200 Accepted (scores published)
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 I nibbled old keys and set them free,

Now tokens buzz from client to me,
Thunder crafts a secret snug,
JWKS checks each little hug,
Hops of joy — secure and keyless glee.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Description check ⚠️ Warning The PR description is minimal and only references the linked issue without following the provided template structure (missing Purpose, Goals, Approach, User Stories, Release Notes, Documentation, etc.). Expand the description to follow the template provided in the repository, including Goals, Approach, User Stories, Documentation impact, Release notes, and other relevant sections.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely summarizes the main change: migrating authentication from a shared secret to client credentials for the publisher API.
Linked Issues check ✅ Passed All code changes comprehensively implement the objectives from issue #436: replacing static PUBLISHER_API_KEY with OAuth2 client credentials flow, validating tokens via JWKS, and removing shared secret references from configs and Helm charts.
Out of Scope Changes check ✅ Passed All changes are directly aligned with issue #436 objectives. However, the docker-compose.yml adds KEY_MANAGER_ISSUER and KEY_MANAGER_AUDIENCE without removing them from other configs, which appears slightly broader than strictly necessary.
Docstring Coverage ✅ Passed Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
agent-manager-service/api/app.go (1)

68-79: ⚠️ Potential issue | 🔴 Critical

Add publisher-specific authorization enforcement.

The PublishScores route currently accepts any valid JWT token without restricting access to publisher-only clients. The middleware validates the token signature and expiry but does not enforce a publisher-specific scope or client identity. The controller does not perform additional claim checks, meaning any bearer token with a valid signature can publish scores.

Implement one of the following:

  • Add a dedicated authorization middleware on the publisher sub-mux that validates a scores:publish scope or publisher client identifier
  • Add an explicit scope/client check inside PublishScores before processing the request
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@agent-manager-service/api/app.go` around lines 68 - 79, The publisher routes
accept any valid JWT because there is no publisher-specific authorization; add
enforcement by either (A) creating and attaching a dedicated middleware on the
publisher sub-router that inspects token claims set by params.AuthMiddleware and
rejects requests missing the scores:publish scope or a specific client_id, and
register that middleware on publisherMux (before RegisterMonitorPublisherRoutes
or when wrapping publisherHandler), or (B) add an explicit claim check at the
start of the PublishScores handler in the MonitorScoresPublisherController to
read the JWT claims from the request context, verify the presence of the
scores:publish scope or allowed publisher client identifier, and return 403 when
absent; reference RegisterMonitorPublisherRoutes, publisherMux,
params.AuthMiddleware, and PublishScores when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In
`@deployments/helm-charts/wso2-amp-evaluation-extension/templates/publisher-secret.yaml`:
- Around line 4-5: The Secret "amp-publisher-credentials" is being created in
.Release.Namespace but the Workflow is rendered into
.Values.ampEvaluation.workflowNamespace | default "default", causing a namespace
mismatch; update the publisher-secret.yaml template (metadata.name
amp-publisher-credentials) to set metadata.namespace to the workflow namespace
expression (use .Values.ampEvaluation.workflowNamespace | default "default")
instead of .Release.Namespace so the Secret is created in the same namespace the
Workflow/job runs in and IDP_CLIENT_SECRET can be resolved at runtime.

In `@deployments/helm-charts/wso2-amp-thunder-extension/values.yaml`:
- Around line 159-163: The values.yaml currently hard-codes
ampPublisherClient.clientSecret which creates a shared default; change
ampPublisherClient to accept an externally supplied or bootstrap-generated
secret instead: remove the committed literal for clientSecret and instead add a
reference field (e.g., secretName or leave clientSecret blank/placeholder) so
the chart consumers must supply a secret or the chart templates can generate one
at install time; ensure the deployment/secret templates consume
ampPublisherClient.secretName (or clientSecret value when provided) and remove
the requirement that it be matched in ampEvaluation.publisher.clientSecret by
documenting the needed value (or using a lookup to the created k8s Secret) so no
well-known default is committed.

---

Outside diff comments:
In `@agent-manager-service/api/app.go`:
- Around line 68-79: The publisher routes accept any valid JWT because there is
no publisher-specific authorization; add enforcement by either (A) creating and
attaching a dedicated middleware on the publisher sub-router that inspects token
claims set by params.AuthMiddleware and rejects requests missing the
scores:publish scope or a specific client_id, and register that middleware on
publisherMux (before RegisterMonitorPublisherRoutes or when wrapping
publisherHandler), or (B) add an explicit claim check at the start of the
PublishScores handler in the MonitorScoresPublisherController to read the JWT
claims from the request context, verify the presence of the scores:publish scope
or allowed publisher client identifier, and return 403 when absent; reference
RegisterMonitorPublisherRoutes, publisherMux, params.AuthMiddleware, and
PublishScores when making the change.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 619d389e-3991-4312-8931-0cc73653947f

📥 Commits

Reviewing files that changed from the base of the PR and between de07609 and c494b25.

📒 Files selected for processing (20)
  • .github/workflows/agent-manager-service-pr-checks.yaml
  • .gitignore
  • agent-manager-service/.env.example
  • agent-manager-service/api/app.go
  • agent-manager-service/config/config.go
  • agent-manager-service/config/config_loader.go
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go
  • agent-manager-service/wiring/wire_gen.go
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/deployment.yaml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/secret.yaml
  • deployments/helm-charts/wso2-agent-manager/values.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/publisher-secret.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/workflow-templates/cluster-workflow-monitor-evaluation.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/workflow-templates/workflow-monitor-evaluation.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/values.yaml
  • deployments/helm-charts/wso2-amp-thunder-extension/templates/amp-thunder-bootstrap.yaml
  • deployments/helm-charts/wso2-amp-thunder-extension/values.yaml
  • deployments/scripts/setup-openchoreo.sh
  • evaluation-job/main.py
  • evaluation-job/test_main.py
💤 Files with no reviewable changes (7)
  • .github/workflows/agent-manager-service-pr-checks.yaml
  • agent-manager-service/.env.example
  • agent-manager-service/config/config.go
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/deployment.yaml
  • deployments/helm-charts/wso2-agent-manager/values.yaml
  • agent-manager-service/config/config_loader.go
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go

Comment thread deployments/helm-charts/wso2-amp-thunder-extension/values.yaml
@nadheesh nadheesh force-pushed the migrate-publisher-key-v2 branch from cc40524 to f016c79 Compare March 23, 2026 10:23
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
agent-manager-service/middleware/jwtassertion/auth.go (1)

199-225: ⚠️ Potential issue | 🔴 Critical

Fail closed when KEY_MANAGER_JWKS_URL is missing.

This branch still authenticates by base64-decoding the payload and trusting its claims. Because KEY_MANAGER_JWKS_URL defaults to empty in agent-manager-service/config/config_loader.go and is not set in deployments/docker-compose.yml, a forged three-part JWT with an allowed iss/aud can pass this middleware without ever proving Thunder issued it. That is an auth bypass against the new publisher flow.

🔒 Proposed fix
-	} else {
-		// No JWKS URL configured - extract claims without signature validation
-		// This is only allowed for tokens from configured trusted issuers
-		extractedClaims, err := extractClaimsFromJWT(tokenString)
-		if err != nil {
-			return nil, fmt.Errorf("failed to extract claims: %w", err)
-		}
-		claims = extractedClaims
-
-		// Ensures we only skip signature validation for configured trusted issuers
-		if err := validateIssuer(claims.Issuer, cfg.KeyManagerConfigurations.Issuer); err != nil {
-			return nil, fmt.Errorf("JWKS signature validation required for issuer '%s'",
-				claims.Issuer)
-		}
-
-		// Validate expiration using RegisteredClaims.ExpiresAt
-		if claims.ExpiresAt != nil {
-			if !claims.ExpiresAt.After(time.Now()) {
-				return nil, fmt.Errorf("token has expired")
-			}
-		}
-
-		// Validate audience
-		if err := validateAudience(claims.Audience, cfg.KeyManagerConfigurations.Audience); err != nil {
-			return nil, err
-		}
-
-		return claims, nil
+	} else {
+		return nil, fmt.Errorf("KEY_MANAGER_JWKS_URL must be configured for JWT validation")
 	}

If you still need an insecure local escape hatch, gate it behind an explicit dev-only flag and keep it off by default.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@agent-manager-service/middleware/jwtassertion/auth.go` around lines 199 -
225, The code currently allows unsigned JWTs when
cfg.KeyManagerConfigurations.JwksURL is empty by calling extractClaimsFromJWT
and then validating issuer/audience only; instead, fail closed: in the auth.go
branch that handles missing JWKS (where extractClaimsFromJWT, validateIssuer,
validateAudience are used and cfg.KeyManagerConfigurations.JwksURL is checked),
return an error if KeyManagerConfigurations.JwksURL is empty unless an explicit
dev-only escape hatch is enabled (e.g., a new config flag like
KeyManagerConfigurations.AllowInsecureLocalTokens), and only permit the insecure
extraction when that flag is true and the environment is dev; otherwise require
a JWKS URL and reject the token. Ensure error messages mention the missing JWKS
(KEY_MANAGER_JWKS_URL) so callers know why authentication failed.
♻️ Duplicate comments (1)
deployments/helm-charts/wso2-amp-thunder-extension/values.yaml (1)

160-163: ⚠️ Potential issue | 🟠 Major

Don't ship a default publisher client secret.

ampPublisherClient.clientSecret is still a committed literal, and the inline note requires OpenBao to mirror the same value. That recreates the well-known/shared-secret and manual cross-chart sync problem this migration is supposed to eliminate.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@deployments/helm-charts/wso2-amp-thunder-extension/values.yaml` around lines
160 - 163, The values.yaml currently ships a committed literal for
ampPublisherClient.clientSecret which recreates a shared-secret drift issue;
remove the hardcoded "amp-publisher-client-secret" value (or set
ampPublisherClient.clientSecret to null/empty) and update the inline comment to
instruct operators to supply the secret via external secret management/OpenBao
(do not hardcode), ensuring the chart consumes the secret name or value injected
at deploy time rather than embedding it in values.yaml for ampPublisherClient.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@deployments/helm-charts/wso2-agent-manager/values.yaml`:
- Around line 123-126: The publisher endpoint currently only uses
JWTAuthMiddleware and lacks scope checks; update the publisher route or the
PublishScores handler to call the existing HasAllScopes(ctx,
[]string{"scores:publish"}) (or equivalent scope-check helper) after JWT
validation and return a 403 on failure. Locate the route registration that wires
JWTAuthMiddleware to the PublishScores controller and add the scope-check
invocation there (or inside PublishScores) ensuring the middleware chain
enforces the "scores:publish" scope before processing the request.

---

Outside diff comments:
In `@agent-manager-service/middleware/jwtassertion/auth.go`:
- Around line 199-225: The code currently allows unsigned JWTs when
cfg.KeyManagerConfigurations.JwksURL is empty by calling extractClaimsFromJWT
and then validating issuer/audience only; instead, fail closed: in the auth.go
branch that handles missing JWKS (where extractClaimsFromJWT, validateIssuer,
validateAudience are used and cfg.KeyManagerConfigurations.JwksURL is checked),
return an error if KeyManagerConfigurations.JwksURL is empty unless an explicit
dev-only escape hatch is enabled (e.g., a new config flag like
KeyManagerConfigurations.AllowInsecureLocalTokens), and only permit the insecure
extraction when that flag is true and the environment is dev; otherwise require
a JWKS URL and reject the token. Ensure error messages mention the missing JWKS
(KEY_MANAGER_JWKS_URL) so callers know why authentication failed.

---

Duplicate comments:
In `@deployments/helm-charts/wso2-amp-thunder-extension/values.yaml`:
- Around line 160-163: The values.yaml currently ships a committed literal for
ampPublisherClient.clientSecret which recreates a shared-secret drift issue;
remove the hardcoded "amp-publisher-client-secret" value (or set
ampPublisherClient.clientSecret to null/empty) and update the inline comment to
instruct operators to supply the secret via external secret management/OpenBao
(do not hardcode), ensuring the chart consumes the secret name or value injected
at deploy time rather than embedding it in values.yaml for ampPublisherClient.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 8cd9f011-65ad-412b-9a7d-d04c759fa254

📥 Commits

Reviewing files that changed from the base of the PR and between c494b25 and f016c79.

📒 Files selected for processing (23)
  • .github/workflows/agent-manager-service-pr-checks.yaml
  • agent-manager-service/.env.example
  • agent-manager-service/api/app.go
  • agent-manager-service/config/config.go
  • agent-manager-service/config/config_loader.go
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go
  • agent-manager-service/middleware/jwtassertion/auth.go
  • agent-manager-service/wiring/wire_gen.go
  • deployments/docker-compose.yml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/configmap.yaml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/deployment.yaml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/secret.yaml
  • deployments/helm-charts/wso2-agent-manager/values.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/workflow-templates/cluster-workflow-monitor-evaluation.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/workflow-templates/workflow-monitor-evaluation.yaml
  • deployments/helm-charts/wso2-amp-evaluation-extension/values.yaml
  • deployments/helm-charts/wso2-amp-thunder-extension/templates/amp-thunder-bootstrap.yaml
  • deployments/helm-charts/wso2-amp-thunder-extension/values.yaml
  • deployments/quick-start/uninstall.sh
  • deployments/scripts/setup-openchoreo.sh
  • deployments/single-cluster/values-openbao.yaml
  • evaluation-job/main.py
  • evaluation-job/test_main.py
💤 Files with no reviewable changes (5)
  • .github/workflows/agent-manager-service-pr-checks.yaml
  • agent-manager-service/.env.example
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/configmap.yaml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/deployment.yaml
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go
✅ Files skipped from review due to trivial changes (2)
  • deployments/single-cluster/values-openbao.yaml
  • agent-manager-service/wiring/wire_gen.go
🚧 Files skipped from review as they are similar to previous changes (7)
  • agent-manager-service/api/app.go
  • deployments/helm-charts/wso2-amp-evaluation-extension/values.yaml
  • deployments/helm-charts/wso2-agent-manager/templates/agent-manager-service/secret.yaml
  • evaluation-job/main.py
  • evaluation-job/test_main.py
  • agent-manager-service/config/config.go
  • deployments/helm-charts/wso2-amp-evaluation-extension/templates/workflow-templates/cluster-workflow-monitor-evaluation.yaml

Comment thread deployments/helm-charts/wso2-agent-manager/values.yaml
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
agent-manager-service/middleware/jwtassertion/auth.go (1)

155-200: ⚠️ Potential issue | 🟠 Major

Refresh JWKS once on kid miss.

Now that Line 199 makes JWKS the only validation path, a normal IdP signing-key rotation can fail closed for up to jwksCacheTTL: the cached set is reused, the new kid is not found, and every publisher request gets a 401 until the cache expires. Retry the lookup once with a forced JWKS refresh before rejecting the token so rotation works without waiting an hour.

🔁 Possible fix
-			jwks, err := fetchJWKS(cfg.KeyManagerConfigurations.JWKSUrl)
+			jwks, err := fetchJWKS(cfg.KeyManagerConfigurations.JWKSUrl, false)
 			if err != nil {
 				return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
 			}

-			// Find the key with matching kid
-			for _, key := range jwks.Keys {
-				if key.Kid == kid {
-					return convertJWKToPublicKey(&key)
-				}
-			}
-
-			return nil, fmt.Errorf("unable to find key with kid: %s", kid)
+			findKey := func(keys []JSONWebKey) (*rsa.PublicKey, error) {
+				for _, key := range keys {
+					if key.Kid == kid {
+						return convertJWKToPublicKey(&key)
+					}
+				}
+				return nil, nil
+			}
+
+			if key, err := findKey(jwks.Keys); err != nil {
+				return nil, err
+			} else if key != nil {
+				return key, nil
+			}
+
+			jwks, err = fetchJWKS(cfg.KeyManagerConfigurations.JWKSUrl, true)
+			if err != nil {
+				return nil, fmt.Errorf("failed to refresh JWKS: %w", err)
+			}
+			if key, err := findKey(jwks.Keys); err != nil {
+				return nil, err
+			} else if key != nil {
+				return key, nil
+			}
+
+			return nil, fmt.Errorf("unable to find key with kid: %s", kid)
-func fetchJWKS(jwksURL string) (*JWKS, error) {
+func fetchJWKS(jwksURL string, forceRefresh bool) (*JWKS, error) {
 	jwksCacheMutex.RLock()
-	if jwksCache != nil && time.Since(jwksCacheTime) < jwksCacheTTL {
+	if !forceRefresh && jwksCache != nil && time.Since(jwksCacheTime) < jwksCacheTTL {
 		defer jwksCacheMutex.RUnlock()
 		return jwksCache, nil
 	}
 	jwksCacheMutex.RUnlock()
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@agent-manager-service/middleware/jwtassertion/auth.go` around lines 155 -
200, The JWKS lookup should retry once with a forced refresh when the token's
kid is not found to avoid failing closed during key rotation: inside the
jwt.ParseWithClaims keyfunc (where you call
fetchJWKS(cfg.KeyManagerConfigurations.JWKSUrl) and loop jwks.Keys to match
key.Kid to kid), on the path that currently returns "unable to find key with
kid" call fetchJWKS again forcing a refresh (or call a new
fetchJWKSForced/fetchJWKS(cfg..., force=true)) and re-scan the returned
jwks.Keys one more time; only if the second scan still misses return the error.
Keep using convertJWKToPublicKey and preserve the existing signing-method/kid
checks and error wrapping.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@agent-manager-service/middleware/jwtassertion/auth.go`:
- Around line 155-200: The JWKS lookup should retry once with a forced refresh
when the token's kid is not found to avoid failing closed during key rotation:
inside the jwt.ParseWithClaims keyfunc (where you call
fetchJWKS(cfg.KeyManagerConfigurations.JWKSUrl) and loop jwks.Keys to match
key.Kid to kid), on the path that currently returns "unable to find key with
kid" call fetchJWKS again forcing a refresh (or call a new
fetchJWKSForced/fetchJWKS(cfg..., force=true)) and re-scan the returned
jwks.Keys one more time; only if the second scan still misses return the error.
Keep using convertJWKToPublicKey and preserve the existing signing-method/kid
checks and error wrapping.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e72ffe0a-07f8-4b91-9006-d4b89de9a2e3

📥 Commits

Reviewing files that changed from the base of the PR and between 58a7d4d and 5713324.

📒 Files selected for processing (3)
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go
  • agent-manager-service/middleware/jwtassertion/auth.go
  • deployments/docker-compose.yml
✅ Files skipped from review due to trivial changes (1)
  • deployments/docker-compose.yml
🚧 Files skipped from review as they are similar to previous changes (1)
  • agent-manager-service/controllers/monitor_scores_publisher_controller.go

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
agent-manager-service/middleware/jwtassertion/auth.go (1)

198-211: Also validate nbf in the unsigned local-dev path.

Fail-closing non-local environments is a good change. The remaining gap is that this branch only rejects expired tokens; a token with a future nbf still passes locally, which can hide time-based auth issues until later.

⏱️ Small parity fix
 		if claims.ExpiresAt != nil && !claims.ExpiresAt.After(time.Now()) {
 			return nil, fmt.Errorf("token has expired")
 		}
+		if claims.NotBefore != nil && claims.NotBefore.After(time.Now()) {
+			return nil, fmt.Errorf("token is not valid yet")
+		}
 	} else {
 		return nil, fmt.Errorf("KEY_MANAGER_JWKS_URL must be configured for JWT validation")
 	}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@agent-manager-service/middleware/jwtassertion/auth.go` around lines 198 -
211, In the local-dev branch (when cfg.IsLocalDevEnv) you currently only check
claims.ExpiresAt; update the unsigned-path logic in the function that calls
extractClaimsFromJWT(tokenString) to also reject tokens whose claims.NotBefore
(nbf) is set in the future: after assigning claims = extractedClaims, add a
check that if claims.NotBefore != nil and claims.NotBefore.After(time.Now())
then return an error like "token not valid yet (nbf)"; ensure you reference the
same claims variable and preserve the existing expiresAt check so both
time-based validations run in extractClaimsFromJWT caller.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@agent-manager-service/middleware/jwtassertion/auth.go`:
- Around line 198-211: In the local-dev branch (when cfg.IsLocalDevEnv) you
currently only check claims.ExpiresAt; update the unsigned-path logic in the
function that calls extractClaimsFromJWT(tokenString) to also reject tokens
whose claims.NotBefore (nbf) is set in the future: after assigning claims =
extractedClaims, add a check that if claims.NotBefore != nil and
claims.NotBefore.After(time.Now()) then return an error like "token not valid
yet (nbf)"; ensure you reference the same claims variable and preserve the
existing expiresAt check so both time-based validations run in
extractClaimsFromJWT caller.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 17b07032-2eeb-46fe-bff2-4851f8caeb00

📥 Commits

Reviewing files that changed from the base of the PR and between 5713324 and 763c7b4.

📒 Files selected for processing (2)
  • agent-manager-service/middleware/jwtassertion/auth.go
  • deployments/docker-compose.yml
✅ Files skipped from review due to trivial changes (1)
  • deployments/docker-compose.yml

--set ampEvaluation.image.repository="amp-evaluation-monitor" \
--set ampEvaluation.publisher.endpoint="http://agent-manager-service:8080" \
--set ampEvaluation.publisher.apiKey="dev-publisher-api-key"
--set ampEvaluation.publisher.idpTokenUrl="http://amp-thunder-extension-service.amp-thunder.svc.cluster.local:8090/oauth2/token" \
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the quick start guide script also updated?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The quick-start install.sh / install-helpers.sh never explicitly set ampEvaluation.publisher.apiKey — that override only existed in setup-openchoreo.sh, which is updated in this PR. The quick-start deploys the evaluation and Thunder extension Helm charts using their default values, which now carry the correct OAuth2 config (idpTokenUrl, clientId). OpenBao is seeded via the values-openbao.yaml fetched from the repo at install time, which already has the amp-publisher-client-secret entry added in this PR. The only quick-start file that needed changing was uninstall.sh (evaluation namespace fix), and that's already included.

if c.apiKey == "" || subtle.ConstantTimeCompare([]byte(providedKey), []byte(c.apiKey)) != 1 {
log.Warn("Unauthorized access attempt to publish scores", "ip", getClientIP(r))
http.Error(w, "Invalid or missing API key", http.StatusUnauthorized)
// Enforce publisher-only access: only the amp-publisher-client may call this endpoint
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we remove this authorization logic from controller and move it to a middleware

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — moved to a PublisherClientAuthMiddleware in the jwtassertion package (commit 3dcdb0c4). The client ID is an internal constant inside the middleware. It is applied inside RegisterMonitorPublisherRoutes itself, so the call site in app.go stays consistent with all other route registrations.

Comment thread agent-manager-service/api/app.go Outdated

// Register publisher routes outside JWT middleware (uses API-key auth instead)
// Register publisher routes with JWT middleware (uses OAuth2 client_credentials token)
publisherMux := http.NewServeMux()
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need a separate mux for the publisher? Can Publisher routes now be part of apiMux?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — the separate publisherMux and its duplicated middleware stack have been removed (commit 3dcdb0c4). Publisher routes are now registered directly on apiMux under the /publisher/ path prefix, so the external URL (/api/v1/publisher/monitors/.../scores) is unchanged.

Comment thread agent-manager-service/middleware/jwtassertion/auth.go
@hanzjk hanzjk merged commit 49c8f55 into wso2:main Mar 24, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Migrate evaluation publisher auth from static API key to client credentials

2 participants