internal(assistant): Add Assistant authentication#1136
Merged
Conversation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
e-fisher
force-pushed
the
internal/a2a-auth
branch
from
731e1bc to
7fa2a3f
Compare
Collaborator
Author
|
cursor review |
allansson
reviewed
Mar 30, 2026
Collaborator
allansson
left a comment
allansson
left a comment
There was a problem hiding this comment.
Everything seems to work as expected.
I had one instance where I got an error in the Grafana stack when trying to authorize but I haven't been able to reproduce it. It seemed to have something to do with the sign-in flow from the dialog.
Have some comments on the code.
| }, | ||
| } | ||
| } | ||
|
|
Collaborator
There was a problem hiding this comment.
For consistency with handlers/ai/a2a/index.ts being imported and initialized in handlers/ai/index.ts, maybe these should be re-exported from handlers/ai/a2a/preload.ts?
export * from "./a2a/preload.ts"| <Text size="2" color="gray"> | ||
| Sign in to Grafana Cloud to use the Grafana Assistant. | ||
| </Text> | ||
| <Button size="3" onClick={openProfileDialog}> |
Collaborator
There was a problem hiding this comment.
I should be possible to use the GrafanaCloudSignIn component to inline the sign in flow in the same dialog.
Comment on lines
+117
to
+229
| /** | ||
| * Starts a temporary local HTTP server to receive the OAuth callback. | ||
| * The server listens on a random port in the 54321-54399 range and | ||
| * shuts down after receiving the callback or when the signal is aborted. | ||
| */ | ||
| export function startCallbackServer( | ||
| signal: AbortSignal | ||
| ): Promise<{ port: number; waitForCallback: () => Promise<CallbackResult> }> { | ||
| return new Promise((resolve, reject) => { | ||
| const server = http.createServer({ keepAliveTimeout: 0 }) | ||
|
|
||
| function closeServer() { | ||
| server.close() | ||
| server.closeAllConnections() | ||
| } | ||
|
|
||
| const callbackPromise = new Promise<CallbackResult>( | ||
| (resolveCallback, rejectCallback) => { | ||
| signal.addEventListener( | ||
| 'abort', | ||
| () => { | ||
| closeServer() | ||
| rejectCallback(new Error('Auth flow aborted')) | ||
| }, | ||
| { once: true } | ||
| ) | ||
|
|
||
| server.on('request', (req, res) => | ||
| handleCallbackRequest( | ||
| req, | ||
| res, | ||
| closeServer, | ||
| resolveCallback, | ||
| rejectCallback | ||
| ) | ||
| ) | ||
| } | ||
| ) | ||
|
|
||
| listenOnAvailablePort(server, resolve, reject, callbackPromise) | ||
| }) | ||
| } | ||
|
|
||
| export function handleCallbackRequest( | ||
| req: http.IncomingMessage, | ||
| res: http.ServerResponse, | ||
| closeServer: () => void, | ||
| resolveCallback: (result: CallbackResult) => void, | ||
| rejectCallback: (error: Error) => void | ||
| ) { | ||
| const url = new URL(req.url ?? '/', 'http://localhost') | ||
|
|
||
| if (url.pathname !== '/callback') { | ||
| res.writeHead(404) | ||
| res.end() | ||
| return | ||
| } | ||
|
|
||
| const code = url.searchParams.get('code') | ||
| const state = url.searchParams.get('state') | ||
| const error = url.searchParams.get('error') | ||
| const endpoint = url.searchParams.get('endpoint') | ||
| const tenant = url.searchParams.get('tenant') | ||
| const email = url.searchParams.get('email') | ||
|
|
||
| const isSuccess = !error && code && state | ||
|
|
||
| res.writeHead(200, { 'Content-Type': 'text/html' }) | ||
| res.end(isSuccess ? successPage() : cancelledPage(), () => { | ||
| if (error) { | ||
| rejectCallback(new Error(`Authorization denied: ${error}`)) | ||
| } else if (code && state) { | ||
| resolveCallback({ code, state, endpoint, tenant, email }) | ||
| } else { | ||
| rejectCallback(new Error('Missing code or state in auth callback')) | ||
| } | ||
|
|
||
| closeServer() | ||
| }) | ||
| } | ||
|
|
||
| function listenOnAvailablePort( | ||
| server: http.Server, | ||
| resolve: (value: { | ||
| port: number | ||
| waitForCallback: () => Promise<CallbackResult> | ||
| }) => void, | ||
| reject: (error: Error) => void, | ||
| callbackPromise: Promise<CallbackResult> | ||
| ) { | ||
| const tryPort = (port: number) => { | ||
| if (port > CALLBACK_PORT_MAX) { | ||
| reject(new Error('No available port for OAuth callback server')) | ||
| return | ||
| } | ||
|
|
||
| server.removeAllListeners('listening') | ||
|
|
||
| server.once('error', () => { | ||
| tryPort(port + 1) | ||
| }) | ||
|
|
||
| server.listen(port, '127.0.0.1', () => { | ||
| server.removeAllListeners('error') | ||
| resolve({ | ||
| port, | ||
| waitForCallback: () => callbackPromise, | ||
| }) | ||
| }) | ||
| } | ||
|
|
||
| tryPort(CALLBACK_PORT_MIN) | ||
| } |
Collaborator
There was a problem hiding this comment.
Passing resolve and reject callbacks is just re-inventing promises. Took the liberty of rewriting it to returning promises instead:
Suggested change
| /** | |
| * Starts a temporary local HTTP server to receive the OAuth callback. | |
| * The server listens on a random port in the 54321-54399 range and | |
| * shuts down after receiving the callback or when the signal is aborted. | |
| */ | |
| export function startCallbackServer( | |
| signal: AbortSignal | |
| ): Promise<{ port: number; waitForCallback: () => Promise<CallbackResult> }> { | |
| return new Promise((resolve, reject) => { | |
| const server = http.createServer({ keepAliveTimeout: 0 }) | |
| function closeServer() { | |
| server.close() | |
| server.closeAllConnections() | |
| } | |
| const callbackPromise = new Promise<CallbackResult>( | |
| (resolveCallback, rejectCallback) => { | |
| signal.addEventListener( | |
| 'abort', | |
| () => { | |
| closeServer() | |
| rejectCallback(new Error('Auth flow aborted')) | |
| }, | |
| { once: true } | |
| ) | |
| server.on('request', (req, res) => | |
| handleCallbackRequest( | |
| req, | |
| res, | |
| closeServer, | |
| resolveCallback, | |
| rejectCallback | |
| ) | |
| ) | |
| } | |
| ) | |
| listenOnAvailablePort(server, resolve, reject, callbackPromise) | |
| }) | |
| } | |
| export function handleCallbackRequest( | |
| req: http.IncomingMessage, | |
| res: http.ServerResponse, | |
| closeServer: () => void, | |
| resolveCallback: (result: CallbackResult) => void, | |
| rejectCallback: (error: Error) => void | |
| ) { | |
| const url = new URL(req.url ?? '/', 'http://localhost') | |
| if (url.pathname !== '/callback') { | |
| res.writeHead(404) | |
| res.end() | |
| return | |
| } | |
| const code = url.searchParams.get('code') | |
| const state = url.searchParams.get('state') | |
| const error = url.searchParams.get('error') | |
| const endpoint = url.searchParams.get('endpoint') | |
| const tenant = url.searchParams.get('tenant') | |
| const email = url.searchParams.get('email') | |
| const isSuccess = !error && code && state | |
| res.writeHead(200, { 'Content-Type': 'text/html' }) | |
| res.end(isSuccess ? successPage() : cancelledPage(), () => { | |
| if (error) { | |
| rejectCallback(new Error(`Authorization denied: ${error}`)) | |
| } else if (code && state) { | |
| resolveCallback({ code, state, endpoint, tenant, email }) | |
| } else { | |
| rejectCallback(new Error('Missing code or state in auth callback')) | |
| } | |
| closeServer() | |
| }) | |
| } | |
| function listenOnAvailablePort( | |
| server: http.Server, | |
| resolve: (value: { | |
| port: number | |
| waitForCallback: () => Promise<CallbackResult> | |
| }) => void, | |
| reject: (error: Error) => void, | |
| callbackPromise: Promise<CallbackResult> | |
| ) { | |
| const tryPort = (port: number) => { | |
| if (port > CALLBACK_PORT_MAX) { | |
| reject(new Error('No available port for OAuth callback server')) | |
| return | |
| } | |
| server.removeAllListeners('listening') | |
| server.once('error', () => { | |
| tryPort(port + 1) | |
| }) | |
| server.listen(port, '127.0.0.1', () => { | |
| server.removeAllListeners('error') | |
| resolve({ | |
| port, | |
| waitForCallback: () => callbackPromise, | |
| }) | |
| }) | |
| } | |
| tryPort(CALLBACK_PORT_MIN) | |
| } | |
| /** | |
| * Starts a temporary local HTTP server to receive the OAuth callback. | |
| * The server listens on a random port in the 54321-54399 range and | |
| * shuts down after receiving the callback or when the signal is aborted. | |
| */ | |
| export async function startCallbackServer( | |
| signal: AbortSignal | |
| ): Promise<{ port: number; result: Promise<CallbackResult> }> { | |
| const server = http.createServer({ keepAliveTimeout: 0 }) | |
| function closeServer() { | |
| if (server.address() === null) { | |
| return | |
| } | |
| server.close() | |
| server.closeAllConnections() | |
| } | |
| // Not awaited because these will occur sometime in the future | |
| // and should be listened to by the caller of this function | |
| const aborted = rejectOnAbort(signal) | |
| const result = handleCallbackRequest(server) | |
| const port = await listenOnAvailablePort(server) | |
| return { | |
| port, | |
| result: Promise.race([result, aborted]).finally(closeServer), // Close server regardless of resolve, reject or abort | |
| } | |
| } | |
| function rejectOnAbort(signal: AbortSignal) { | |
| // Typed as `never` to guarantee it will never be resolved and | |
| // can be ruled out during type inference of `Promise.race`. | |
| const { promise, reject } = Promise.withResolvers<never>() | |
| function abort() { | |
| reject(new Error('Auth flow aborted')) | |
| } | |
| if (signal.aborted) { | |
| abort() | |
| return promise | |
| } | |
| signal.addEventListener('abort', abort) | |
| return promise | |
| } | |
| export function handleCallbackRequest(server: http.Server) { | |
| const { promise, resolve, reject } = Promise.withResolvers<CallbackResult>() | |
| server.on('request', (req, res) => { | |
| const url = new URL(req.url ?? '/', 'http://localhost') | |
| if (url.pathname !== '/callback') { | |
| res.writeHead(404) | |
| res.end() | |
| return | |
| } | |
| const code = url.searchParams.get('code') | |
| const state = url.searchParams.get('state') | |
| const error = url.searchParams.get('error') | |
| const endpoint = url.searchParams.get('endpoint') | |
| const tenant = url.searchParams.get('tenant') | |
| const email = url.searchParams.get('email') | |
| const isSuccess = !error && code && state | |
| res.writeHead(200, { 'Content-Type': 'text/html' }) | |
| res.end(isSuccess ? successPage() : cancelledPage(), () => { | |
| if (error) { | |
| reject(new Error(`Authorization denied: ${error}`)) | |
| } else if (code && state) { | |
| resolve({ code, state, endpoint, tenant, email }) | |
| } else { | |
| reject(new Error('Missing code or state in auth callback')) | |
| } | |
| }) | |
| }) | |
| return promise | |
| } | |
| function listenOnAvailablePort(server: http.Server) { | |
| const { promise, resolve, reject } = Promise.withResolvers<number>() | |
| const tryPort = (port: number) => { | |
| if (port > CALLBACK_PORT_MAX) { | |
| reject(new Error('No available port for OAuth callback server')) | |
| return | |
| } | |
| server.removeAllListeners('listening') | |
| server.once('error', () => { | |
| tryPort(port + 1) | |
| }) | |
| server.listen(port, '127.0.0.1', () => { | |
| server.removeAllListeners('error') | |
| resolve(port) | |
| }) | |
| } | |
| tryPort(CALLBACK_PORT_MIN) | |
| return promise | |
| } |
NOTE: I changed waitForCallback to just a Promise<CallbackResult>.
| return | ||
| } | ||
|
|
||
| const code = url.searchParams.get('code') |
Collaborator
There was a problem hiding this comment.
We're using openid-client when signing in to Grafana Cloud. Did you consider using it for this as well?
https://github.com/panva/openid-client?tab=readme-ov-file#authorization-code-flow
Collaborator
Author
There was a problem hiding this comment.
I've considered it, but went with manual approach because assistant's implementation doesn't follow oauth2 flow:
- Custom auth endpoint (
/a/grafana-assistant-app/cli/auth) with non-standard params (callback_port, scopes (plural), device_name) - Extra callback params
endpoint, tenant, email - Token exchange uses JSON POST instead of form-encoded POST
- Custom token response
Collaborator
Author
|
@allansson all great points, thank you for taking time to rewrite CallBack server 🙌 Ready for another review |
allansson
approved these changes
Apr 1, 2026
Collaborator
allansson
left a comment
allansson
left a comment
There was a problem hiding this comment.
Everything works great!
One note on the UX: it would be nice to have a more distinct "pending approval" state.
Knowing that there's more development and polish to come, I'm willing to leave this be for now though.
|
|
||
| export function useAssistantAuthStatus() { | ||
| const isProfileOpen = useStudioUIStore((s) => s.isProfileDialogOpen) | ||
| const wasOpen = useRef(false) |
Collaborator
There was a problem hiding this comment.
Nit:
Suggested change
| const wasOpen = useRef(false) | |
| const wasOpen = useRef(isProfileOpen) |
Comment on lines
+14
to
+22
| const isProfileOpen = useStudioUIStore((s) => s.isProfileDialogOpen) | ||
| const wasOpen = useRef(false) | ||
|
|
||
| useEffect(() => { | ||
| if (wasOpen.current && !isProfileOpen) { | ||
| void invalidateAssistantAuthStatus() | ||
| } | ||
| wasOpen.current = isProfileOpen | ||
| }, [isProfileOpen]) |
Collaborator
There was a problem hiding this comment.
Is this logic still applicable now that we don't use the profile dialog?
Collaborator
Author
|
I'll address last comments in the follow-up PR, thanks! |
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.
Description
This PR sets up the authentication layer only. A follow-up PR will wire the autocorrelation feature to use Grafana Assistant through this connection.
How the auth flow works
When the user clicks "Connect to Grafana Assistant", k6 Studio generates a PKCE challenge and starts a temporary HTTP server on localhost (port range 54321-54399) to receive the OAuth callback. The default browser opens to the Grafana Assistant CLI auth page where the user grants access. Grafana redirects back to the local server with an authorization code, which is exchanged for access and refresh tokens. Tokens are encrypted and stored per stack in the user data directory.
Known limitations
First-time Grafana Assistant users will see a "Terms acceptance required" page during OAuth instead of the consent screen. They need to open the Assistant chat in Grafana and send a message to accept terms, then retry. There's not much we can do about this from studio side, same limitation applies for Assistant CLI auth, will need to revisit this before removing feature toggle - attempt to fix it on the Assistant side.
How to Test
Also:
Disable Assistant feature toggle and verify existing OpenAI integration is unaffected and works as expected.
Checklist
Related PR(s)/Issue(s)
Note
Medium Risk
Introduces new auth/token-handling code (PKCE, localhost callback server, encrypted persistence) and Electron IPC handlers, which can impact security and login reliability if misconfigured.
Overview
Adds an end-to-end OAuth PKCE authentication layer to connect k6 Studio to Grafana Assistant, including a localhost callback server, code exchange, and per-stack encrypted token persistence.
Wires the new auth flow through Electron IPC (
AssistantAuthHandler) with cancel/sign-out support, exposes renderer-side helpers/hooks (useAssistantAuth*), and updates the autocorrelation intro UI to drive Grafana Cloud sign-in plus connect/disconnect/error states. Also centralizes Profile dialog open state inuseStudioUIStoreand enables thegrafana-assistantfeature by default in dev, with new unit tests for the auth exchange/callback logic.Written by Cursor Bugbot for commit 828add8. This will update automatically on new commits. Configure here.