feat(runner): adds sql query runner ai tab auto completion

apps/api/orpc/routers/ai/code-completion.ts

@@ -0,0 +1,78 @@
+import { google } from '@ai-sdk/google'
+import { ConnectionType } from '@conar/shared/enums/connection-type'
+import { generateText } from 'ai'
+import { type } from 'arktype'
+import { orpc, requireSubscriptionMiddleware } from '~/orpc'
+
+export const codeCompletion = orpc
+  .use(requireSubscriptionMiddleware)
+  .input(type({
+    context: 'string',
+    suffix: 'string?',
+    instruction: 'string?',
+    fileContent: 'string',
+    type: type.valueOf(ConnectionType),
+    schemaContext: 'string?',
+  }))
+  .output(type({
+    completion: 'string',
+  }))
+  .handler(async ({ input }) => {
+    try {
+      const enhancedPrompt = [
+        'You are an expert SQL code completion assistant.',
+        '',
+        `Database Type: ${input.type}`,
+        '',
+        'Database Schema:',
+        input.schemaContext || '',
+        '',
+        'Current Code BEFORE Cursor:',
+        input.context || '',
+        '',
+        'Current Code AFTER Cursor:',
+        input.suffix || '',
+        '',
+        'Full File Content:',
+        input.fileContent || '',
+        '',
+        `Task: ${input.instruction || 'Complete the SQL query'}`,
+        '',
+        'General Safety Rules:',
+        '- Prioritize SELECT queries (read-only) when the intent is ambiguous.',
+        '- If the user explicitly starts writing a data modification query (INSERT, UPDATE, etc.), complete it validly.',
+        '- Do NOT generate destructive schema changes (DROP TABLE, TRUNCATE, ALTER) unless explicitly requested in a generic comment.',
+        '',
+        'Specific Instructions:',
+        '- If the context is empty (start of file), return an empty string.',
+        '- For UUID/CUID: Use generation functions (e.g. gen_random_uuid()) instead of hardcoded strings if possible.',
+        '- For Foreign Keys: Do NOT plain guess IDs. Use subqueries (e.g. "(SELECT id FROM users ... LIMIT 1)") to make the query robust without needing external tools.',
+        '- Look at "Current Code AFTER Cursor". Do NOT generate code that is already present there (like closing parenthesis or semicolons). Make your completion join smoothly with the suffix.',
+        '- Intelligent Spacing: If the code before the cursor is a keyword or identifier (e.g., \'SELECT\', \'FROM\', \'tableName\') and your completion starts with a keyword/identifier, ensure there is a space at the beginning of your completion.',
+        '',
+        'Generate ONLY the completion text that should be inserted (or used to replace the last word).',
+        '- Do NOT wrap in markdown code blocks.',
+        '- If no completion is needed, return an empty string.',
+      ].join('\n')
+
+      const result = await generateText({
+        model: google('gemini-2.5-flash'),
+        prompt: enhancedPrompt,
+      })
+
+      let cleanText = result.text.trim()
+      if (cleanText.startsWith('```')) {
+        cleanText = cleanText.replace(/^```\w*\s*/, '').replace(/\s*```$/, '')
+      }
+
+      return {
+        completion: cleanText,
+      }
+    }
+    catch (error) {
+      console.error('Error in completion handler:', error instanceof Error ? error.message : 'Unknown error')
+      return {
+        completion: '',
+      }
+    }
+  })

apps/api/orpc/routers/ai/index.ts

@@ -1,4 +1,5 @@
 export { chat } from './chat'
+export { codeCompletion } from './code-completion'
 export { enhancePrompt } from './enhance-prompt'
 export { filters } from './filters'
 export { fixSQL } from './fix-sql'

apps/api/package.json

@@ -39,6 +39,7 @@
     "hono": "catalog:",
     "ioredis": "catalog:",
     "ioredis-mock": "catalog:",
+    "monacopilot": "catalog:",
     "nanoid": "catalog:",
     "pg": "catalog:",
     "posthog-node": "catalog:",

apps/desktop/package.json

@@ -31,6 +31,7 @@
     "@clickhouse/client": "catalog:",
     "@todesktop/runtime": "catalog:",
     "electron-store": "catalog:",
+    "monacopilot": "catalog:",
     "mssql": "catalog:",
     "mysql2": "catalog:",
     "pg": "catalog:"

apps/desktop/src/entities/connection/utils/monaco.ts

@@ -171,3 +171,51 @@ export function connectionCompletionService(connection: typeof connections.$infe
     return [...items, ...keywordItems]
   }
 }
+
+export async function connectionAICompletionContext(connection: typeof connections.$inferSelect) {
+  const store = connectionStore(connection.id)
+  queryClient.prefetchQuery(connectionTablesAndSchemasQuery({ connection, showSystem: store.state.showSystem }))
+  queryClient.prefetchQuery(connectionEnumsQuery({ connection }))
+
+  const buildSchemaContext = async (): Promise<string> => {
+    const [tablesAndSchemas, enums] = await Promise.all([
+      queryClient.ensureQueryData(connectionTablesAndSchemasQuery({ connection, showSystem: store.state.showSystem })),
+      queryClient.ensureQueryData(connectionEnumsQuery({ connection })),
+    ])
+
+    const contextLines: string[] = []
+
+    for (const schema of tablesAndSchemas.schemas) {
+      contextLines.push('', `Schema: ${schema.name}`, '')
+
+      for (const tableName of schema.tables) {
+        const columns = await queryClient.ensureQueryData(
+          connectionTableColumnsQuery({
+            connection,
+            schema: schema.name,
+            table: tableName,
+          }),
+        )
+
+        contextLines.push(`  Table: ${tableName}`, '')
+        columns.forEach((col) => {
+          contextLines.push(`    - ${col.id}: ${col.type} ${col.isNullable ? '(nullable)' : '(not null)'}`)
+        })
+      }
+    }
+
+    if (enums?.length > 0) {
+      contextLines.push('', 'Enums:', '')
+      enums.forEach((enumItem) => {
+        contextLines.push(`  ${enumItem.schema}.${enumItem.name}: [${enumItem.values.join(', ')}]`)
+      })
+    }
+
+    return contextLines.join('\n')
+  }
+
+  return {
+    type: connection.type,
+    buildSchemaContext,
+  }
+}

apps/desktop/src/routes/_protected/database/$id/sql/-components/runner/ai-tab-completion.tsx

@@ -0,0 +1,231 @@
+import type { ConnectionType } from '@conar/shared/enums/connection-type'
+import type { editor } from 'monaco-editor'
+import type { CompletionRegistration } from 'monacopilot'
+import type { RefObject } from 'react'
+import * as monaco from 'monaco-editor'
+import { LanguageIdEnum } from 'monaco-sql-languages'
+import { registerCompletion } from 'monacopilot'
+import { useCallback, useEffect, useRef } from 'react'
+import { connectionAICompletionContext } from '~/entities/connection/utils/monaco'
+import { orpc } from '~/lib/orpc'
+import { Route } from '../..'
+
+const dialectsMap = {
+  postgres: LanguageIdEnum.PG,
+  mysql: LanguageIdEnum.MYSQL,
+  mssql: LanguageIdEnum.PG,
+  clickhouse: LanguageIdEnum.MYSQL,
+} satisfies Record<ConnectionType, LanguageIdEnum>
+
+interface Cache {
+  before: string
+  completion: string
+}
+
+function getCache(
+  model: editor.ITextModel,
+  position: monaco.Position,
+  cache: Cache | null,
+): { completion: string } | null {
+  if (!cache)
+    return null
+
+  const offset = model.getOffsetAt(position)
+  const textBefore = model.getValue().substring(0, offset)
+
+  if (!textBefore.startsWith(cache.before))
+    return null
+
+  const addedText = textBefore.slice(cache.before.length)
+
+  if (!cache.completion.startsWith(addedText))
+    return null
+
+  const remainingCompletion = cache.completion.slice(addedText.length)
+  return { completion: remainingCompletion }
+}
+
+async function fetchCompletion(
+  model: editor.ITextModel,
+  position: monaco.Position,
+  aiConfig: Awaited<ReturnType<typeof connectionAICompletionContext>>,
+  signal: AbortSignal,
+): Promise<{ result: { completion: string }, context: string } | null> {
+  const offset = model.getOffsetAt(position)
+  const context = model.getValue().substring(0, offset)
+  const suffix = model.getValue().substring(offset)
+
+  if (context.trim().length < 2)
+    return null
+
+  const schemaContext = await aiConfig.buildSchemaContext()
+  const result = await orpc.ai.codeCompletion({
+    context,
+    suffix,
+    instruction: 'Complete the SQL query with secure and safe optimised version',
+    fileContent: model.getValue(),
+    type: aiConfig.type,
+    schemaContext,
+  }, { signal })
+
+  return { result, context }
+}
+
+export function useRunnerEditorAiTabCompletion(monacoRef: RefObject<editor.IStandaloneCodeEditor | null>) {
+  const { connection } = Route.useRouteContext()
+
+  const completionRef = useRef<CompletionRegistration | null>(null)
+  const completionCacheRef = useRef<Cache | null>(null)
+  const debounceTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
+  const pendingResolveRef = useRef<((value: { completion: string }) => void) | null>(null)
+  const aiConfigRef = useRef<Awaited<ReturnType<typeof connectionAICompletionContext>> | null>(null)
+  const abortControllerRef = useRef<AbortController | null>(null)
+
+  const clearPending = useCallback(() => {
+    if (debounceTimerRef.current) {
+      clearTimeout(debounceTimerRef.current)
+      debounceTimerRef.current = null
+    }
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort()
+      abortControllerRef.current = null
+    }
+    if (pendingResolveRef.current) {
+      pendingResolveRef.current({ completion: '' })
+      pendingResolveRef.current = null
+    }
+  }, [])
+
+  const requestHandler = useCallback(async () => {
+    const model = monacoRef.current?.getModel()
+    const position = monacoRef.current?.getPosition()
+    const aiConfig = aiConfigRef.current
+
+    if (!model || !position || !aiConfig)
+      return { completion: '' }
+
+    const cachedResult = getCache(model, position, completionCacheRef.current)
+    if (cachedResult)
+      return cachedResult
+
+    if (completionCacheRef.current)
+      completionCacheRef.current = null
+
+    clearPending()
+
+    return new Promise<{ completion: string }>((resolve) => {
+      pendingResolveRef.current = resolve
+      debounceTimerRef.current = setTimeout(async () => {
+        if (pendingResolveRef.current !== resolve) {
+          resolve({ completion: '' })
+          return
+        }
+
+        const model = monacoRef.current?.getModel()
+        const position = monacoRef.current?.getPosition()
+
+        if (!model || !position || !aiConfigRef.current) {
+          resolve({ completion: '' })
+          return
+        }
+
+        const controller = new AbortController()
+        abortControllerRef.current = controller
+
+        try {
+          const fetchResult = await fetchCompletion(
+            model,
+            position,
+            aiConfigRef.current,
+            controller.signal,
+          )
+
+          if (!fetchResult) {
+            resolve({ completion: '' })
+            return
+          }
+
+          completionCacheRef.current = {
+            before: fetchResult.context,
+            completion: fetchResult.result.completion,
+          }
+
+          if (pendingResolveRef.current === resolve) {
+            resolve({ completion: fetchResult.result.completion })
+            pendingResolveRef.current = null
+          }
+          else {
+            resolve({ completion: '' })
+          }
+        }
+        catch (error) {
+          if (error instanceof Error && error.name === 'AbortError') {
+            resolve({ completion: '' })
+            return
+          }
+          console.error('AI completion error:', error)
+          resolve({ completion: '' })
+        }
+        finally {
+          if (pendingResolveRef.current === resolve)
+            pendingResolveRef.current = null
+        }
+      }, 600)
+    })
+  }, [clearPending, monacoRef])
+
+  useEffect(() => {
+    let isMounted = true
+
+    const setupCompletion = async () => {
+      if (!monacoRef.current)
+        return
+
+      const config = await connectionAICompletionContext(connection)
+      if (!isMounted)
+        return
+      aiConfigRef.current = config
+
+      if (completionRef.current) {
+        completionRef.current.deregister()
+      }
+
+      // Monkey patch to fix 'disposeInlineCompletions is not a function' error from monacopilot
+      const originalRegister = monaco.languages.registerInlineCompletionsProvider
+      // @ts-expect-error - patching internal api
+      if (!originalRegister.__patched) {
+        monaco.languages.registerInlineCompletionsProvider = (languageSelector, provider) => {
+          if (!provider.disposeInlineCompletions) {
+            // eslint-disable-next-line ts/ban-ts-comment
+            // @ts-ignore
+            provider.disposeInlineCompletions = () => {}
+          }
+          return originalRegister.call(monaco.languages, languageSelector, provider)
+        }
+        // @ts-expect-error - patching internal api
+        monaco.languages.registerInlineCompletionsProvider.__patched = true
+      }
+
+      completionRef.current = registerCompletion(
+        monaco,
+        monacoRef.current,
+        {
+          trigger: 'onTyping',
+          language: dialectsMap[connection.type],
+          requestHandler,
+        },
+      )
+    }
+
+    setupCompletion()
+
+    return () => {
+      isMounted = false
+      clearPending()
+      if (completionRef.current) {
+        completionRef.current.deregister()
+        completionRef.current = null
+      }
+    }
+  }, [clearPending, connection, monacoRef, requestHandler])
+}