Feat/add oauth2 and jwt supports#34
Conversation
Standardize AFM authentication field schemas (issue wso2#35), Phase 1. - parser.bal: validate every authentication block at parse time against its type's field schema, across all three sites (model, MCP transport, webhook subscription). Reject unknown types, enforce required fields, and reject fields that do not belong to the type. bearer/basic/api-key are fully validated; jwt/oauth2 are recognized with field validation deferred to their phase. - agent.bal: add an explicit api-key case to mapToHttpClientAuth with a clear message. Ballerina's http client has no raw-header auth variant, so api-key is not yet supported for MCP/webhook transport (it remains supported for model providers). - tests: cover schema validation, the api-key mapping case, and parse-level rejection.
Standardize AFM authentication field schemas (issue wso2#35), Phase 1. - models.py: ClientAuthentication now validates at parse time - reject unknown types, enforce per-type required fields, and reject fields that do not belong to the type (catches typos and cross-type fields). Add an optional header_name for api-key. bearer/basic/api-key are fully validated; jwt/oauth2 are recognized with field validation deferred to their phase. - mcp.py: wire header_name into ApiKeyAuth so api-key can target custom headers (e.g. X-API-Key) instead of always Authorization. - tests: cover the validation rules and the custom header_name wiring.
Standardize AFM authentication field schemas (issue wso2#35), Phase 2. The runtime now mints and signs a JWT from a key plus claims and sends it as a bearer token, for MCP and webhook transports. - agent.bal: map a jwt authentication block to http:JwtIssuerConfig (issuer->iss, subject->username/sub, key_id->kid header, custom_claims, expiry_seconds->expTime). HMAC algorithms (HS256/384/512) use signing_key as a shared secret; asymmetric algorithms use it as a PEM key file path. - parser.bal: validate jwt fields at parse time (issuer, audience and signing_key required); jwt is no longer treated as deferred. - tests: HMAC and RS256 mapping, plus jwt schema validation.
Standardize AFM authentication field schemas (issue wso2#35), Phase 2. - models.py: add the jwt fields to ClientAuthentication (issuer, audience, signing_key, algorithm, key_id, subject, custom_claims, expiry_seconds) and fold jwt into the parse-time validator (issuer, audience and signing_key required; unknown jwt fields rejected). oauth2 stays deferred. - tests: jwt valid/invalid schema cases, audience list, unknown field.
Standardize AFM authentication field schemas (issue wso2#35), Phase 2. - mcp.py: add JwtAuth (httpx.Auth) that mints and signs a JWT per request via PyJWT and sends it as a bearer token, and wire it into build_httpx_auth. HMAC algorithms use signing_key as a shared secret; asymmetric algorithms read it as a PEM key file. Claims are assembled to match the Ballerina runtime (sets nbf; custom_claims merged last). Only oauth2 remains not yet supported. - pyproject.toml / uv.lock: declare the pyjwt[crypto] dependency. - tests: HMAC sign+decode round-trip, RS256 with a generated key file, default-RS256, and custom header_name.
Standardize AFM authentication field schemas (issue wso2#35), Phase 3. - agent.bal: add buildOAuth2GrantConfig, mapping an oauth2 authentication block to the matching http:OAuth2*GrantConfig so the HTTP client performs the token exchange (and refresh) natively. Supports the client_credentials, password, refresh_token and jwt_bearer grants; field names are mapped from the AFM snake_case names to the connector camelCase names per grant. - parser.bal: validate oauth2 at parse time. grant_type is a required discriminator that selects the required/optional fields; unknown grants and fields not allowed for the grant are rejected. All five auth types are now recognized and validated.
Standardize AFM authentication field schemas (issue wso2#35), Phase 3. - models.py: validate oauth2 at parse time. grant_type is a required discriminator selecting the required/optional fields per grant (client_credentials, password, refresh_token, jwt_bearer); unknown grants and fields not allowed for the grant are rejected. - mcp.py: add OAuth2Auth (httpx.Auth) that obtains an access token via a token exchange, caches it until expiry, and sends it as a bearer token. Implements both sync and async flows; client credentials are sent as HTTP Basic to match the Ballerina runtime, and jwt_bearer uses the RFC 7523 grant URN with a user-supplied assertion. No new dependency. All five auth types are now supported (none remain not yet supported).
Standardize AFM authentication field schemas (issue wso2#35), Phase 3. - cover parse-time validation for all four grants (valid configs, missing grant_type, unsupported grant, missing required fields, fields not allowed for the grant, case-insensitive grant_type) and the runtime token-exchange mapping, in both runtimes.
|
Warning Review limit reached
More reviews will be available in 49 minutes and 9 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more credits in the billing tab to continue. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (7)
📝 WalkthroughWalkthroughJWT and OAuth2 authentication support is added to the Ballerina and Python AFM interpreters. Ballerina now validates authentication metadata during parsing and converts JWT/OAuth2 auth into HTTP client config. Python extends Sequence Diagram(s)sequenceDiagram
participant Parser
participant ClientAuthentication
participant JwtAuth
participant OAuth2Auth
participant TokenEndpoint
Parser->>ClientAuthentication: validate auth fields
ClientAuthentication-->>Parser: valid or error
Parser->>JwtAuth: build JWT request auth
Parser->>OAuth2Auth: build OAuth2 request auth
OAuth2Auth->>TokenEndpoint: request access token
TokenEndpoint-->>OAuth2Auth: access_token
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
python-interpreter/packages/afm-langchain/tests/test_mcp.py (1)
188-212: ⚡ Quick winMake JWT TTL assertion tolerant to second-boundary timing.
The strict
exp - iat == 600check can be intermittently brittle. Prefer a small tolerance window.Proposed test adjustment
- assert decoded["exp"] - decoded["iat"] == 600 + ttl = decoded["exp"] - decoded["iat"] + assert 599 <= ttl <= 600🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/tests/test_mcp.py` around lines 188 - 212, The assertion `assert decoded["exp"] - decoded["iat"] == 600` uses a strict equality check which is brittle due to potential timing variations between when the `iat` and `exp` timestamps are set. Replace this strict equality check with a tolerance-based assertion that allows for a small time window (e.g., checking that the difference is within approximately 600 seconds plus or minus a small tolerance, such as 1-2 seconds) to make the test more reliable and tolerant of second-boundary timing edge cases.python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py (2)
111-111: 💤 Low value
custom_claimscan overwrite standard JWT claims.Using
claims.update(self.custom_claims)allowscustom_claimsto overwrite standard claims likeiss,aud,exp,iat, ornbf. If this is intentional for flexibility, consider documenting it. Otherwise, consider updating only keys not already present.Option to preserve standard claims
- claims.update(self.custom_claims) + for key, value in self.custom_claims.items(): + if key not in claims: + claims[key] = value🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py` at line 111, The `claims.update(self.custom_claims)` call at line 111 allows custom claims to overwrite standard JWT claims such as iss, aud, exp, iat, and nbf, which could be unintended. Instead of using update which unconditionally overwrites, modify the code to only add custom claims that are not already present in the claims dictionary. This preserves the standard JWT claims while still allowing custom claims to be added without conflicts.
156-157: ⚖️ Poor tradeoffToken cache is not thread-safe.
The
_tokenand_expires_atfields are accessed without synchronization. If multiple threads callsync_auth_flowconcurrently when the cache is expired, they may all fetch tokens simultaneously. This causes redundant requests but not data corruption. Consider documenting this limitation or adding a lock if high-concurrency scenarios are expected.Also applies to: 251-259
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py` around lines 156 - 157, The `_token` and `_expires_at` instance variables are accessed without synchronization, creating a race condition where multiple threads calling `sync_auth_flow` concurrently when the cache is expired can fetch tokens simultaneously. Either add thread-safe synchronization (such as a threading.Lock) around the token cache access and update logic in the method that reads and writes these fields, or add clear documentation to the class explaining that this implementation is not thread-safe and should only be used in single-threaded contexts. If choosing synchronization, ensure the lock protects all read-modify-write operations on the token cache to prevent redundant concurrent token requests.ballerina-interpreter/agent.bal (1)
443-448: 💤 Low valueMinor inconsistency in
passwordgrant handling.The
passwordgrant usesaddOptionalforclientIdandclientSecret, but perparser.balthese are required fields for this grant type. The code works correctly since parse-time validation guarantees their presence, but for consistency withclient_credentials(which assigns them directly), consider assigning them directly here as well.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ballerina-interpreter/agent.bal` around lines 443 - 448, In the "password" grant case within the switch statement, the clientId and clientSecret fields are being added to grantConfig using the addOptional function, but since these are required fields for this grant type (validated at parse time), they should be assigned directly to the grantConfig map instead. Replace the two addOptional calls for clientId and clientSecret with direct assignments to grantConfig, matching the pattern used in the client_credentials grant for consistency.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ballerina-interpreter/parser.bal`:
- Around line 364-369: The required-field validation in the auth parsing logic
currently only checks if keys are present in auth.keys(), but does not validate
that those fields have non-null values. This allows null-valued fields to pass
parse-time validation even though they would be rejected later at runtime by
ensureType() in getAuthTokenOrApiKey(). Modify the provided array derivation to
filter out keys whose corresponding values are null, ensuring required fields
both exist and have non-null values. Apply this same filtering pattern at both
validation locations: in the requiredAuthFields loop check around lines 364-369
and in the OAuth2 grant validation around lines 429-434.
In `@python-interpreter/packages/afm-core/src/afm/models.py`:
- Around line 139-143: The current logic building the `provided` set at the
initialization filtering excludes `None` values, preventing detection of fields
explicitly set to null. Replace the set comprehension that skips `None` values
with `model_fields_set` to capture all fields that were explicitly provided by
the user. Then update the validation logic that checks for unsupported fields
(which currently relies on the `provided` set) to use this new approach based on
`model_fields_set`, ensuring that fields like `token` explicitly set to `null`
are properly detected and rejected. This change should be consistently applied
wherever the `provided` set is used for validation checks.
In `@python-interpreter/packages/afm-langchain/pyproject.toml`:
- Line 20: The pyjwt[crypto] dependency in the pyproject.toml file is pinned to
version 2.10.0, which contains security vulnerabilities including
CVE-2024-53861. Update the version constraint for pyjwt[crypto] from >=2.10.0 to
>=2.13.0 to address these security advisories and ensure a secure version is
used.
---
Nitpick comments:
In `@ballerina-interpreter/agent.bal`:
- Around line 443-448: In the "password" grant case within the switch statement,
the clientId and clientSecret fields are being added to grantConfig using the
addOptional function, but since these are required fields for this grant type
(validated at parse time), they should be assigned directly to the grantConfig
map instead. Replace the two addOptional calls for clientId and clientSecret
with direct assignments to grantConfig, matching the pattern used in the
client_credentials grant for consistency.
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py`:
- Line 111: The `claims.update(self.custom_claims)` call at line 111 allows
custom claims to overwrite standard JWT claims such as iss, aud, exp, iat, and
nbf, which could be unintended. Instead of using update which unconditionally
overwrites, modify the code to only add custom claims that are not already
present in the claims dictionary. This preserves the standard JWT claims while
still allowing custom claims to be added without conflicts.
- Around line 156-157: The `_token` and `_expires_at` instance variables are
accessed without synchronization, creating a race condition where multiple
threads calling `sync_auth_flow` concurrently when the cache is expired can
fetch tokens simultaneously. Either add thread-safe synchronization (such as a
threading.Lock) around the token cache access and update logic in the method
that reads and writes these fields, or add clear documentation to the class
explaining that this implementation is not thread-safe and should only be used
in single-threaded contexts. If choosing synchronization, ensure the lock
protects all read-modify-write operations on the token cache to prevent
redundant concurrent token requests.
In `@python-interpreter/packages/afm-langchain/tests/test_mcp.py`:
- Around line 188-212: The assertion `assert decoded["exp"] - decoded["iat"] ==
600` uses a strict equality check which is brittle due to potential timing
variations between when the `iat` and `exp` timestamps are set. Replace this
strict equality check with a tolerance-based assertion that allows for a small
time window (e.g., checking that the difference is within approximately 600
seconds plus or minus a small tolerance, such as 1-2 seconds) to make the test
more reliable and tolerant of second-boundary timing edge cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 80f2db5c-afa3-457c-97b2-2604a8994880
⛔ Files ignored due to path filters (1)
python-interpreter/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (8)
ballerina-interpreter/agent.balballerina-interpreter/parser.balballerina-interpreter/tests/main_test.balpython-interpreter/packages/afm-core/src/afm/models.pypython-interpreter/packages/afm-core/tests/test_parser.pypython-interpreter/packages/afm-langchain/pyproject.tomlpython-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.pypython-interpreter/packages/afm-langchain/tests/test_mcp.py
| string[] provided = from string key in auth.keys() where key != "type" select key; | ||
|
|
||
| foreach string required in requiredAuthFields(authType) { | ||
| if provided.indexOf(required) is () { | ||
| return error(string `type '${authType}' requires '${required}' field`); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's locate and examine the parser.bal file
fd -t f "parser.bal" | head -20Repository: wso2/reference-implementations-afm
Length of output: 111
🏁 Script executed:
# Check the file size to ensure we can read it
wc -l ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 115
🏁 Script executed:
# Read the specific lines mentioned in the review (364-369 and 429-434)
sed -n '360,375p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 613
🏁 Script executed:
# Also read the second location
sed -n '425,440p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 691
🏁 Script executed:
# Search for the function definitions to understand context
rg -n "validateMetadataAuthentication|validateAuthentication|validateOAuth2" ballerina-interpreter/parser.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 646
🏁 Script executed:
# Look for test files related to auth validation
fd -t f ".*test.*\.bal" | grep -i "auth\|parse" | head -20Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Find ClientAuthentication type definition
rg -n "type ClientAuthentication" ballerina-interpreter/Repository: wso2/reference-implementations-afm
Length of output: 148
🏁 Script executed:
# Search for test files
find . -type f -name "*.bal" | xargs grep -l "test\|Test" | grep -i "auth\|parse" | head -10Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Look at the broader context of validateAuthentication function
sed -n '345,380p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 1195
🏁 Script executed:
# Check if there's any null checking already
rg -n "is \(\)" ballerina-interpreter/parser.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 678
🏁 Script executed:
# Look for auth field initialization patterns
rg -n "authentication.*=" ballerina-interpreter/parser.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 136
🏁 Script executed:
# Check the ClientAuthentication record definition
sed -n '49,80p' ballerina-interpreter/types.balRepository: wso2/reference-implementations-afm
Length of output: 603
🏁 Script executed:
# Look for any existing parse tests
find . -type f -name "*.bal" -o -name "*.yaml" -o -name "*.json" | xargs grep -l "bearer\|basic\|api-key" 2>/dev/null | head -10Repository: wso2/reference-implementations-afm
Length of output: 236
🏁 Script executed:
# Check for test resources or test data
find . -type d -name "*test*" | head -10Repository: wso2/reference-implementations-afm
Length of output: 203
🏁 Script executed:
# Check the test files for auth validation testing
cat -n ballerina-interpreter/tests/agent_test.bal | head -50Repository: wso2/reference-implementations-afm
Length of output: 2020
🏁 Script executed:
# Check main_test.bal
cat -n ballerina-interpreter/tests/main_test.bal | head -50Repository: wso2/reference-implementations-afm
Length of output: 1932
🏁 Script executed:
# Look for YAML/JSON test files with authentication configs
find . -type f \( -name "*.yaml" -o -name "*.yml" -o -name "*.json" \) | xargs grep -l "bearer\|basic\|oauth2" 2>/dev/null | head -10Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check if there are AFM YAML files with auth sections
find . -type f -name "*.yaml" -o -name "*.yml" | xargs grep -A5 "authentication:" 2>/dev/null | head -40Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check for sample AFM files that might contain auth configs
find . -type f -name "*.afm*" -o -name "*.md" | xargs grep -l "authentication" 2>/dev/null | head -5Repository: wso2/reference-implementations-afm
Length of output: 367
🏁 Script executed:
# Look at sample_agent.afm.md if it exists
find . -name "sample_agent.afm.md" -exec cat {} \;Repository: wso2/reference-implementations-afm
Length of output: 3252
🏁 Script executed:
# Check requiredAuthFields and allowedAuthFields implementations
sed -n '380,410p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 872
🏁 Script executed:
# Check what fields are actually required for each auth type
rg -n "requiredAuthFields|allowedAuthFields" -A 20 ballerina-interpreter/parser.bal | head -80Repository: wso2/reference-implementations-afm
Length of output: 1945
🏁 Script executed:
# Check how auth records are created from YAML to understand null-field behavior
sed -n '200,250p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 1898
🏁 Script executed:
# Look for the YAML parsing logic that creates ClientAuthentication records
rg -n "ClientAuthentication\|fromYaml\|fromJson" ballerina-interpreter/parser.bal -B2 -A5 | head -60Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Verify Ballerina record key behavior with null values - check related code
rg -n "\.keys()" ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 260
🏁 Script executed:
# Check how the YAML/JSON is parsed into ClientAuthentication records
rg -n "parsing\|Parser\|yaml\|json" ballerina-interpreter/parser.bal | head -30Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Look for the from/map conversion that creates ClientAuthentication
sed -n '1,100p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 3406
🏁 Script executed:
# Verify if there are any null value tests in the test files
grep -n "null\|:\s*\(\)" ballerina-interpreter/tests/main_test.bal | head -30Repository: wso2/reference-implementations-afm
Length of output: 1427
🏁 Script executed:
# Check how the deserialized auth records are actually used in the runtime
grep -n "auth\[" ballerina-interpreter/parser.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 131
🏁 Script executed:
# Look for any usage of auth fields after validation
rg -n "authentication\." ballerina-interpreter/ | head -30Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check agent.bal to see how auth is used at runtime
head -100 ballerina-interpreter/agent.balRepository: wso2/reference-implementations-afm
Length of output: 3340
🏁 Script executed:
# Check the full ClientAuthentication type definition and its usage
sed -n '49,60p' ballerina-interpreter/types.balRepository: wso2/reference-implementations-afm
Length of output: 274
🏁 Script executed:
# Look for how auth values are actually used at runtime
rg -n "transport\.authentication\|model\.authentication" ballerina-interpreter/ -A3 -B1 | head -40Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check the mapToHttpClientAuth function to see if it handles null values
rg -n "mapToHttpClientAuth" ballerina-interpreter/ -A20 | head -50Repository: wso2/reference-implementations-afm
Length of output: 4060
🏁 Script executed:
# Find the mapToHttpClientAuth function implementation
rg -n "function mapToHttpClientAuth" ballerina-interpreter/ -A30 | head -50Repository: wso2/reference-implementations-afm
Length of output: 2319
🏁 Script executed:
# Check specifically for how token/password fields are accessed
rg -n "auth\[\"token\"\]|auth\[\"password\"\]|auth\[\"username\"\]|auth\[\"client_" ballerina-interpreter/ -B2 -A2 | head -60Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check if there's any test with null auth fields
rg -n "token.*\(\)|password.*\(\)|null" ballerina-interpreter/tests/main_test.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 149
🏁 Script executed:
# Check if ClientAuthentication is actually a record with optional fields or how it's defined
sed -n '49,65p' ballerina-interpreter/types.balRepository: wso2/reference-implementations-afm
Length of output: 351
🏁 Script executed:
# Look at what fromJsonWithType does - check the imports and usage
head -20 ballerina-interpreter/parser.bal | grep importRepository: wso2/reference-implementations-afm
Length of output: 127
🏁 Script executed:
# Check if there's a test with explicitly null auth fields that should fail
grep -n "null\|:\s*\(\)" ballerina-interpreter/tests/main_test.bal | grep -i "auth\|bearer\|basic" | head -20Repository: wso2/reference-implementations-afm
Length of output: 1564
🏁 Script executed:
# Verify how the YAML is parsed and if null fields are even possible
sed -n '48,70p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 711
🏁 Script executed:
# Check if there are any Ballerina tests that actually set auth fields to null
rg -n "null|():" ballerina-interpreter/tests/main_test.bal | grep -i "authentication\|bearer\|basic\|oauth" -B2 -A2 | head -40Repository: wso2/reference-implementations-afm
Length of output: 2107
🏁 Script executed:
# Check the complete test for missing fields to understand validation behavior
sed -n '418,426p' ballerina-interpreter/tests/main_test.balRepository: wso2/reference-implementations-afm
Length of output: 402
🏁 Script executed:
# Search for any null-field handling in mapToHttpClientAuth or related functions
rg -n "is \(\)|!is \(\)" ballerina-interpreter/agent.bal | head -30Repository: wso2/reference-implementations-afm
Length of output: 534
🏁 Script executed:
# Verify if YAML deserialization includes null fields or skips them
python3 << 'PY'
# Check if Ballerina's YAML parsing would include null fields
# by examining how records handle null in the context
# The key question: does fromJsonWithType include null fields in the record?
# In Ballerina, when a YAML field is null, does it:
# 1. Omit the field from the record?
# 2. Include the field with null value?
# Let's check the test to see if we can infer behavior
test_code = """
ClientAuthentication auth = {'type: "api-key"};
"""
print("Test shows creation without optional fields - suggests fields are optional")
print("This implies missing fields don't cause issues, only present ones with wrong values")
PYRepository: wso2/reference-implementations-afm
Length of output: 237
🏁 Script executed:
# Check if there's documentation or examples showing how Ballerina handles null in YAML->record conversion
rg -n "null|null:" ballerina-interpreter/tests/ | head -20Repository: wso2/reference-implementations-afm
Length of output: 191
🏁 Script executed:
# Look at how the error message "api_key not found" is generated
rg -n "not found in" ballerina-interpreter/Repository: wso2/reference-implementations-afm
Length of output: 947
🏁 Script executed:
# Check the getApiKey function that generates that error
rg -n "function getApiKey" ballerina-interpreter/ -A10Repository: wso2/reference-implementations-afm
Length of output: 923
🏁 Script executed:
# Look for any actual null field tests in validation tests
sed -n '418,453p' ballerina-interpreter/tests/main_test.balRepository: wso2/reference-implementations-afm
Length of output: 1299
🏁 Script executed:
# Search for tests that might create auth with explicit null fields
rg -n "token.*null|\(\)" ballerina-interpreter/tests/ | grep -i "auth\|bearer\|basic" | head -20Repository: wso2/reference-implementations-afm
Length of output: 2428
🏁 Script executed:
# Check if there's documentation about how null values behave in Ballerina record deserialization
rg -n "cloneWithType" ballerina-interpreter/ -B3 -A3 | head -40Repository: wso2/reference-implementations-afm
Length of output: 3101
🏁 Script executed:
# Look for how auth is actually constructed from YAML in edge cases
sed -n '320,340p' ballerina-interpreter/agent.balRepository: wso2/reference-implementations-afm
Length of output: 751
🏁 Script executed:
# Check the complete getAuthTokenOrApiKey to understand null handling at usage time
sed -n '321,335p' ballerina-interpreter/agent.balRepository: wso2/reference-implementations-afm
Length of output: 608
🏁 Script executed:
# Check if there are actual tests showing null field behavior in AFM YAML
find . -type f -name "*.md" -o -name "*.afm" | xargs grep -l "null" 2>/dev/null | head -5Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Look at how the OAuth2Config and other config types handle null/missing fields
rg -n "type OAuth2Config|type BearerTokenConfig|type CredentialsConfig" ballerina-interpreter/Repository: wso2/reference-implementations-afm
Length of output: 142
🏁 Script executed:
# Check if there's any Ballerina documentation or test showing null field behavior in record deserialization
sed -n '1,50p' ballerina-interpreter/tests/agent_test.balRepository: wso2/reference-implementations-afm
Length of output: 1670
🏁 Script executed:
# Look for any test that actually creates an auth record with explicit null value
rg -n ":\s*\(\)\|null" ballerina-interpreter/tests/agent_test.bal | head -20Repository: wso2/reference-implementations-afm
Length of output: 60
🏁 Script executed:
# Check if Ballerina's YAML module preserves null values or omits them
rg -n "data.yaml|fromJson" ballerina-interpreter/parser.bal -B2 -A2 | head -30Repository: wso2/reference-implementations-afm
Length of output: 399
🏁 Script executed:
# Check the OAuth2Config and other config types to understand field definitions
sed -n '422,440p' ballerina-interpreter/agent.balRepository: wso2/reference-implementations-afm
Length of output: 712
🏁 Script executed:
# Create a test to understand how Ballerina handles null in record deserialization
python3 << 'PY'
import json
# Simulate what happens in Ballerina:
# 1. YAML with null field: {type: "bearer", token: null}
# 2. Deserialize to ClientAuthentication record
# In Ballerina, when deserializing to a typed record:
# - If a field is declared as string, null cannot be assigned
# - The deserialization would fail OR omit the null field
# - Extra fields go into ...rest as map<anydata>
# The key question: does auth.keys() include null-valued fields?
# In Ballerina maps, a key with null value IS included in keys()
# But ClientAuthentication is a record, not a raw map
# The validation uses:
# 1. auth.keys() - gets keys from the record's dynamic fields
# 2. indexOf() - checks key presence only
# The real issue: if auth has {token: ()} (null), keys() includes "token"
# But then auth.get("token") returns null
print("In Ballerina:")
print("- auth.keys() includes keys even if value is null")
print("- The parser validation only checks key presence")
print("- Runtime validation auth.get(key).ensureType() would fail on null")
print("")
print("The concern is valid: parse-time validation allows null values to pass")
print("Runtime would catch it, but parse-time should be stricter")
PYRepository: wso2/reference-implementations-afm
Length of output: 388
🏁 Script executed:
# Look for how fields are actually used at parse time vs runtime
rg -n "ensureType" ballerina-interpreter/ -B2 -A2Repository: wso2/reference-implementations-afm
Length of output: 2939
🏁 Script executed:
# Check the complete validation flow - from YAML to validation
sed -n '27,35p' ballerina-interpreter/parser.balRepository: wso2/reference-implementations-afm
Length of output: 353
Required-field validation currently checks key presence, not non-nil values.
At lines 364–369 and 429–434, provided is derived from auth.keys(), allowing required fields that are explicitly set to null to pass parse-time validation. The runtime validation via ensureType() in getAuthTokenOrApiKey() (line 334) would later reject null values, but parse-time validation should catch this earlier for fail-fast feedback.
Add filtering to exclude null-valued fields from the required-field check:
Proposed approach
string[] provided = from string key in auth.keys() where key != "type" select key;
+ string[] providedWithValue = [];
+ foreach string key in provided {
+ if auth[key] !is () {
+ providedWithValue.push(key);
+ }
+ }
foreach string required in requiredAuthFields(authType) {
- if provided.indexOf(required) is () {
+ if providedWithValue.indexOf(required) is () {
return error(string `type '${authType}' requires '${required}' field`);
}
}Apply the same pattern at line 429–434 for OAuth2 grant validation.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| string[] provided = from string key in auth.keys() where key != "type" select key; | |
| foreach string required in requiredAuthFields(authType) { | |
| if provided.indexOf(required) is () { | |
| return error(string `type '${authType}' requires '${required}' field`); | |
| } | |
| string[] provided = from string key in auth.keys() where key != "type" select key; | |
| string[] providedWithValue = []; | |
| foreach string key in provided { | |
| if auth[key] !is () { | |
| providedWithValue.push(key); | |
| } | |
| } | |
| foreach string required in requiredAuthFields(authType) { | |
| if providedWithValue.indexOf(required) is () { | |
| return error(string `type '${authType}' requires '${required}' field`); | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballerina-interpreter/parser.bal` around lines 364 - 369, The required-field
validation in the auth parsing logic currently only checks if keys are present
in auth.keys(), but does not validate that those fields have non-null values.
This allows null-valued fields to pass parse-time validation even though they
would be rejected later at runtime by ensureType() in getAuthTokenOrApiKey().
Modify the provided array derivation to filter out keys whose corresponding
values are null, ensuring required fields both exist and have non-null values.
Apply this same filtering pattern at both validation locations: in the
requiredAuthFields loop check around lines 364-369 and in the OAuth2 grant
validation around lines 429-434.
Source: Coding guidelines
| "langchain-anthropic>=1.3.1", | ||
| "mcp>=1.26.0", | ||
| "langchain-mcp-adapters>=0.2.1", | ||
| "pyjwt[crypto]>=2.10.0", |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check available PyJWT versions on PyPI
curl -s https://pypi.org/pypi/PyJWT/json | jq -r '.releases | keys | .[]' | sort -V | tail -10
# Check for security advisories
gh api graphql -f query='
{
securityVulnerabilities(first: 5, ecosystem: PIP, package: "pyjwt") {
nodes {
advisory {
summary
severity
publishedAt
}
vulnerableVersionRange
firstPatchedVersion {
identifier
}
}
}
}'Repository: wso2/reference-implementations-afm
Length of output: 321
🌐 Web query:
PyJWT 2.10.0 security vulnerabilities CVE advisories
💡 Result:
PyJWT version 2.10.0 was introduced with at least one confirmed security vulnerability and is affected by several subsequent advisories affecting the 2.x line [1][2][3]. Security Vulnerabilities Affecting PyJWT 2.10.0: 1. CVE-2024-53861 (GHSA-75c5-xw7c-p5pm): This vulnerability involves incorrect string comparison when checking the "iss" (issuer) claim [4][2]. The bug was introduced in version 2.10.0 and allows partial matching of the issuer field (e.g., "acb" could be accepted for "abc") [4][2]. This was patched in version 2.10.1 [4][5]. 2. Vulnerabilities Fixed in 2.13.0 (Also affecting 2.10.0): Because version 2.10.0 is an older release, it is affected by a range of vulnerabilities discovered and patched in later versions, including those addressed in the 2.13.0 release (May 2026) [3][6]: - Algorithm Confusion (CVE-2026-48526 / GHSA-xgmm-8j9v-c9wx): A vulnerability where a public-key JWK accepted as an HMAC secret could enable forged HS256 tokens [7][6]. - Algorithm Allow-list Bypass (GHSA-jq35-7prp-9v3f): An issue where the caller's algorithm allow-list could be bypassed when decoding with a PyJWK or PyJWKClient key [6][8]. - SSRF in PyJWKClient (GHSA-993g-76c3-p5m4): Lack of URI scheme validation in PyJWKClient allowed unauthorized access to local files or other schemes via urllib (e.g., file://, ftp://) [6][9]. - Denial of Service (GHSA-w7vc-732c-9m39): An unauthenticated DoS amplifier related to base64 decoding of unused payload segments when b64=false [6][10]. Recommendation: Users of PyJWT 2.10.0 are strongly advised to upgrade to the latest version (2.13.0 or newer) to mitigate these and other security risks [3][6][11]. There are no known workarounds for several of these vulnerabilities [2]. Always review the official changelog and GitHub security advisories for the most comprehensive list of patches [11][10].
Citations:
- 1: https://pypi.org/project/PyJWT/2.10.0/
- 2: https://osv.dev/vulnerability/CVE-2024-53861
- 3: https://security.snyk.io/package/pip/pyjwt
- 4: GHSA-75c5-xw7c-p5pm
- 5: https://pyjwt.readthedocs.io/en/2.10.1/changelog.html
- 6: https://github.com/jpadilla/pyjwt/releases/tag/2.13.0
- 7: GHSA-xgmm-8j9v-c9wx
- 8: GHSA-jq35-7prp-9v3f
- 9: GHSA-993g-76c3-p5m4
- 10: https://github.com/jpadilla/pyjwt/blob/master/CHANGELOG.rst
- 11: https://pyjwt.readthedocs.io/en/stable/changelog.html
Upgrade PyJWT dependency to address security advisories.
Version 2.10.0 is affected by multiple security advisories, including CVE-2024-53861 (introduced in that version and patched in 2.10.1) and several others addressed in later releases. Upgrade to version 2.13.0 or newer.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python-interpreter/packages/afm-langchain/pyproject.toml` at line 20, The
pyjwt[crypto] dependency in the pyproject.toml file is pinned to version 2.10.0,
which contains security vulnerabilities including CVE-2024-53861. Update the
version constraint for pyjwt[crypto] from >=2.10.0 to >=2.13.0 to address these
security advisories and ensure a secure version is used.
There was a problem hiding this comment.
Pull request overview
This PR standardizes AFM authentication schemas across the Python and Ballerina reference runtimes, adds parse-time validation for all supported auth types, and implements runtime support for jwt and oauth2 authentication (plus api-key header customization in Python, and explicit transport limitation in Ballerina).
Changes:
- Add
jwtsigning andoauth2grant handling to the Python MCP HTTP client auth wiring (and tests). - Introduce strict, type-aware authentication schema validation at parse/load time in both runtimes.
- Extend Ballerina transport auth mapping to support
jwt/oauth2and explicitly rejectapi-keyfor MCP/webhook transport.
Reviewed changes
Copilot reviewed 8 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| python-interpreter/uv.lock | Adds pyjwt[crypto] to the locked dependency set. |
| python-interpreter/packages/afm-langchain/pyproject.toml | Adds pyjwt[crypto]>=2.10.0 dependency. |
| python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py | Implements JwtAuth + OAuth2Auth and wires them via build_httpx_auth. |
| python-interpreter/packages/afm-langchain/tests/test_mcp.py | Adds unit tests for api-key header override, JWT signing, and OAuth2 token request behavior/caching. |
| python-interpreter/packages/afm-core/src/afm/models.py | Expands ClientAuthentication fields and adds strict per-type/grant validation. |
| python-interpreter/packages/afm-core/tests/test_parser.py | Adds validation tests ensuring auth errors are raised at parse time. |
| ballerina-interpreter/parser.bal | Adds parse-time validation for authentication blocks across metadata locations. |
| ballerina-interpreter/agent.bal | Maps oauth2/jwt auth into Ballerina http client auth configs; rejects api-key for transport. |
| ballerina-interpreter/tests/main_test.bal | Adds tests for new mapping + validation behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| string? keyId = jwtConfig?.key_id; | ||
| if keyId is string { | ||
| issuerConfig["keyId"] = keyId; | ||
| } | ||
| string? subject = jwtConfig?.subject; | ||
| if subject is string { | ||
| issuerConfig["username"] = subject; | ||
| } | ||
| map<json>? customClaims = jwtConfig?.custom_claims; |
| map<json> grantConfig = {tokenUrl: cfg?.token_url, clientId: cfg?.client_id, clientSecret: cfg?.client_secret}; | ||
| addScopes(grantConfig, cfg?.scopes); | ||
| return wrapOAuth2(grantConfig.cloneWithType(http:OAuth2ClientCredentialsGrantConfig)); | ||
| } | ||
| "password" => { | ||
| map<json> grantConfig = {tokenUrl: cfg?.token_url, username: cfg?.username, password: cfg?.password}; | ||
| addOptional(grantConfig, "clientId", cfg?.client_id); | ||
| addOptional(grantConfig, "clientSecret", cfg?.client_secret); | ||
| addScopes(grantConfig, cfg?.scopes); | ||
| return wrapOAuth2(grantConfig.cloneWithType(http:OAuth2PasswordGrantConfig)); | ||
| } | ||
| "refresh_token" => { | ||
| map<json> grantConfig = {refreshUrl: cfg?.refresh_url, refreshToken: cfg?.refresh_token, clientId: cfg?.client_id, clientSecret: cfg?.client_secret}; | ||
| addScopes(grantConfig, cfg?.scopes); | ||
| return wrapOAuth2(grantConfig.cloneWithType(http:OAuth2RefreshTokenGrantConfig)); | ||
| } | ||
| "jwt_bearer" => { | ||
| map<json> grantConfig = {tokenUrl: cfg?.token_url, assertion: cfg?.assertion}; | ||
| addOptional(grantConfig, "clientId", cfg?.client_id); | ||
| addOptional(grantConfig, "clientSecret", cfg?.client_secret); | ||
| addScopes(grantConfig, cfg?.scopes); | ||
| return wrapOAuth2(grantConfig.cloneWithType(http:OAuth2JwtBearerGrantConfig)); |
| self.issuer = issuer | ||
| self.audience = audience | ||
| self.signing_key = signing_key | ||
| self.algorithm = algorithm | ||
| self.key_id = key_id | ||
| self.subject = subject | ||
| self.custom_claims = custom_claims or {} | ||
| self.expiry_seconds = expiry_seconds | ||
|
|
||
| def _resolve_key(self) -> str: | ||
| if self.algorithm.upper() in _HMAC_JWT_ALGORITHMS: | ||
| return self.signing_key | ||
| try: | ||
| return Path(self.signing_key).read_text() | ||
| except OSError as e: | ||
| raise MCPAuthenticationError( | ||
| f"Could not read JWT signing key file '{self.signing_key}': {e}" | ||
| ) from e | ||
|
|
There was a problem hiding this comment.
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 (2)
python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py (2)
201-208: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate OAuth2 token response shape before storing it.
Line 207 can raise non-auth errors for malformed
expires_in, and non-object payloads can fail before the intendedMCPAuthenticationError.Proposed fix
- def _store_token(self, payload: dict) -> str: + def _store_token(self, payload: object) -> str: + if not isinstance(payload, dict): + raise MCPAuthenticationError("OAuth2 token response was not a JSON object") token = payload.get("access_token") - if not token: + if not isinstance(token, str) or not token: raise MCPAuthenticationError( "OAuth2 token response did not contain 'access_token'" ) - expires_in = float(payload.get("expires_in", 3600)) + try: + expires_in = float(payload.get("expires_in", 3600)) + except (TypeError, ValueError) as e: + raise MCPAuthenticationError( + "OAuth2 token response contained an invalid 'expires_in' value" + ) from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py` around lines 201 - 208, The _store_token method in MCPClient should validate the OAuth2 token payload shape before reading fields so malformed responses fail with MCPAuthenticationError instead of generic exceptions. In _store_token, first verify payload is a dict-like object and that expires_in is present and numeric before converting it; if not, raise MCPAuthenticationError with a clear message. Keep the existing access_token check and update the token expiry logic only after the payload has been validated.
75-84: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winEnforce the JWT algorithm allow-list in the signer.
JwtAuthstores and uses the configured algorithm directly. Add the same RS/HS allow-list here so direct adapter construction stays aligned with parse-time validation.Proposed fix
+_JWT_ALGORITHMS = {"RS256", "RS384", "RS512", "HS256", "HS384", "HS512"} + class JwtAuth(httpx.Auth): @@ - self.algorithm = algorithm + self.algorithm = algorithm.upper() + if self.algorithm not in _JWT_ALGORITHMS: + raise MCPAuthenticationError( + f"JWT algorithm '{algorithm}' is not supported" + )Also applies to: 115-117
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py` around lines 75 - 84, JwtAuth currently accepts and stores the configured JWT algorithm directly, so direct construction can bypass the same RS/HS allow-list enforced during parsing. Update JwtAuth’s initialization and any related algorithm handling in the signer path to validate the algorithm against the existing allowed RS/HS set before assigning or using it, matching the behavior used by the parse-time validation logic.
🧹 Nitpick comments (1)
python-interpreter/packages/afm-langchain/tests/test_mcp.py (1)
250-312: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the OAuth2 password grant.
The OAuth2 request-shape tests skip the
passwordgrant even though this PR supports it. Add one focused case beside the other_token_request()tests.Proposed test addition
def test_oauth2_client_credentials_token_request(self): result = build_httpx_auth( ClientAuthentication( type="oauth2", grant_type="client_credentials", @@ assert data["grant_type"] == "client_credentials" assert data["scope"] == "read write" assert basic == ("id", "secret") + def test_oauth2_password_token_request(self): + result = build_httpx_auth( + ClientAuthentication( + type="oauth2", + grant_type="password", + token_url="https://auth.example.com/token", + username="user", + password="pass", + client_id="id", + client_secret="secret", + scopes=["read"], + ) + ) + assert isinstance(result, OAuth2Auth) + url, data, basic = result._token_request() + assert url == "https://auth.example.com/token" + assert data["grant_type"] == "password" + assert data["username"] == "user" + assert data["password"] == "pass" + assert data["scope"] == "read" + assert basic == ("id", "secret") + def test_oauth2_refresh_token_uses_token_url(self):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/tests/test_mcp.py` around lines 250 - 312, The OAuth2 request-shape tests in test_mcp.py are missing coverage for the password grant even though build_httpx_auth and OAuth2Auth now support it. Add a focused test alongside test_oauth2_client_credentials_token_request, test_oauth2_refresh_token_uses_token_url, and test_oauth2_jwt_bearer_token_request that builds a ClientAuthentication with type="oauth2" and grant_type="password", then asserts OAuth2Auth is returned and _token_request() produces the expected token URL and password-grant payload fields, including username/password and any required client auth behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py`:
- Around line 145-156: OAuth2Auth currently accepts any credential_bearer value
and silently treats non-"post_body" values as header auth, which bypasses the
schema contract. Update the OAuth2Auth initializer and related validation/path
that handles credential_bearer so only the supported values are accepted, and
reject unsupported values with a clear error before they can be used. Keep the
check close to the constructor/normalization logic so programmatic instantiation
and schema-based usage stay consistent.
- Around line 102-112: Prevent custom_claims from overwriting registered JWT
claims in the token builder. In the claims assembly inside the MCP JWT creation
logic, keep iss, iat, nbf, exp, aud, and sub set only from the configured fields
on the class, and merge only non-reserved entries from custom_claims so those
standard claims cannot be replaced. Update the claim-building code path in the
method that constructs the JWT payload to filter or reject reserved keys before
applying custom_claims.
---
Outside diff comments:
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py`:
- Around line 201-208: The _store_token method in MCPClient should validate the
OAuth2 token payload shape before reading fields so malformed responses fail
with MCPAuthenticationError instead of generic exceptions. In _store_token,
first verify payload is a dict-like object and that expires_in is present and
numeric before converting it; if not, raise MCPAuthenticationError with a clear
message. Keep the existing access_token check and update the token expiry logic
only after the payload has been validated.
- Around line 75-84: JwtAuth currently accepts and stores the configured JWT
algorithm directly, so direct construction can bypass the same RS/HS allow-list
enforced during parsing. Update JwtAuth’s initialization and any related
algorithm handling in the signer path to validate the algorithm against the
existing allowed RS/HS set before assigning or using it, matching the behavior
used by the parse-time validation logic.
---
Nitpick comments:
In `@python-interpreter/packages/afm-langchain/tests/test_mcp.py`:
- Around line 250-312: The OAuth2 request-shape tests in test_mcp.py are missing
coverage for the password grant even though build_httpx_auth and OAuth2Auth now
support it. Add a focused test alongside
test_oauth2_client_credentials_token_request,
test_oauth2_refresh_token_uses_token_url, and
test_oauth2_jwt_bearer_token_request that builds a ClientAuthentication with
type="oauth2" and grant_type="password", then asserts OAuth2Auth is returned and
_token_request() produces the expected token URL and password-grant payload
fields, including username/password and any required client auth behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2c8f1496-12c6-400f-beb7-d7da84e9b8a7
📒 Files selected for processing (7)
ballerina-interpreter/agent.balballerina-interpreter/parser.balballerina-interpreter/tests/agent_test.balpython-interpreter/packages/afm-core/src/afm/models.pypython-interpreter/packages/afm-core/tests/test_parser.pypython-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.pypython-interpreter/packages/afm-langchain/tests/test_mcp.py
🚧 Files skipped from review as they are similar to previous changes (4)
- ballerina-interpreter/parser.bal
- ballerina-interpreter/agent.bal
- python-interpreter/packages/afm-core/src/afm/models.py
- python-interpreter/packages/afm-core/tests/test_parser.py
| claims: dict = { | ||
| "iss": self.issuer, | ||
| "iat": now, | ||
| "nbf": now, | ||
| "exp": now + self.expiry_seconds, | ||
| } | ||
| if self.audience is not None: | ||
| claims["aud"] = self.audience | ||
| if self.subject is not None: | ||
| claims["sub"] = self.subject | ||
| claims.update(self.custom_claims) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Prevent custom claims from replacing registered JWT claims.
Line 112 lets custom_claims override iss, iat, nbf, exp, aud, or sub, which can make the generated token differ from the configured auth metadata.
Proposed fix
if self.audience is not None:
claims["aud"] = self.audience
if self.subject is not None:
claims["sub"] = self.subject
+ reserved_claims = {"iss", "iat", "nbf", "exp", "aud", "sub"}
+ conflicting_claims = reserved_claims.intersection(self.custom_claims)
+ if conflicting_claims:
+ raise MCPAuthenticationError(
+ "JWT custom_claims cannot override registered claims: "
+ + ", ".join(sorted(conflicting_claims))
+ )
claims.update(self.custom_claims)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| claims: dict = { | |
| "iss": self.issuer, | |
| "iat": now, | |
| "nbf": now, | |
| "exp": now + self.expiry_seconds, | |
| } | |
| if self.audience is not None: | |
| claims["aud"] = self.audience | |
| if self.subject is not None: | |
| claims["sub"] = self.subject | |
| claims.update(self.custom_claims) | |
| claims: dict = { | |
| "iss": self.issuer, | |
| "iat": now, | |
| "nbf": now, | |
| "exp": now + self.expiry_seconds, | |
| } | |
| if self.audience is not None: | |
| claims["aud"] = self.audience | |
| if self.subject is not None: | |
| claims["sub"] = self.subject | |
| reserved_claims = {"iss", "iat", "nbf", "exp", "aud", "sub"} | |
| conflicting_claims = reserved_claims.intersection(self.custom_claims) | |
| if conflicting_claims: | |
| raise MCPAuthenticationError( | |
| "JWT custom_claims cannot override registered claims: " | |
| ", ".join(sorted(conflicting_claims)) | |
| ) | |
| claims.update(self.custom_claims) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py`
around lines 102 - 112, Prevent custom_claims from overwriting registered JWT
claims in the token builder. In the claims assembly inside the MCP JWT creation
logic, keep iss, iat, nbf, exp, aud, and sub set only from the configured fields
on the class, and merge only non-reserved entries from custom_claims so those
standard claims cannot be replaced. Update the claim-building code path in the
method that constructs the JWT payload to filter or reject reserved keys before
applying custom_claims.
| credential_bearer: str = "auth_header", | ||
| ) -> None: | ||
| self.grant_type = grant_type.lower() | ||
| self.token_url = token_url | ||
| self.client_id = client_id | ||
| self.client_secret = client_secret | ||
| self.username = username | ||
| self.password = password | ||
| self.refresh_token = refresh_token | ||
| self.assertion = assertion | ||
| self.scopes = scopes | ||
| self.credential_bearer = credential_bearer |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject unsupported credential_bearer values in the adapter.
Any value other than "post_body" currently falls through to header credentials. Validate this in OAuth2Auth so programmatic construction matches the schema contract.
Proposed fix
- self.credential_bearer = credential_bearer
+ self.credential_bearer = credential_bearer.lower()
+ if self.credential_bearer not in {"auth_header", "post_body"}:
+ raise MCPAuthenticationError(
+ f"oauth2 'credential_bearer' '{credential_bearer}' is not supported"
+ )Also applies to: 192-198
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/tools/mcp.py`
around lines 145 - 156, OAuth2Auth currently accepts any credential_bearer value
and silently treats non-"post_body" values as header auth, which bypasses the
schema contract. Update the OAuth2Auth initializer and related validation/path
that handles credential_bearer so only the supported values are accepted, and
reject unsupported values with a clear error before they can be used. Keep the
check close to the constructor/normalization logic so programmatic instantiation
and schema-based usage stay consistent.
- jwt: validate algorithm against RS/HS allow-list and reject 'none' - oauth2: use token_url for the refresh_token grant (drop refresh_url) - jwt: make audience optional per RFC 7519 - auth: allow x-* extension types (pass-through; runtime decides support) - oauth2: add credential_bearer (auth_header|post_body) for token-endpoint client auth
Thareesha98
force-pushed
the
feat/add-oauth2-and-jwt-supports
branch
from
0e026fb to
132f3d2
Compare
Purpose
The AFM spec recognizes five authentication types (
bearer,basic,api-key,jwt,oauth2) but never defines the required or optional fields for any of them — the spec onlystates that "additional fields are authentication-type specific." This leaves several gaps:
tokenon abearerblock is only discovered at request time.the same type.
jwtandoauth2were recognized but had no runtime support becausetheir field schemas were never agreed.
Resolves #35.
Goals
authenticationappears.jwt(runtime signing) andoauth2(grant flows) in bothreference runtimes, keeping them at full parity.
Approach
Where it applies. The schema is identical in all three places
authenticationcan appear:model.authentication,tools.mcp[].transport.authentication, andinterfaces[].subscription.authentication.Common rules (enforced at parse time, in both runtimes).
typeis required and matched case-insensitively against the recognized set; any othervalue is rejected with a clear error listing the supported types.
tokn:.(
oauth2is the documented exception — its allowed fields depend ongrant_type.)${env:VAR}substitution; secrets SHOULD use${env:...}.Field schemas
bearertokenbasicusername,passwordapi-keyapi_keyheader_name(defaultAuthorization)jwtissuer,audience,signing_keyalgorithm(defaultRS256),key_id,subject,custom_claims,expiry_seconds(default300)oauth2grant_type+ per-grant fieldsoauth2grants:grant_typeclient_credentialstoken_url,client_id,client_secretscopespasswordtoken_url,username,password,client_id,client_secretscopesrefresh_tokenrefresh_url,refresh_token,client_id,client_secretscopesjwt_bearertoken_url,assertionclient_id,client_secret,scopesImplementation (both runtimes use one open model + a parse-time validator, for parity and
minimal churn).
afm-core/afm-langchain):ClientAuthenticationcarries the typed fields and a@model_validatorenforcing the rules above;build_httpx_authwires each type to anhttpx.Auth.jwtsigns per request via PyJWT (HMAC usessigning_keyas the shared secret;asymmetric algorithms read it as a PEM key file).
oauth2performs a real token exchange(sync and async), caches the token until expiry, sends client credentials as HTTP Basic,
and uses the RFC 7523 URN for
jwt_bearer. No new runtime dependency for oauth2.parser.balvalidates everyauthenticationblock at parse time;agent.balmaps
jwt→http:JwtIssuerConfigandoauth2→the matchinghttp:OAuth2*GrantConfig, so theHTTP client performs signing / token exchange / refresh natively.
Resolution of the open questions. (1)
jwt= runtime signing; pre-signed tokens usebearer. (2)api-keygains an optionalheader_name. (3) unknowntypevalues are rejectedat parse time.
Known limitations (documented).
api-keyis not applied to Ballerina MCP/webhook transport — thehttpclient has noraw-header auth variant — so it returns a clear "not yet supported for transport" error
(it remains supported for Ballerina model providers, and everywhere in Python).
subscription.authentication(pre-existing;affects all auth types) — tracked as a follow-up.
User stories
authenticationblock has an unknown type, a missing required field, or a typo'd field — instead of a late,
cryptic request-time failure.
via
header_name.jwtto have the runtime mint and sign a JWT from a key andclaims, without pre-generating tokens.
oauth2(client_credentials / password / refresh_token /jwt_bearer) and have the runtime fetch, cache, and attach the access token automatically.
implementations.
Release note
Standardized AFM authentication field schemas with parse-time validation across both reference
Release note
Standardized AFM authentication field schemas with parse-time validation across both reference
runtimes. Added an optional
header_nameforapi-key, runtime JWT signing forjwt, andOAuth2 grant flows (
client_credentials,password,refresh_token,jwt_bearer) foroauth2. Unknown authentication types and unknown/missing fields are now rejected at load timewith clear, consistent error messages.
Documentation
N/A — reference implementation. The authentication field schemas standardized here are defined
in the proposal on #35.