internal: Add assistant feature toggle (#1100)

Author: e-fisher

src/handlers/ai/grafanaAssistantProvider.ts

@@ -0,0 +1,24 @@
+import { LanguageModelV2, LanguageModelV2CallOptions } from '@ai-sdk/provider'
+
+/**
+ * Placeholder implementation of the Grafana Assistant language model.
+ * The actual A2A integration will be implemented in a follow-up PR.
+ */
+export class GrafanaAssistantLanguageModel implements LanguageModelV2 {
+  readonly specificationVersion = 'v2' as const
+  readonly provider = 'grafana-assistant'
+  readonly modelId = 'grafana_assistant_k6_studio'
+  readonly supportedUrls = {}
+
+  doGenerate(
+    _options: LanguageModelV2CallOptions
+  ): ReturnType<LanguageModelV2['doGenerate']> {
+    throw new Error('not implemented')
+  }
+
+  doStream(
+    _options: LanguageModelV2CallOptions
+  ): ReturnType<LanguageModelV2['doStream']> {
+    throw new Error('not implemented')
+  }
+}

src/handlers/ai/index.ts

@@ -2,7 +2,7 @@ import { OpenAIResponsesProviderOptions } from '@ai-sdk/openai'
 import { convertToModelMessages, streamText } from 'ai'
 import { ipcMain, IpcMainEvent } from 'electron'
 
-import { setupAiModel } from './model'
+import { getGrafanaAssistantModel, getOpenAiModel } from './model'
 import { streamMessages } from './streamMessages'
 import { tools } from './tools'
 import { AiHandler, StreamChatRequest, AbortStreamChatRequest } from './types'
@@ -19,31 +19,47 @@ async function handleStreamChat(
   event: IpcMainEvent,
   request: StreamChatRequest
 ) {
-  const aiModel = await setupAiModel()
+  const provider = request.provider ?? 'openai'
   const messages = convertToModelMessages(request.messages)
 
   const abortController = new AbortController()
   activeAbortControllers.set(request.id, abortController)
 
   try {
-    const response = streamText({
-      model: aiModel,
-      toolChoice: 'required',
-      messages,
-      tools,
-      abortSignal: abortController.signal,
-      providerOptions: {
-        openai: {
-          parallelToolCalls: false,
-          reasoningEffort: 'low',
-          textVerbosity: 'low',
-          // Disable storing of conversations, required for orgs with zero data retention
-          store: false,
-        } satisfies OpenAIResponsesProviderOptions,
-      },
-    })
+    if (provider === 'grafana-assistant') {
+      const aiModel = getGrafanaAssistantModel()
 
-    await streamMessages(event.sender, response, request.id)
+      const response = streamText({
+        model: aiModel,
+        toolChoice: 'required',
+        messages,
+        tools,
+        abortSignal: abortController.signal,
+      })
+
+      await streamMessages(event.sender, response, request.id, false)
+    } else {
+      const aiModel = await getOpenAiModel()
+
+      const response = streamText({
+        model: aiModel,
+        toolChoice: 'required',
+        messages,
+        tools,
+        abortSignal: abortController.signal,
+        providerOptions: {
+          openai: {
+            parallelToolCalls: false,
+            reasoningEffort: 'low',
+            textVerbosity: 'low',
+            // Disable storing of conversations, required for orgs with zero data retention
+            store: false,
+          } satisfies OpenAIResponsesProviderOptions,
+        },
+      })
+
+      await streamMessages(event.sender, response, request.id, true)
+    }
   } finally {
     // Clean up the AbortController after streaming completes or fails
     activeAbortControllers.delete(request.id)

src/handlers/ai/model.ts

@@ -3,20 +3,27 @@ import { LanguageModel } from 'ai'
 
 import { getDecryptedAiKey } from '@/main/settings'
 
-let model: LanguageModel | null = null
+import { GrafanaAssistantLanguageModel } from './grafanaAssistantProvider'
 
-export async function setupAiModel() {
-  if (model) {
-    return model
+let openAiModel: LanguageModel | null = null
+const grafanaAssistantModel = new GrafanaAssistantLanguageModel()
+
+export async function getOpenAiModel() {
+  if (openAiModel) {
+    return openAiModel
   }
 
   const apiKey = await getDecryptedAiKey()
   const provider = createOpenAI({ apiKey })
-  model = provider('gpt-5')
+  openAiModel = provider('gpt-5')
+
+  return openAiModel
+}
 
-  return model
+export function resetOpenAiModel() {
+  openAiModel = null
 }
 
-export function resetAiModel() {
-  model = null
+export function getGrafanaAssistantModel() {
+  return grafanaAssistantModel
 }

src/handlers/ai/streamMessages.ts

@@ -5,7 +5,8 @@ import { AiHandler } from './types'
 export async function streamMessages<Tools extends ToolSet, PARTIAL_OUTPUT>(
   webContents: Electron.WebContents,
   response: StreamTextResult<Tools, PARTIAL_OUTPUT>,
-  requestId: string
+  requestId: string,
+  includeUsage: boolean
 ) {
   const stream = response.toUIMessageStream({})
 
@@ -16,7 +17,7 @@ export async function streamMessages<Tools extends ToolSet, PARTIAL_OUTPUT>(
     })
   }
 
-  const usageData = await response.usage
+  const usageData = includeUsage ? await response.usage : undefined
 
   webContents.send(AiHandler.StreamChatEnd, {
     id: requestId,

src/handlers/ai/types.ts

@@ -1,5 +1,7 @@
 import { LanguageModelUsage, UIMessage, UIMessageChunk } from 'ai'
 
+import { AiProvider } from '@/types/features'
+
 export enum AiHandler {
   StreamChat = 'ai:streamChat',
   StreamChatChunk = 'ai:streamChatChunk',
@@ -14,6 +16,7 @@ export interface StreamChatRequest {
   messages: UIMessage[]
   headers?: Record<string, string>
   body?: object
+  provider?: AiProvider
 }
 
 export interface StreamChatChunk {

src/main/settings.ts

@@ -4,7 +4,7 @@ import { existsSync, readFileSync } from 'fs'
 import { writeFile, open } from 'fs/promises'
 import path from 'path'
 
-import { resetAiModel } from '@/handlers/ai/model'
+import { resetOpenAiModel } from '@/handlers/ai/model'
 import { configureSystemProxy } from '@/services/http'
 
 import { AppSettingsSchema } from '../schemas/settings'
@@ -185,7 +185,7 @@ export async function applySettings(
 
   if (modifiedSettings.ai) {
     k6StudioState.appSettings.ai = modifiedSettings.ai
-    resetAiModel()
+    resetOpenAiModel()
   }
 }
 

src/store/features/useFeaturesStore.ts

@@ -12,6 +12,7 @@ export const defaultFeatures: Record<Feature, boolean> = {
   'dummy-feature': false,
   'typeahead-json': false,
   'browser-test-editor': import.meta.env.DEV,
+  'grafana-assistant': false,
 }
 
 export const useFeaturesStore = create<FeaturesStore>()(

src/types/features.ts

@@ -1 +1,7 @@
-export type Feature = 'dummy-feature' | 'typeahead-json' | 'browser-test-editor'
+export type Feature =
+  | 'dummy-feature'
+  | 'typeahead-json'
+  | 'browser-test-editor'
+  | 'grafana-assistant'
+
+export type AiProvider = 'openai' | 'grafana-assistant'

src/views/Generator/AutoCorrelation/ErrorMessage.tsx

@@ -3,6 +3,7 @@ import { ExternalLink, KeyIcon, RefreshCw } from 'lucide-react'
 
 import grotCrashed from '@/assets/grot-crashed.svg'
 import { useSettingsChanged } from '@/hooks/useSettings'
+import { useFeaturesStore } from '@/store/features'
 import { useStudioUIStore } from '@/store/ui'
 
 interface AutoCorrelationErrorProps {
@@ -11,6 +12,18 @@ interface AutoCorrelationErrorProps {
 }
 
 export function ErrorMessage({ error, onRetry }: AutoCorrelationErrorProps) {
+  const isGrafanaAssistant = useFeaturesStore(
+    (state) => state.features['grafana-assistant']
+  )
+
+  if (isGrafanaAssistant) {
+    return <GrafanaAssistantError error={error} onRetry={onRetry} />
+  }
+
+  return <OpenAiError error={error} onRetry={onRetry} />
+}
+
+function OpenAiError({ error, onRetry }: AutoCorrelationErrorProps) {
   const errorMessage = error.message.toLowerCase()
   const openSettingsDialog = useStudioUIStore(
     (state) => state.openSettingsDialog
@@ -83,6 +96,34 @@ export function ErrorMessage({ error, onRetry }: AutoCorrelationErrorProps) {
   )
 }
 
+function GrafanaAssistantError({ onRetry }: AutoCorrelationErrorProps) {
+  const retryButton = (
+    <Button onClick={onRetry}>
+      <RefreshCw />
+      Retry
+    </Button>
+  )
+
+  const reportIssueButton = (
+    <Button onClick={() => window.studio.ui.reportIssue()} variant="outline">
+      <ExternalLink />
+      Report issue
+    </Button>
+  )
+
+  // TODO: Add Assistant specific error handling
+
+  return (
+    <MessageContent
+      title="Something went wrong"
+      message="An unexpected error occurred during autocorrelation. Click retry to try again or report an issue if problem persists."
+    >
+      {retryButton}
+      {reportIssueButton}
+    </MessageContent>
+  )
+}
+
 function MessageContent({
   title,
   message,

src/views/Generator/AutoCorrelation/IntroductionMessage.tsx

@@ -4,19 +4,31 @@ import { CheckCircleIcon, KeyIcon, WandSparkles } from 'lucide-react'
 import grotIllustration from '@/assets/grot-magic.svg'
 import { useProxyStatus } from '@/hooks/useProxyStatus'
 import { useSettings } from '@/hooks/useSettings'
+import { useFeaturesStore } from '@/store/features'
 import { useStudioUIStore } from '@/store/ui'
 
 interface IntroductionMessageProps {
   onStart: () => void
 }
 
 export function IntroductionMessage({ onStart }: IntroductionMessageProps) {
+  const isGrafanaAssistant = useFeaturesStore(
+    (state) => state.features['grafana-assistant']
+  )
+
+  if (isGrafanaAssistant) {
+    return <GrafanaAssistantIntro onStart={onStart} />
+  }
+
+  return <OpenAiIntro onStart={onStart} />
+}
+
+function OpenAiIntro({ onStart }: IntroductionMessageProps) {
   const openSettingsDialog = useStudioUIStore(
     (state) => state.openSettingsDialog
   )
   const { data: settings } = useSettings()
   const isAiConfigured = !!settings?.ai.apiKey
-
   const proxyStatus = useProxyStatus()
 
   return (
@@ -93,6 +105,52 @@ export function IntroductionMessage({ onStart }: IntroductionMessageProps) {
   )
 }
 
+function GrafanaAssistantIntro({ onStart }: IntroductionMessageProps) {
+  return (
+    <Flex
+      direction="column"
+      align="center"
+      gap="6"
+      justify="center"
+      height="100%"
+    >
+      <img
+        src={grotIllustration}
+        role="img"
+        aria-label="Grafana mascot illustration"
+        css={{ maxWidth: 250 }}
+      />
+
+      <Flex
+        direction="column"
+        align="center"
+        gap="4"
+        maxWidth="600px"
+        css={{ textAlign: 'center' }}
+      >
+        <Badge color="orange" variant="soft">
+          Coming Soon
+        </Badge>
+        <Text size="3" weight="bold">
+          Automatically correlate dynamic values
+        </Text>
+        <Text size="2" color="gray" mb="2">
+          Powered by Grafana Assistant
+        </Text>
+
+        <Button onClick={onStart} size="3" disabled>
+          <WandSparkles />
+          Analyze recording
+        </Button>
+
+        <Text size="1" color="gray" mt="1">
+          This feature is in public preview and subject to change.
+        </Text>
+      </Flex>
+    </Flex>
+  )
+}
+
 function ListItem({ children }: { children: React.ReactNode }) {
   return (
     <Flex align="center" gap="2">