internal(assistant): Add assistant backend (#1154)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

Author: e-fisher

src/handlers/ai/a2a/assistantAuth.ts

@@ -1,4 +1,4 @@
-import { app, ipcMain, shell } from 'electron'
+import { app, BrowserWindow, ipcMain, shell } from 'electron'
 import log from 'electron-log/main'
 
 import { getProfileData } from '@/handlers/auth/fs'
@@ -10,17 +10,19 @@ import {
   generateState,
   startCallbackServer,
 } from '@/services/grafana/assistantAuth'
+import { browserWindowFromEvent } from '@/utils/electron'
 
+import { abortAllActiveAssistantSessions } from '../grafanaAssistantProvider'
 import { AssistantAuthHandler } from '../types'
 
+import { LOG_PREFIX } from './constants'
 import {
   clearAssistantTokens,
   hasAssistantTokens,
+  mapTokenResponse,
   saveAssistantTokens,
 } from './tokenStore'
 
-const PREFIX = '[GrafanaAssistant]'
-
 export type AssistantAuthResult =
   | { type: 'authenticated' }
   | { type: 'error'; error: string }
@@ -47,11 +49,23 @@ function isAllowedEndpoint(endpoint: string, stackUrl: string): boolean {
   }
 }
 
+function verificationCode(codeChallenge: string): string {
+  try {
+    const hex = Buffer.from(codeChallenge, 'base64url')
+      .subarray(0, 4)
+      .toString('hex')
+    return hex.slice(0, 4) + '-' + hex.slice(4)
+  } catch {
+    return '----'
+  }
+}
+
 let pendingAbortController: AbortController | null = null
 
 async function performSignIn(
   stackId: string,
-  stackUrl: string
+  stackUrl: string,
+  browserWindow: BrowserWindow
 ): Promise<AssistantAuthResult> {
   const abortController = new AbortController()
   pendingAbortController = abortController
@@ -65,22 +79,25 @@ async function performSignIn(
     const authUrl = buildAssistantAuthUrl(stackUrl, codeChallenge, state, port)
 
     log.info(
-      PREFIX,
+      LOG_PREFIX,
       'Initiating assistant auth for stack',
       stackId,
       'on port',
       port
     )
     void shell.openExternal(authUrl)
 
+    const code = verificationCode(codeChallenge)
+    browserWindow.webContents.send(AssistantAuthHandler.VerificationCode, code)
+
     const callback = await result
     app.focus({ steal: true })
 
     if (callback.state !== state) {
-      log.error(PREFIX, 'State mismatch in assistant auth callback')
+      log.error(LOG_PREFIX, 'State mismatch in assistant auth callback')
       return {
         type: 'error',
-        error: 'State mismatch — possible CSRF attack. Please try again.',
+        error: 'State mismatch, possible CSRF attack. Please try again.',
       }
     }
 
@@ -93,7 +110,7 @@ async function performSignIn(
 
     if (!isAllowedEndpoint(callback.endpoint, stackUrl)) {
       log.error(
-        PREFIX,
+        LOG_PREFIX,
         'Callback endpoint does not match expected stack URL:',
         callback.endpoint
       )
@@ -114,22 +131,18 @@ async function performSignIn(
       return { type: 'aborted' }
     }
 
-    await saveAssistantTokens(stackId, {
-      accessToken: tokenResponse.token,
-      refreshToken: tokenResponse.refresh_token,
-      apiEndpoint: tokenResponse.api_endpoint,
-      expiresAt: new Date(tokenResponse.expires_at).getTime(),
-      refreshExpiresAt: new Date(tokenResponse.refresh_expires_at).getTime(),
-    })
+    const tokens = mapTokenResponse(tokenResponse, tokenResponse.api_endpoint)
+
+    await saveAssistantTokens(stackId, tokens)
 
-    log.info(PREFIX, 'Assistant auth completed for stack', stackId)
+    log.info(LOG_PREFIX, 'Assistant auth completed for stack', stackId)
     return { type: 'authenticated' }
   } catch (error) {
     if (abortController.signal.aborted) {
       return { type: 'aborted' }
     }
 
-    log.error(PREFIX, 'Assistant auth failed:', error)
+    log.error(LOG_PREFIX, 'Assistant auth failed:', error)
     return {
       type: 'error',
       error: error instanceof Error ? error.message : 'Authentication failed',
@@ -144,7 +157,8 @@ async function performSignIn(
 export function initialize() {
   ipcMain.handle(
     AssistantAuthHandler.SignIn,
-    async (): Promise<AssistantAuthResult> => {
+    async (event): Promise<AssistantAuthResult> => {
+      const browserWindow = browserWindowFromEvent(event)
       const profile = await getProfileData()
       const stackId = profile.profiles.currentStack
 
@@ -170,7 +184,7 @@ export function initialize() {
         pendingAbortController.abort()
       }
 
-      return performSignIn(stackId, stack.url)
+      return performSignIn(stackId, stack.url, browserWindow)
     }
   )
 
@@ -207,8 +221,9 @@ export function initialize() {
     const stackId = profile.profiles.currentStack
 
     if (stackId) {
+      abortAllActiveAssistantSessions()
       await clearAssistantTokens(stackId)
-      log.info(PREFIX, 'Cleared assistant tokens for stack', stackId)
+      log.info(LOG_PREFIX, 'Cleared assistant tokens for stack', stackId)
     }
   })
 }

src/handlers/ai/a2a/cancelTask.ts

@@ -0,0 +1,32 @@
+import log from 'electron-log/main'
+
+import type { A2ASessionConfig } from './config'
+import { LOG_PREFIX } from './constants'
+import { buildA2AHeaders, safeResponseText } from './helpers'
+
+export async function sendTaskCancel(
+  config: A2ASessionConfig,
+  taskId: string
+): Promise<void> {
+  const body = {
+    jsonrpc: '2.0',
+    id: crypto.randomUUID(),
+    method: 'tasks/cancel',
+    params: { id: taskId },
+  }
+
+  const response = await fetch(`${config.baseUrl}/agents/${config.agentId}`, {
+    method: 'POST',
+    headers: buildA2AHeaders(config),
+    body: JSON.stringify(body),
+  })
+
+  if (!response.ok) {
+    const text = await safeResponseText(response)
+    log.error(
+      LOG_PREFIX,
+      `Failed to cancel task ${taskId} (${response.status}):`,
+      text
+    )
+  }
+}

src/handlers/ai/a2a/config.ts

@@ -0,0 +1,44 @@
+import { getProfileData } from '@/handlers/auth/fs'
+
+import { getValidAssistantTokens } from './tokenRefresh'
+
+export interface A2AConfig {
+  baseUrl: string
+  agentId: string
+  extensions: string
+  bearerToken: string
+}
+
+export type A2ASessionConfig = Omit<A2AConfig, 'extensions'>
+
+const AGENT_ID = 'grafana_assistant_k6_studio'
+const REMOTE_TOOL_EXTENSION =
+  'https://grafana.com/extensions/remote-tool-execution/v1'
+const TOKEN_STREAMING_EXTENSION =
+  'https://grafana.com/extensions/token-streaming/v1'
+
+export async function getA2AConfig(): Promise<A2AConfig> {
+  const profile = await getProfileData()
+  const stackId = profile.profiles.currentStack
+
+  if (!stackId) {
+    throw new Error(
+      'No Grafana Cloud stack selected. Please sign in to Grafana Cloud first.'
+    )
+  }
+
+  const tokens = await getValidAssistantTokens(stackId)
+
+  if (!tokens) {
+    throw new Error(
+      'Not authenticated with Grafana Assistant. Please connect to Grafana Assistant first.'
+    )
+  }
+
+  return {
+    baseUrl: `${tokens.apiEndpoint}/api/cli/v1/a2a`,
+    agentId: AGENT_ID,
+    extensions: [REMOTE_TOOL_EXTENSION, TOKEN_STREAMING_EXTENSION].join(', '),
+    bearerToken: tokens.accessToken,
+  }
+}

src/handlers/ai/a2a/constants.ts

@@ -0,0 +1,18 @@
+export const LOG_PREFIX = '[GrafanaAssistant]'
+
+export const NO_USAGE = {
+  inputTokens: undefined,
+  outputTokens: undefined,
+  totalTokens: undefined,
+} as const
+
+/** A2A artifact name constants used in SSE events */
+export const ARTIFACT_NAME = {
+  STEP_TOOL_CALL: 'step.toolCall',
+  STEP_COMPLETE: 'step.complete',
+  STEP_MESSAGE: 'step.message',
+  STEP_TOOL_RESULT: 'step.toolResult',
+  MESSAGE_STREAM_START: 'message.stream.start',
+  MESSAGE_CONTENT_DELTA: 'message.content.delta',
+  MESSAGE_STREAM_COMPLETE: 'message.stream.complete',
+} as const