Add AI Studio user ID support

Author: robertpiosik

packages/browser/src/content-scripts/send-prompt-content-script/send-prompt-content-script.ts

@@ -34,7 +34,7 @@ const is_cwc_hash = hash.startsWith(hash_prefix_new)
 
 const batch_id = hash.substring(hash_prefix_new.length + 1)
 
-const ai_studio_url = 'https://aistudio.google.com/prompts/new_chat'
+const ai_studio_url = 'https://aistudio.google.com/'
 const is_ai_studio = current_url.startsWith(ai_studio_url)
 
 const gemini_url = 'https://gemini.google.com/'

packages/browser/src/manifest.json

@@ -24,7 +24,7 @@
   "content_scripts": [
     {
       "matches": [
-        "https://aistudio.google.com/prompts/new_chat*",
+        "https://aistudio.google.com/*",
         "https://chatgpt.com/*",
         "https://claude.ai/*",
         "https://copilot.microsoft.com/",

packages/vscode/package.json

@@ -600,6 +600,13 @@
       "type": "object",
       "title": "Code Web Chat Settings",
       "properties": {
+        "codeWebChat.aiStudioUserId": {
+          "type": [
+            "integer",
+            "null"
+          ],
+          "default": null
+        },
         "codeWebChat.geminiUserId": {
           "type": [
             "integer",

packages/vscode/src/services/websocket-manager.ts

@@ -279,6 +279,7 @@ export class WebSocketManager {
     const web_chat_presets =
       config.get<ConfigPresetFormat[]>(params.presets_config_key) ?? []
     const gemini_user_id = config.get<number | null>('geminiUserId')
+    const ai_studio_user_id = config.get<number | null>('aiStudioUserId')
 
     for (const chat of params.chats) {
       const preset = web_chat_presets.find((p) => p.name == chat.preset_name)
@@ -289,10 +290,17 @@ export class WebSocketManager {
       const chatbot = CHATBOTS[preset.chatbot as keyof typeof CHATBOTS]
       let url: string
       if (preset.chatbot == 'AI Studio') {
+        let base_url = chatbot.url
+        if (ai_studio_user_id !== null && ai_studio_user_id !== undefined) {
+          base_url = base_url.replace(
+            'https://aistudio.google.com/',
+            `https://aistudio.google.com/u/${ai_studio_user_id}/`
+          )
+        }
         if (preset.model) {
-          url = `${chatbot.url}?model=${preset.model}`
+          url = `${base_url}?model=${preset.model}`
         } else {
-          url = chatbot.url
+          url = base_url
         }
       } else if (preset.chatbot == 'Open WebUI') {
         if (preset.port) {
@@ -362,14 +370,22 @@ export class WebSocketManager {
 
     const config = vscode.workspace.getConfiguration('codeWebChat')
     const gemini_user_number = config.get<number | null>('geminiUserNumber')
+    const ai_studio_user_id = config.get<number | null>('aiStudioUserId')
 
     const chatbot = CHATBOTS[params.preset.chatbot as keyof typeof CHATBOTS]
     let url: string
     if (params.preset.chatbot == 'AI Studio') {
+      let base_url = chatbot.url
+      if (ai_studio_user_id !== null && ai_studio_user_id !== undefined) {
+        base_url = base_url.replace(
+          'https://aistudio.google.com/',
+          `https://aistudio.google.com/u/${ai_studio_user_id}/`
+        )
+      }
       if (params.preset.model) {
-        url = `${chatbot.url}?model=${params.preset.model}`
+        url = `${base_url}?model=${params.preset.model}`
       } else {
-        url = chatbot.url
+        url = base_url
       }
     } else if (params.preset.chatbot == 'Open WebUI') {
       if (params.preset.port) {

packages/vscode/src/views/settings/backend/message-handlers/handle-get-ai-studio-user-id.ts

@@ -0,0 +1,13 @@
+import * as vscode from 'vscode'
+import { SettingsProvider } from '@/views/settings/backend/settings-provider'
+
+export const handle_get_ai_studio_user_id = async (
+  provider: SettingsProvider
+): Promise<void> => {
+  const config = vscode.workspace.getConfiguration('codeWebChat')
+  const aiStudioUserId = config.get<number | null>('aiStudioUserId') ?? null
+  provider.postMessage({
+    command: 'AI_STUDIO_USER_ID',
+    aiStudioUserId
+  })
+}

packages/vscode/src/views/settings/backend/message-handlers/handle-update-ai-studio-user-id.ts

@@ -0,0 +1,14 @@
+import * as vscode from 'vscode'
+import { UpdateAiStudioUserIdMessage } from '@/views/settings/types/messages'
+
+export const handle_update_ai_studio_user_id = async (
+  message: UpdateAiStudioUserIdMessage
+): Promise<void> => {
+  await vscode.workspace
+    .getConfiguration('codeWebChat')
+    .update(
+      'aiStudioUserId',
+      message.aiStudioUserId || undefined,
+      vscode.ConfigurationTarget.Global
+    )
+}

packages/vscode/src/views/settings/backend/message-handlers/index.ts

@@ -31,6 +31,8 @@ export * from './handle-get-clear-checks-in-workspace-behavior'
 export * from './handle-update-clear-checks-in-workspace-behavior'
 export * from './handle-get-gemini-user-id'
 export * from './handle-update-gemini-user-id'
+export * from './handle-get-ai-studio-user-id'
+export * from './handle-update-ai-studio-user-id'
 export * from './handle-get-context-size-warning-threshold'
 export * from './handle-update-context-size-warning-threshold'
 export * from './handle-get-are-automatic-checkpoints-disabled'

packages/vscode/src/views/settings/backend/settings-provider.ts

@@ -23,6 +23,7 @@ import {
   handle_get_edit_context_system_instructions,
   handle_get_edit_format_instructions,
   handle_get_gemini_user_id,
+  handle_get_ai_studio_user_id,
   handle_get_intelligent_update_configurations,
   handle_get_model_providers,
   handle_get_checkpoint_lifespan,
@@ -42,6 +43,7 @@ import {
   handle_update_edit_context_system_instructions,
   handle_update_edit_format_instructions,
   handle_update_gemini_user_id,
+  handle_update_ai_studio_user_id,
   handle_upsert_configuration,
   handle_open_ignore_patterns_settings
 } from './message-handlers'
@@ -199,6 +201,10 @@ export class SettingsProvider {
           await handle_get_gemini_user_id(this)
         } else if (message.command == 'UPDATE_GEMINI_USER_ID') {
           await handle_update_gemini_user_id(message)
+        } else if (message.command == 'GET_AI_STUDIO_USER_ID') {
+          await handle_get_ai_studio_user_id(this)
+        } else if (message.command == 'UPDATE_AI_STUDIO_USER_ID') {
+          await handle_update_ai_studio_user_id(message)
         } else if (message.command == 'GET_CHECKPOINT_LIFESPAN') {
           await handle_get_checkpoint_lifespan(this)
         } else if (
@@ -246,6 +252,7 @@ export class SettingsProvider {
           void handle_get_are_automatic_checkpoints_disabled(this)
           void handle_get_checkpoint_lifespan(this)
           void handle_get_gemini_user_id(this)
+          void handle_get_ai_studio_user_id(this)
         }
       })
     )

packages/vscode/src/views/settings/frontend/Home/Home.tsx

@@ -78,6 +78,7 @@ type Props = {
   commit_message_auto_accept_after: number | null
   context_size_warning_threshold: number
   gemini_user_id: number | null
+  ai_studio_user_id: number | null
   are_automatic_checkpoints_disabled: boolean
   checkpoint_lifespan: number
   edit_format_instructions: EditFormatInstructions
@@ -99,6 +100,7 @@ type Props = {
   on_automatic_checkpoints_toggle: (disabled: boolean) => void
   on_checkpoint_lifespan_change: (hours: number) => void
   on_gemini_user_id_change: (id: number | null) => void
+  on_ai_studio_user_id_change: (id: number | null) => void
   on_clear_checks_in_workspace_behavior_change: (
     value: 'ignore-open-editors' | 'uncheck-all'
   ) => void
@@ -314,7 +316,9 @@ export const Home: React.FC<Props> = (props) => {
         <PresetsSection
           ref={(el) => (section_refs.current['presets'] = el)}
           gemini_user_id={props.gemini_user_id}
+          ai_studio_user_id={props.ai_studio_user_id}
           on_gemini_user_id_change={props.on_gemini_user_id_change}
+          on_ai_studio_user_id_change={props.on_ai_studio_user_id_change}
           on_stuck_change={presets_on_stuck_change}
         />
         <Section

packages/vscode/src/views/settings/frontend/Home/sections/PresetsSection.tsx

@@ -6,13 +6,16 @@ import { Group } from '@ui/components/editor/settings/Group/Group'
 
 type Props = {
   gemini_user_id: number | null
+  ai_studio_user_id: number | null
   on_gemini_user_id_change: (id: number | null) => void
+  on_ai_studio_user_id_change: (id: number | null) => void
   on_stuck_change: (is_stuck: boolean) => void
 }
 
 export const PresetsSection = forwardRef<HTMLDivElement, Props>(
   (props, ref) => {
     const [gemini_user_id_str, set_gemini_user_id_str] = useState('')
+    const [ai_studio_user_id_str, set_ai_studio_user_id_str] = useState('')
 
     useEffect(() => {
       set_gemini_user_id_str(
@@ -22,6 +25,15 @@ export const PresetsSection = forwardRef<HTMLDivElement, Props>(
       )
     }, [props.gemini_user_id])
 
+    useEffect(() => {
+      set_ai_studio_user_id_str(
+        props.ai_studio_user_id === null ||
+          props.ai_studio_user_id === undefined
+          ? ''
+          : String(props.ai_studio_user_id)
+      )
+    }, [props.ai_studio_user_id])
+
     const handle_gemini_user_id_blur = () => {
       if (gemini_user_id_str === '') {
         props.on_gemini_user_id_change(null)
@@ -33,6 +45,17 @@ export const PresetsSection = forwardRef<HTMLDivElement, Props>(
       }
     }
 
+    const handle_ai_studio_user_id_blur = () => {
+      if (ai_studio_user_id_str === '') {
+        props.on_ai_studio_user_id_change(null)
+        return
+      }
+      const num_id = parseInt(ai_studio_user_id_str, 10)
+      if (!isNaN(num_id) && num_id >= 0) {
+        props.on_ai_studio_user_id_change(num_id)
+      }
+    }
+
     return (
       <Section
         ref={ref}
@@ -54,6 +77,19 @@ export const PresetsSection = forwardRef<HTMLDivElement, Props>(
               />
             }
           />
+          <Item
+            title="AI Studio User ID"
+            description="Run AI Studio chatbot as non-default user. Check URL for the numeric ID."
+            slot_right={
+              <Input
+                type="number"
+                value={ai_studio_user_id_str}
+                on_change={set_ai_studio_user_id_str}
+                on_blur={handle_ai_studio_user_id_blur}
+                max_width={60}
+              />
+            }
+          />
         </Group>
       </Section>
     )