Add authorization to agents using cotext token #1741
Conversation
jezekra1
changed the title
Summary of ChangesHello @jezekra1, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly overhauls the platform's authentication infrastructure, moving to a more secure RSA-based JWT system. It ensures that agent interactions are properly authenticated and authorized by integrating a new authentication backend, refining how tokens are managed and exchanged between services, and updating both the core SDK and user interface to support these changes. The update also includes general improvements to dependency handling and type consistency across the codebase. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and wide-ranging feature to add authorization to agents using context tokens. The changes span across the Python SDK, CLI, server, and the TypeScript UI, as well as configuration files and Helm charts. Key changes include the introduction of a PlatformAuthBackend for JWT-based authentication, a switch from symmetric (HS256) to asymmetric (RS256) JWT signing, and a substantial refactoring of how agent extensions handle request context and authentication. The overall approach is solid and well-implemented. However, I've identified a critical bug in the UI related to context token handling, a potential issue with query caching in React, and a configuration change in the Python SDK that could hide future bugs. Addressing these points will improve the robustness of this new authentication system.
| export async function createContextToken(createContextTokenParams: CreateContextTokenParams) { | ||
| console.log(createContextTokenParams) | ||
| const result = await agentstackClient.createContextToken(createContextTokenParams); | ||
| return result.token; |
There was a problem hiding this comment.
This function has two issues:
- A
console.logstatement is present, which appears to be for debugging and should be removed. - The function returns
result.token(a string), but the caller inAgentRunProvider.tsxexpects the fullContextTokenobject ({ token: string, expires_at: string }). This will cause a runtime error or incorrect behavior. The function should return the entireresultobject.
export async function createContextToken(createContextTokenParams: CreateContextTokenParams) {
const result = await agentstackClient.createContextToken(createContextTokenParams);
return result;
}| }: Props<UIGenericPart>) { | ||
| const { data: agentClient } = useQuery({ | ||
| queryKey: runKeys.client(providerId), | ||
| queryKey: runKeys.client(`${providerId}${Boolean(authToken)}`), |
There was a problem hiding this comment.
The queryKey for useBuildA2AClient includes Boolean(authToken). This is insufficient to detect changes in the authentication token. If the token is refreshed, Boolean(authToken) will remain true, and the query will not be re-triggered, causing the old agentClient with a stale token to be used. The query key should incorporate the token value itself to ensure the client is rebuilt when the token changes.
| queryKey: runKeys.client(`${providerId}${Boolean(authToken)}`), | |
| queryKey: runKeys.client(`${providerId}${authToken?.token}`), |
| ignore = ["tests/**", "examples/cli.py"] | ||
| venvPath = "." | ||
| venv = ".venv" | ||
| reportUnusedCallResult = "none" |
There was a problem hiding this comment.
Disabling reportUnusedCallResult globally can be risky. While it cleans up the need for _ = ... for functions where the return value is intentionally ignored, it also silences warnings for forgotten await statements on coroutines. This can lead to subtle and hard-to-debug concurrency issues. A safer approach is to keep this check enabled and explicitly ignore unused results where necessary (e.g., by assigning to _).
| token, _ = exchange_internal_jwt(header_token, self._configuration, audience=[audience]) | ||
| headers["authorization"] = f"Bearer {token}" | ||
| except Exception: | ||
| headers.pop("authorization", None) # forward header only if it's a valid context token |
There was a problem hiding this comment.
Forwarding the request without the Authorization header when the token exchange fails is a reasonable fallback. However, it might be beneficial to log a warning in this except block. This would make it easier to debug situations where agent calls are unexpectedly unauthenticated because the context token exchange failed for some reason.
Signed-off-by: Radek Ježek <[email protected]>
jezekra1
force-pushed
the
platform-auth-wip
branch
from
94e52de to
3db4e6c
Compare
Summary
Important changes:
audclaim.PLATFORM_AUTH__PUBLIC_URLenv for managed providers and validated.task/listortask/getmethods from a2a, which are not properly scoped per user. The a2a proxy is tracking task ownership and won't allow these methods if unauthorized.Todos:
Linked Issues
Ref: #1637
Documentation
If this PR adds new feature or changes existing. Make sure documentation is adjusted accordingly. If the docs is not needed, please explain why.