Add Agent Identity#121
Open
kalaiyarasiganeshalingam wants to merge 13 commits intoballerina-platform:mainfrom
Open
Add Agent Identity#121kalaiyarasiganeshalingam wants to merge 13 commits intoballerina-platform:mainfrom
kalaiyarasiganeshalingam wants to merge 13 commits intoballerina-platform:mainfrom
Conversation
[Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions Add token validation Add token validation Add dependency toml file Add dependencies toml file Update test [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions [Automated] Update the native jar versions Update implementation Update code Undo the changes Undo the changes
kalaiyarasiganeshalingam
requested review from
MohamedSabthar and
shafreenAnfar
as code owners
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds OAuth2/OIDC auth and per-tool token lifecycle: new auth config types, token cache, PKCE support, token acquisition/validation, per-tool scopes and scope-aware validation, and plumbing to pass a token manager into agents, tool stores, and the compiler plugin. Changes
Sequence DiagramsequenceDiagram
participant Agent as Agent
participant TokenMgr as TokenManager (cache)
participant TokenFlow as Token Flow
participant Validator as Validator (JWKS/Cert/Introspect)
participant ToolExec as Tool Executor
Agent->>TokenMgr: initialize with auth config
Agent->>TokenFlow: request token for tool (toolName, scopes)
TokenFlow->>TokenMgr: check cached token for toolName
alt cached token valid
TokenMgr-->>TokenFlow: return token + scopes
else missing/expired
TokenFlow->>TokenFlow: generate PKCE (verifier/challenge)
TokenFlow->>TokenFlow: start auth flow -> obtain code
TokenFlow->>TokenFlow: exchange code for token
TokenFlow->>TokenMgr: store token + scopes
end
TokenFlow->>Validator: validate token (JWKS / cert / introspect)
Validator-->>TokenFlow: validation result
TokenFlow->>TokenFlow: compare cached scopes vs required scopes
alt scopes OK
TokenFlow-->>Agent: return token + scopes
Agent->>ToolExec: execute tool with token
else scopes mismatch
TokenFlow-->>Agent: return MismatchScopeError
end
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested Reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 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: 15
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
compiler-plugin/src/main/java/io/ballerina/stdlib/ai/plugin/AiSourceModifier.java (1)
200-212:⚠️ Potential issue | 🟠 MajorGuard against empty scopes when adding missing fields, and include scopes in the mapping constructor when present.
Scopes defaults to empty string
""when not provided in annotations. However,getMissingFields()unconditionally treats SCOPES as required and callsNodeParser.parseExpression(config.get(fieldName))(line 273) without checking for blank values, which will fail. Additionally,generateConfigMappingConstructor()(lines 205-212) never includes scopes in the generated mapping, so configured scopes are silently dropped.🔧 Proposed fix
@@ private String generateConfigMappingConstructor(ToolAnnotationConfig config, String openBraceSource, String closeBraceSource) { String name = config.name().replaceAll("\\R", " "); - return openBraceSource + String.format("name:%s,description:%s,parameters:%s", + String scopes = config.scopes(); + String scopesPart = (scopes != null && !scopes.isBlank()) ? ",scopes:" + scopes : ""; + return openBraceSource + String.format("name:%s,description:%s,parameters:%s%s", name, config.description() != null ? config.description().replaceAll("\\R", " ") : name, - config.parameterSchema()) + closeBraceSource; + config.parameterSchema(), + scopesPart) + closeBraceSource; } @@ - List<String> requiredFields = List.of(NAME_FIELD_NAME, DESCRIPTION_FIELD_NAME, PARAMETERS_FIELD_NAME, SCOPES); + List<String> requiredFields = new ArrayList<>( + List.of(NAME_FIELD_NAME, DESCRIPTION_FIELD_NAME, PARAMETERS_FIELD_NAME)); + if (config.scopes() != null && !config.scopes().isBlank()) { + requiredFields.add(SCOPES); + } List<MappingFieldNode> missingFields = new ArrayList<>(); for (String fieldName : requiredFields) { if (!existingFieldNames.contains(fieldName)) { + String fieldValue = config.get(fieldName); + if (fieldValue == null || fieldValue.isBlank()) { + continue; + } missingFields.add(NodeFactory.createSpecificFieldNode( null, NodeFactory.createIdentifierToken(fieldName), NodeFactory.createToken(COLON_TOKEN), - NodeParser.parseExpression(config.get(fieldName)) + NodeParser.parseExpression(fieldValue) )); } }Also applies to: 264-275
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@compiler-plugin/src/main/java/io/ballerina/stdlib/ai/plugin/AiSourceModifier.java` around lines 200 - 212, The getMissingFields logic should guard against empty/blank SCOPES before calling NodeParser.parseExpression so it doesn't attempt to parse an empty string: update the code that iterates missing fields (referencing getMissingFields, SCOPES and NodeParser.parseExpression) to check for null/blank (e.g., config.get(fieldName) == null or trim().isEmpty()) and skip parsing/adding that field when blank; and update both generateConfigMappingConstructor overloads (generateConfigMappingConstructor(ToolAnnotationConfig) and generateConfigMappingConstructor(ToolAnnotationConfig, String, String)) to include scopes in the produced mapping when config.scopes() is non-blank—normalize newlines to spaces like other fields and add a scopes:%s entry in the formatted mapping only when present.
🧹 Nitpick comments (9)
ballerina/tests/toolkit-mcp-test.bal (1)
194-194: Consider renaming variabletoolStoretotoolManagerfor consistency.The type was renamed to
ToolManagerbut the variable names still usetoolStore. While not blocking, aligning variable names would improve readability.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/tests/toolkit-mcp-test.bal` at line 194, Rename the local variable named toolStore to toolManager to match the type ToolManager and improve readability; update the instantiation line that currently reads "ToolManager toolStore = check new (mcpToolKit);" and any other references to toolStore in this scope to toolManager so the variable name aligns with the ToolManager type and the mcpToolKit initializer.ballerina/tests/tool-test.bal (1)
275-277: Test function name still referencesToolStore.The function
testInitializingToolStoreWithoutNoToolsshould be renamed totestInitializingToolManagerWithoutNoToolsfor consistency with the API rename.📝 Proposed fix
`@test`:Config -isolated function testInitializingToolStoreWithoutNoTools() returns error? { +isolated function testInitializingToolManagerWithoutNoTools() returns error? { ToolManager toolStore = check new (); test:assertEquals(toolStore.tools.toArray().length(), 0); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/tests/tool-test.bal` around lines 275 - 277, Rename the test function testInitializingToolStoreWithoutNoTools to testInitializingToolManagerWithoutNoTools to match the API rename from ToolStore to ToolManager; update the function declaration (isolated function name) and any references to this test within the test suite so the new name is used consistently (look for the function declaration and any invocation or annotation that references testInitializingToolStoreWithoutNoTools).ballerina/constants.bal (1)
48-48: Misleading constant name:AUTHN_HEADERis a URL path, not a header.This constant is used as a URL path segment (
baseUrl.concat(AUTHN_HEADER)in token.bal), but its name suggests it's an HTTP header. Consider renaming for clarity.♻️ Suggested rename
-const AUTHN_HEADER = "authn"; +const AUTHN_PATH = "authn";🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/constants.bal` at line 48, The constant AUTHN_HEADER is misnamed — it represents a URL path segment, not an HTTP header; rename the constant (for example to AUTHN_PATH or AUTHN_SEGMENT) in constants.bal and update all references (e.g., the baseUrl.concat(AUTHN_HEADER) usage in token.bal) to the new identifier so the name accurately reflects its purpose.ballerina/tool.bal (1)
273-293: Code duplication ingetToolConfigreturn paths.The two return paths (with and without scopes) duplicate the same field assignments. Consider consolidating.
♻️ Suggested simplification
string?|string[] scopes = config?.scopes; do { - if scopes !is () { - return { - name: check config?.name.ensureType(), - description: check config?.description.ensureType(), - parameters: check config?.parameters.ensureType(), - scopes: check config?.scopes.ensureType(), - caller: tool - }; - } else { - return { - name: check config?.name.ensureType(), - description: check config?.description.ensureType(), - parameters: check config?.parameters.ensureType(), - caller: tool - }; + ToolConfig toolConfig = { + name: check config?.name.ensureType(), + description: check config?.description.ensureType(), + parameters: check config?.parameters.ensureType(), + caller: tool + }; + if scopes !is () { + toolConfig.scopes = check config?.scopes.ensureType(); } + return toolConfig; } on fail error e {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/tool.bal` around lines 273 - 293, The two return branches in getToolConfig duplicate the same fields; instead construct the result once (e.g., a var/record like result with name: check config?.name.ensureType(), description: check config?.description.ensureType(), parameters: check config?.parameters.ensureType(), caller: tool) and then, only if scopes are present, add scopes: check config?.scopes.ensureType() before returning result; keep the existing on fail error e catch that returns error Error("Unable to register the function '" + getFunctionName(tool) + "' as agent tool", e).ballerina/token.bal (2)
102-103: Confusing mutation: overwritingexpires_inwith epoch time.The
Token.expires_infield typically represents seconds-until-expiry per OAuth2 spec, but here it's being overwritten withvalidateTokenResult.expwhich is epoch time. While this works becauseTokenCache.update()uses it as epoch time, this semantic overloading is confusing. Consider using a separate variable or adding a comment.♻️ Suggested clarification
ValidationResponse|error validateTokenResult = validateToken(auth, baseUrl, freshToken.access_token, freshToken.token_type); if validateTokenResult is error { log:printError("Token validation failed", 'error = validateTokenResult, agentId = auth.agentId, toolName = toolName); return error TokenValidationError("Token validation failed: ", detail = {cause: validateTokenResult}); } - freshToken.expires_in = validateTokenResult.exp; - freshToken.scope = validateTokenResult.scope; + // Override expires_in with epoch expiry time from validation response + // TokenCache.update() expects epoch time, not seconds-until-expiry + freshToken.expires_in = validateTokenResult.exp; + freshToken.scope = validateTokenResult.scope;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/token.bal` around lines 102 - 103, The code overwrites Token.expires_in with validateTokenResult.exp (epoch time), which conflates OAuth2's seconds-until-expiry with an epoch timestamp and is confusing; change the mutation in the freshToken assignment so that expires_in retains the original semantics and instead set a new field or local variable (e.g., tokenExpiryEpoch or expires_at) from validateTokenResult.exp, and update any consumers such as TokenCache.update() to read the new field (or add a brief comment on freshToken/expires_in explaining the epoch usage if changing callers is infeasible) so the intent is explicit and unambiguous.
273-290: Redundant cast and unclear fallthrough error.Line 274 has a redundant cast
validation = <Jwt>validation;since the type is already narrowed. Also, line 289's error is reached when Jwt validation has neitherjwksConfignorcertFile, but the error message doesn't make this clear.♻️ Suggested improvement
if validation is Jwt { - validation = <Jwt>validation; record {string url;}? jwks = validation?.jwksConfig; if (jwks is record {string url;}) { return usingJwks(accessToken, jwks.url, auth); } else { string|crypto:PublicKey? certFile = validation?.certFile; if (certFile is string|crypto:PublicKey) { return usingCertificate(accessToken,certFile); } } } else { string? url = validation.introspectionUrl; string introspectUrl = url is string ? url : baseUrl.concat(INTROSPECT); return usingIntrospection(introspectUrl, accessToken, auth.agentId, validation.clientConfig, tokenTypeHint); } - return error TokenValidationError("No valid token validation configuration found"); + return error TokenValidationError("JWT validation requires either jwksConfig or certFile to be configured");🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/token.bal` around lines 273 - 290, Remove the redundant cast "validation = <Jwt>validation;" (type is already narrowed) and instead keep using the narrowed "validation" value; add an explicit error return when a Jwt has neither jwksConfig nor certFile so the fallthrough error is clearer—replace the final generic TokenValidationError with a specific message like "Jwt validation missing jwksConfig and certFile" (referencing Jwt, validation.jwksConfig, certFile, usingJwks, usingCertificate) and ensure usingIntrospection remains unchanged for non-Jwt branches.ballerina/agent.bal (2)
261-274: Inconsistent logging detail between authenticated and non-authenticated paths.When authentication is enabled, the success log omits
stepsandanswerthat are logged in the non-authenticated path. This reduces observability for authenticated agents. Consider logging the same details atprintInfolevel for authenticated paths, or using a consistent approach.♻️ Suggested improvement
if self.isAuthEnabled { log:printInfo("Agent execution completed successfully", executionId = executionId, - agentId = self.agentId + agentId = self.agentId, + steps = executionTrace.steps.toString(), + answer = answer ); } else {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/agent.bal` around lines 261 - 274, The authenticated success branch currently logs less detail than the non-authenticated branch: inside the lock when checking self.isAuthEnabled you call log:printInfo("Agent execution completed successfully", executionId = executionId, agentId = self.agentId) but omit steps and answer; update that branch (the lock / self.isAuthEnabled path) to include the same fields as the non-authenticated path (executionTrace.steps.toString() as steps and answer) so both log:printInfo and log:printDebug produce consistent observability (i.e., add steps and answer parameters to the log:printInfo call while preserving executionId and agentId).
175-177: Consider makingagentIdafinalfield for consistency.The
agentIdfield is set once during initialization and never modified afterward. Making itfinal(similar toisAuthEnabled) would be more consistent and eliminate the need for locks when accessing it.♻️ Suggested improvement
private final cache:Cache tokenManager = new (); -private string agentId = ""; -private boolean & readonly isAuthEnabled = false; +private final string agentId; +private final boolean isAuthEnabled;Then initialize in
init():if auth is AuthConfig { - self.isAuthEnabled = true; - self.agentId = auth.agentId.cloneReadOnly(); + self.isAuthEnabled = true; + self.agentId = auth.agentId; +} else { + self.isAuthEnabled = false; + self.agentId = ""; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/agent.bal` around lines 175 - 177, Make agentId immutable by changing its declaration to a final field and initialize it in the init() routine where it's currently set; specifically update the agentId field declaration (currently "private string agentId = \"\";") to be final and move its assignment into the init() method so you can remove any locking around agentId reads (references: agentId, init()). Ensure any code that previously mutated agentId is removed and compilation updated for final semantics.ballerina/agent-utils.bal (1)
154-160: Inconsistent error logging inreason()method.When
isAuthEnabledis false and the task is already completed, no error is logged before returningTaskCompletedError. Consider logging in both cases for consistent observability.♻️ Suggested improvement
if self.isCompleted { - if self.isAuthEnabled { - log:printError("Task is already completed. No more reasoning is needed.", - executionId = self.progress.executionId, - agentId = self.agentId - ); - } + log:printDebug("Task is already completed. No more reasoning is needed.", + executionId = self.progress.executionId, + agentId = self.isAuthEnabled ? self.agentId : () + ); return error TaskCompletedError("Task is already completed. No more reasoning is needed."); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/agent-utils.bal` around lines 154 - 160, In reason() the TaskCompletedError is only logged when self.isAuthEnabled is true; when false the function returns the error without logging which breaks observability—ensure you call log:printError with the same message and context (use self.progress.executionId and self.agentId) immediately before returning TaskCompletedError regardless of the value of self.isAuthEnabled so that both branches log consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ballerina/agent-utils.bal`:
- Around line 209-228: The code redundantly runs the same tool existence and
input merging/validation in toolManager.validateTool and then again inside
toolManager.execute; extract the shared validation/merging logic into a private
helper (e.g., a new function called validateAndMergeToolInput or similar) and
have validateTool and execute call that helper so execute assumes inputs are
pre-validated; update usages around validateTool, parsedOutput, and
toolManager.execute to pass the merged/validated result and remove duplicate
checks in execute to avoid the redundant work.
In `@ballerina/context.bal`:
- Around line 94-98: The getAccessToken function currently uses
self.toolAccessToken.get(toolName) which panics if the key is missing; change
getAccessToken to return an optional string (string?) and use a safe lookup on
self.toolAccessToken inside the lock (e.g., check existence or use a retrieval
form that returns an optional) so it returns nil when the token is absent, or
alternatively add a hasAccessToken(toolName) method that checks the map inside a
lock and keep getAccessToken unchanged; update callers to handle the optional or
to call hasAccessToken first.
In `@ballerina/error.bal`:
- Around line 71-75: Update the duplicate doc comment for TokenValidationError
so it correctly describes validation failures instead of acquisition;
specifically, change the documentation above the public type
TokenValidationError distinct ToolExecutionError to something like "Represents
an error that occurs when validating token for tool" (or similar wording
indicating token validation), leaving TokenAcquisitionError's comment unchanged.
- Around line 86-87: Rename the exported error type MissMatchScopeError to the
correct MismatchScopeError across the codebase: update the declaration public
type MissMatchScopeError distinct ToolInvalidInputError in ballerina/error.bal
to public type MismatchScopeError ..., then find and replace all references
(including uses and imports) in ballerina/token.bal and ballerina/tool.bal so
the identifier matches exactly (there are 5 occurrences total); ensure public
API signatures and any documentation/comments are updated accordingly and run
tests/compilation to catch any missed references.
In `@ballerina/token.bal`:
- Around line 132-133: The error message constructed in the return of
MissMatchScopeError concatenates "does not" and "match" without a space,
producing "does notmatch"; update the string concatenation in the return
statement (the MissMatchScopeError return in token.bal) to include a space
between "does not" and "match" so the message reads "...does not match the
existing token scopes: " + scope.
- Around line 351-362: The generatePKCE function currently swallows errors from
generateVerifier and returns nil; change its signature to return Pkce|error and
propagate the error from generateVerifier (i.e., when generateVerifier returns
an error, return that error instead of `return`), keep the successful path
building the {verifier, challenge} Pkce; then update the caller getFreshToken to
accept and handle the Pkce|error result (handle the error branch and bubble or
convert it appropriately). This touches generatePKCE, generateVerifier and
getFreshToken.
- Around line 197-202: The code creates a new http:Client for each call and
parses JSON without checking the response status; update the logic around
httpclient, post, and res to (1) reuse or accept an existing http:Client instead
of instantiating one per request (remove/replace the direct check
new(authorizeUrl) usage), (2) after httpclient->post(...) call check
res.statusCode() (or res.statusCode) and ensure it's a 2xx success (e.g.,
200–299) before calling res.getJsonPayload(), and (3) when the status is
non-2xx, read the body as text (res.getTextPayload() or similar) and return or
throw a descriptive error that includes the status code and body instead of
directly calling outputRes.cloneWithType(AuthResponse). Ensure AuthResponse
parsing only runs on successful responses.
- Around line 253-258: The code parses the response JSON immediately without
checking the HTTP status; update the sequence around httpclient, res and
getJsonPayload() to first validate the HTTP response status (e.g., ensure
res.statusCode() == 200 or within expected 2xx), and if not, return or throw a
descriptive error (including status and body) instead of calling
res.getJsonPayload(); keep the existing Token cloneWithType(Token) flow for
successful responses and reuse baseUrl, TOKEN, httpclient, res, getJsonPayload,
and Token identifiers when making the change.
- Around line 222-228: Validate the HTTP response status from the http:Response
(res) before calling getJsonPayload()—check res.statusCode (or res.status) and
return/throw a descriptive error when it's not a 200/expected status instead of
proceeding to parse; after parsing, avoid hardcoding the first authenticator by
not directly indexing authData[CODE]—inspect authData (e.g., ensure it is the
expected map/array), select the correct authenticator entry (by key or matching
criteria) and then return its CODE value via toString(); update the logic around
httpclient->post, res, getJsonPayload, obj["authData"], and authData[CODE]
accordingly.
- Around line 317-332: In usingIntrospection, ensure HTTP error status codes are
handled before parsing and validate the introspection response's required
boolean active field per RFC 7662: after httpclient->post and before
res.getJsonPayload, check res.statusCode and return an error for non-2xx
responses; then parse JSON, verify the presence and boolean type of the active
field (reject/return an error when active is false or missing) before calling
outputRes.cloneWithType(ValidationResponse) so inactive tokens are not accepted;
reference the usingIntrospection function, httpclient->post call,
res.getJsonPayload(), and ValidationResponse conversion when making the changes.
In `@ballerina/tool.bal`:
- Around line 110-112: The error message in the ToolNotFoundError return (the
return error ToolNotFoundError(...) call) contains an extra closing brace in the
string interpolation for self.tools.keys().toString() — remove the stray '}'
from the instruction string so the interpolation is
${self.tools.keys().toString()} and the message concatenation produces a valid
string.
- Around line 126-130: The debug log inside validateTool is misleadingly labeled
"Executing tool" — update the log:printDebug call in the validateTool function
to reflect validation (e.g., "Validating tool" or similar) and keep the same
structured fields (toolName, isMcpTool(self.isMcpTool(toolName)),
arguments=inputValues) so only the message string changes; ensure execute()
retains its own "Executing tool" log if present.
- Around line 53-56: Remove the dead boolean field hasAuthConfig from the
ToolManager class and delete any assignment to it in validateTool; update the
class declaration (remove "private boolean hasAuthConfig = false;") and remove
the line in validateTool that sets hasAuthConfig so there are no remaining
references. After removal, run the build/tests to ensure no other code depends
on hasAuthConfig and clean up any imports or comments that mentioned it.
- Line 22: The global declaration "configurable record {||} config = {};" is
unused and is shadowed by a local variable named "config" later; either delete
the global configurable declaration to avoid dead code and shadowing, or keep it
but add a clear comment above "configurable record {||} config = {};" describing
its intended purpose and why it differs from the local "config" variable (or
rename one of them to avoid shadowing). Update references accordingly so there's
no ambiguity between the global "config" and the local "config".
In `@ballerina/toolkit.bal`:
- Around line 543-545: The code risks a runtime panic by using bracket indexing
([SCOPES]) on a ToolConfig that may omit the optional scopes field; update the
ternary branches that set scopes (the expression referencing permittedTools,
FunctionTool, getToolConfig, SCOPES, and tool.name) to safely access the
optional field instead of direct bracket access — e.g., call getToolConfig(...)
and then use optional field access (or check existence) for .scopes and provide
a safe default (nil or empty list) when absent; also ensure the existing
check(...) only handles getToolConfig errors and does not force a panic when
scopes is missing.
---
Outside diff comments:
In
`@compiler-plugin/src/main/java/io/ballerina/stdlib/ai/plugin/AiSourceModifier.java`:
- Around line 200-212: The getMissingFields logic should guard against
empty/blank SCOPES before calling NodeParser.parseExpression so it doesn't
attempt to parse an empty string: update the code that iterates missing fields
(referencing getMissingFields, SCOPES and NodeParser.parseExpression) to check
for null/blank (e.g., config.get(fieldName) == null or trim().isEmpty()) and
skip parsing/adding that field when blank; and update both
generateConfigMappingConstructor overloads
(generateConfigMappingConstructor(ToolAnnotationConfig) and
generateConfigMappingConstructor(ToolAnnotationConfig, String, String)) to
include scopes in the produced mapping when config.scopes() is
non-blank—normalize newlines to spaces like other fields and add a scopes:%s
entry in the formatted mapping only when present.
---
Nitpick comments:
In `@ballerina/agent-utils.bal`:
- Around line 154-160: In reason() the TaskCompletedError is only logged when
self.isAuthEnabled is true; when false the function returns the error without
logging which breaks observability—ensure you call log:printError with the same
message and context (use self.progress.executionId and self.agentId) immediately
before returning TaskCompletedError regardless of the value of
self.isAuthEnabled so that both branches log consistently.
In `@ballerina/agent.bal`:
- Around line 261-274: The authenticated success branch currently logs less
detail than the non-authenticated branch: inside the lock when checking
self.isAuthEnabled you call log:printInfo("Agent execution completed
successfully", executionId = executionId, agentId = self.agentId) but omit steps
and answer; update that branch (the lock / self.isAuthEnabled path) to include
the same fields as the non-authenticated path (executionTrace.steps.toString()
as steps and answer) so both log:printInfo and log:printDebug produce consistent
observability (i.e., add steps and answer parameters to the log:printInfo call
while preserving executionId and agentId).
- Around line 175-177: Make agentId immutable by changing its declaration to a
final field and initialize it in the init() routine where it's currently set;
specifically update the agentId field declaration (currently "private string
agentId = \"\";") to be final and move its assignment into the init() method so
you can remove any locking around agentId reads (references: agentId, init()).
Ensure any code that previously mutated agentId is removed and compilation
updated for final semantics.
In `@ballerina/constants.bal`:
- Line 48: The constant AUTHN_HEADER is misnamed — it represents a URL path
segment, not an HTTP header; rename the constant (for example to AUTHN_PATH or
AUTHN_SEGMENT) in constants.bal and update all references (e.g., the
baseUrl.concat(AUTHN_HEADER) usage in token.bal) to the new identifier so the
name accurately reflects its purpose.
In `@ballerina/tests/tool-test.bal`:
- Around line 275-277: Rename the test function
testInitializingToolStoreWithoutNoTools to
testInitializingToolManagerWithoutNoTools to match the API rename from ToolStore
to ToolManager; update the function declaration (isolated function name) and any
references to this test within the test suite so the new name is used
consistently (look for the function declaration and any invocation or annotation
that references testInitializingToolStoreWithoutNoTools).
In `@ballerina/tests/toolkit-mcp-test.bal`:
- Line 194: Rename the local variable named toolStore to toolManager to match
the type ToolManager and improve readability; update the instantiation line that
currently reads "ToolManager toolStore = check new (mcpToolKit);" and any other
references to toolStore in this scope to toolManager so the variable name aligns
with the ToolManager type and the mcpToolKit initializer.
In `@ballerina/token.bal`:
- Around line 102-103: The code overwrites Token.expires_in with
validateTokenResult.exp (epoch time), which conflates OAuth2's
seconds-until-expiry with an epoch timestamp and is confusing; change the
mutation in the freshToken assignment so that expires_in retains the original
semantics and instead set a new field or local variable (e.g., tokenExpiryEpoch
or expires_at) from validateTokenResult.exp, and update any consumers such as
TokenCache.update() to read the new field (or add a brief comment on
freshToken/expires_in explaining the epoch usage if changing callers is
infeasible) so the intent is explicit and unambiguous.
- Around line 273-290: Remove the redundant cast "validation = <Jwt>validation;"
(type is already narrowed) and instead keep using the narrowed "validation"
value; add an explicit error return when a Jwt has neither jwksConfig nor
certFile so the fallthrough error is clearer—replace the final generic
TokenValidationError with a specific message like "Jwt validation missing
jwksConfig and certFile" (referencing Jwt, validation.jwksConfig, certFile,
usingJwks, usingCertificate) and ensure usingIntrospection remains unchanged for
non-Jwt branches.
In `@ballerina/tool.bal`:
- Around line 273-293: The two return branches in getToolConfig duplicate the
same fields; instead construct the result once (e.g., a var/record like result
with name: check config?.name.ensureType(), description: check
config?.description.ensureType(), parameters: check
config?.parameters.ensureType(), caller: tool) and then, only if scopes are
present, add scopes: check config?.scopes.ensureType() before returning result;
keep the existing on fail error e catch that returns error Error("Unable to
register the function '" + getFunctionName(tool) + "' as agent tool", e).
There was a problem hiding this comment.
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)
ballerina/Dependencies.toml (1)
1-4:⚠️ Potential issue | 🟡 MinorRemove unused dependency:
lang.stringis declared in Dependencies.toml but not imported anywhere in the codebase.This file is auto-generated and should not be manually modified. The new dependencies added to the
aipackage have been verified:
- ✓
cacheis imported in 5 files (tool.bal, token.bal, function-call-agent.bal, agent.bal, agent-utils.bal)- ✓
cryptois imported in 2 files (token.bal, agent.bal)- ✓
jwtis imported in token.bal- ✓
randomis imported in token.bal- ✗
lang.stringis NOT imported anywhere and should be removed from Dependencies.tomlEnsure
lang.stringis removed or, if intentionally added, add the corresponding import in the source files.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballerina/Dependencies.toml` around lines 1 - 4, Dependencies.toml currently lists the unused dependency lang.string for the ai package; remove the lang.string entry from Dependencies.toml (or if it was intentionally required, add the corresponding import lang.string statement in the source file(s) that need it) so the declared dependencies match actual imports — look for the Dependencies.toml file and the lang.string dependency entry and either delete that entry or add import lang.string in the relevant ai package source files.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ballerina/agent-utils.bal`:
- Around line 209-225: When validateTool returns an Error (validateRes) you must
close the observe span before returning/continuing: in the validateRes is Error
branch, after logging the error (log:printError) and before setting
observation/executionResult, call the agent's observe span close method (e.g.,
self.progress.observeSpan.end() or the appropriate close/end method on
self.progress's span) so the trace is ended and error attribution is recorded;
ensure the call uses the actual span property on self.progress and then proceed
to set observation and executionResult as currently done.
---
Outside diff comments:
In `@ballerina/Dependencies.toml`:
- Around line 1-4: Dependencies.toml currently lists the unused dependency
lang.string for the ai package; remove the lang.string entry from
Dependencies.toml (or if it was intentionally required, add the corresponding
import lang.string statement in the source file(s) that need it) so the declared
dependencies match actual imports — look for the Dependencies.toml file and the
lang.string dependency entry and either delete that entry or add import
lang.string in the relevant ai package source files.
---
Duplicate comments:
In `@ballerina/agent-utils.bal`:
- Around line 209-211: The code currently calls
ToolManager.validateTool(parsedOutput, ...) before calling
ToolManager.execute(...), which duplicates validation if execute already
re-validates; inspect the ToolManager implementation (methods validateTool and
execute) to confirm whether execute performs the same existence/input checks,
and then remove the duplicate: either remove the pre-call to validateTool and
rely on execute to validate internally, or refactor execute to accept a
skipValidation flag and call validateTool only when needed; update callers (the
current call site using parsedOutput, self.agent.auth, self.agent.tokenManager,
self.progress.context, isMcpTool) to use the new behavior so validation runs
exactly once and no redundant validation logic remains.
compiler-plugin/src/main/java/io/ballerina/stdlib/ai/plugin/ToolAnnotationAnalysisTask.java
Show resolved
Hide resolved
compiler-plugin/src/main/java/io/ballerina/stdlib/ai/plugin/AiSourceModifier.java
Outdated
Show resolved
Hide resolved
| # Function that should be called to execute the tool | ||
| isolated function caller; | ||
| # Scopes required to invoke this tool | ||
| string|string[] scopes?; |
Member
There was a problem hiding this comment.
why do we need scopes here?
There was a problem hiding this comment.
This stores all information required to execute the tool in subsequent processes, including its scope.
| string agentId = auth.agentId; | ||
| boolean needsRefresh = true; | ||
| string[] scopeInToken = []; | ||
| string & readonly tokenValue = EMPTY_STRING; |
Member
There was a problem hiding this comment.
shouldn't this be string? tokenValue = ();
There was a problem hiding this comment.
In the successful execution path, the access token is guaranteed to be non-null. In unsuccessful scenarios, the flow returns an error and does not continue with a null access token.
ballerina/token.bal
Outdated
| return scopeInToken; | ||
| } | ||
|
|
||
| isolated function validateToolScope(string[] cachedScopes, string toolName, string|string[] scopes, string agentId) |
Member
There was a problem hiding this comment.
Is it possible to pass a map<()> scopes. That way we can reduce time complexity.
There was a problem hiding this comment.
No, these are the scopes defined for a particular tool.
Member
There was a problem hiding this comment.
No, these are the scopes defined for a particular tool.
Correct. what I meant is that since Ballerina doesn’t have a set type, we can use a map to bring the lookup time down to O(n) in the loop below, instead of doing index search in array
| return httpclient->post(EMPTY_STRING, req); | ||
| } | ||
|
|
||
| isolated function generateVerifier(int length) returns string|error { |
Member
There was a problem hiding this comment.
It's seems like there is a recommended practice to generate PKCE verifier:
kalaiyarasiganeshalingam
force-pushed
the
main
branch
from
5629474 to
fa53231
Compare
Fix missing comments Remove scope Clean the code
kalaiyarasiganeshalingam
force-pushed
the
main
branch
from
9dc68b2 to
d98f8c8
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fixes:
Examples
Checklist
This PR adds agent identity and authentication features to the AI module, enabling agents to manage credentials, obtain and validate access tokens, and enforce scope-based tool access.
Key changes
Outcomes
Reviewer suggestion (follow-up)