Bring to head of internal repo

package.json

@@ -1,6 +1,6 @@
 {
   "name": "chrome-extension-webpack",
-  "version": "1.2.18",
+  "version": "1.2.26",
   "description": "Get started with Chrome extensions development using webpack, Typescript, Sass, and more",
   "scripts": {
     "generate": "rm -rf ./proto && node generate.js",

src/auth.ts

@@ -1,7 +1,9 @@
 import { v4 as uuidv4 } from 'uuid';
 
-export const PROFILE_URL = 'https://www.codeium.com/profile';
+import { getGeneralProfileUrl } from './storage';
 
+// Runs in popup.
+// TODO(prem): Move to a popup-specific source file.
 export async function openAuthTab(): Promise<void> {
   const uuid = uuidv4();
   await chrome.runtime.sendMessage({
@@ -10,13 +12,25 @@ export async function openAuthTab(): Promise<void> {
       state: uuid,
     },
   });
+  const profileUrl = await getGeneralProfileUrl();
+
   await chrome.tabs.create({
-    url: `${PROFILE_URL}?redirect_uri=chrome-extension://${chrome.runtime.id}&state=${uuid}`,
+    url: `${profileUrl}?redirect_uri=chrome-extension://${chrome.runtime.id}&state=${uuid}`,
   });
 }
 
-export async function registerUser(token: string): Promise<{ api_key: string; name: string }> {
-  const url = new URL('register_user/', 'https://api.codeium.com');
+// Runs in service worker.
+// TODO(prem): Move to a service worker-specific source file.
+export async function registerUser(
+  token: string,
+  portalUrl: string | undefined
+): Promise<{ api_key: string; name: string }> {
+  const url = ((): URL => {
+    if (portalUrl === undefined) {
+      return new URL('register_user/', 'https://api.codeium.com');
+    }
+    return new URL('_route/api_server/exa.api_server_pb.ApiServerService/RegisterUser', portalUrl);
+  })();
   const response = await fetch(url, {
     body: JSON.stringify({ firebase_id_token: token }),
     method: 'POST',
@@ -25,7 +39,7 @@ export async function registerUser(token: string): Promise<{ api_key: string; na
     },
   });
   if (!response.ok) {
-    throw new Error(response.statusText);
+    throw new Error(`${url}: ${response.statusText}`);
   }
   const user = await response.json();
   return user as { api_key: string; name: string };

src/common.ts

@@ -13,8 +13,7 @@ import {
 } from '../proto/exa/language_server_pb/language_server_pb';
 
 const EXTENSION_NAME = 'chrome';
-const EXTENSION_VERSION = '1.2.18';
-const BASE_URL = 'https://server.codeium.com';
+const EXTENSION_VERSION = '1.2.26';
 
 export const CODEIUM_DEBUG = false;
 
@@ -27,9 +26,9 @@ async function getApiKey(extensionId: string): Promise<string | undefined> {
   return user?.apiKey;
 }
 
-function languageServerClient(): PromiseClient<typeof LanguageServerService> {
+function languageServerClient(baseUrl: string): PromiseClient<typeof LanguageServerService> {
   const transport = createConnectTransport({
-    baseUrl: BASE_URL,
+    baseUrl,
     useBinaryFormat: true,
   });
   return createPromiseClient(LanguageServerService, transport);
@@ -53,10 +52,15 @@ export interface IdeInfo {
 }
 
 export class LanguageServerServiceWorkerClient {
-  client = languageServerClient();
+  // Note that the URL won't refresh post-initialization.
+  client: Promise<PromiseClient<typeof LanguageServerService>>;
   private abortController?: AbortController;
 
-  constructor(readonly sessionId: string) {}
+  constructor(baseUrlPromise: Promise<string>, private readonly sessionId: string) {
+    this.client = (async (): Promise<PromiseClient<typeof LanguageServerService>> => {
+      return languageServerClient(await baseUrlPromise);
+    })();
+  }
 
   getHeaders(apiKey: string | undefined): Record<string, string> {
     if (apiKey === undefined) {
@@ -72,7 +76,7 @@ export class LanguageServerServiceWorkerClient {
     this.abortController?.abort();
     this.abortController = new AbortController();
     const signal = this.abortController.signal;
-    const getCompletionsPromise = this.client.getCompletions(request, {
+    const getCompletionsPromise = (await this.client).getCompletions(request, {
       signal,
       headers: this.getHeaders(request.metadata?.apiKey),
     });
@@ -105,7 +109,9 @@ export class LanguageServerServiceWorkerClient {
     acceptCompletionRequest: PartialMessage<AcceptCompletionRequest>
   ): Promise<void> {
     try {
-      await this.client.acceptCompletion(acceptCompletionRequest, {
+      await (
+        await this.client
+      ).acceptCompletion(acceptCompletionRequest, {
         headers: this.getHeaders(acceptCompletionRequest.metadata?.apiKey),
       });
     } catch (err) {

src/component/Options.tsx

@@ -9,8 +9,13 @@ import Divider from '@mui/material/Divider';
 import React, { createRef, useEffect, useState } from 'react';
 import { v4 as uuidv4 } from 'uuid';
 
-import { PROFILE_URL } from '../auth';
-import { computeAllowlist, defaultAllowlist, getStorageItem, setStorageItem } from '../storage';
+import {
+  computeAllowlist,
+  defaultAllowlist,
+  getGeneralProfileUrl,
+  getStorageItem,
+  setStorageItem,
+} from '../storage';
 
 const EditableList = () => {
   const [text, setText] = useState('');
@@ -118,6 +123,7 @@ const EditableList = () => {
   );
 };
 
+const profileUrl = getGeneralProfileUrl();
 const openTokenPage = async () => {
   const params = new URLSearchParams({
     response_type: 'token',
@@ -127,11 +133,18 @@ const openTokenPage = async () => {
     redirect_parameters_type: 'query',
     state: uuidv4(),
   });
-  await chrome.tabs.create({ url: `${PROFILE_URL}?${params}` });
+  await chrome.tabs.create({ url: `${await profileUrl}?${params}` });
 };
 
 const Options = () => {
-  const ref = createRef<HTMLInputElement>();
+  const tokenRef = createRef<HTMLInputElement>();
+  const portalUrlRef = createRef<HTMLInputElement>();
+  const [portalUrlText, setPortalUrlText] = useState('');
+  useEffect(() => {
+    (async () => {
+      setPortalUrlText((await getStorageItem('portalUrl')) ?? '');
+    })();
+  }, []);
   return (
     <Box sx={{ width: '100%', maxWidth: 400, bgcolor: 'background.paper' }}>
       <Typography variant="body2">
@@ -161,14 +174,14 @@ const Options = () => {
         }}
       />
       <Box sx={{ my: 2, mx: 2 }}>
-        <Typography variant="h6"> Alternative Ways to Log in </Typography>
+        <Typography variant="h6"> Alternative ways to log in </Typography>
         <TextField
           id="token"
           label="Token"
           variant="standard"
           fullWidth
           type="password"
-          inputRef={ref}
+          inputRef={tokenRef}
         />
         <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
           <Button variant="text" onClick={openTokenPage} sx={{ textTransform: 'none' }}>
@@ -179,7 +192,7 @@ const Options = () => {
             variant="text"
             onClick={async () => {
               // get token from input
-              const token = ref.current?.value;
+              const token = tokenRef.current?.value;
               await chrome.runtime.sendMessage({ type: 'manual', token: token });
             }}
             sx={{ textTransform: 'none' }}
@@ -193,6 +206,37 @@ const Options = () => {
           padding: '0.5em',
         }}
       />
+      <Box sx={{ my: 2, mx: 2 }}>
+        <Typography variant="h6"> Enterprise settings </Typography>
+        <TextField
+          id="portal"
+          label="Portal URL"
+          variant="standard"
+          fullWidth
+          type="url"
+          inputRef={portalUrlRef}
+          value={portalUrlText}
+          onChange={(e) => setPortalUrlText(e.target.value)}
+        />
+        <Box sx={{ display: 'flex', justifyContent: 'flex-end' }}>
+          <Button
+            variant="text"
+            onClick={async () => {
+              const portalUrl = portalUrlRef.current?.value;
+              console.log(portalUrl);
+              await setStorageItem('portalUrl', portalUrl);
+            }}
+            sx={{ textTransform: 'none' }}
+          >
+            Enter Portal URL <LoginIcon />
+          </Button>
+        </Box>
+      </Box>
+      <Divider
+        sx={{
+          padding: '0.5em',
+        }}
+      />
       <Box sx={{ my: 2, mx: 2 }}>
         <EditableList />
       </Box>

src/popup.ts

@@ -14,6 +14,19 @@ getStorageItem('user')
     const usernameP = document.getElementById('username');
     if (usernameP !== null && user !== undefined) {
       usernameP.textContent = `Welcome, ${user.name}`;
+      if (user.userPortalUrl !== undefined && user.userPortalUrl !== '') {
+        const br = document.createElement('br');
+        usernameP.appendChild(br);
+        const a = document.createElement('a');
+        const linkText = document.createTextNode('Portal');
+        a.appendChild(linkText);
+        a.title = 'Portal';
+        a.href = user.userPortalUrl;
+        a.addEventListener('click', () => {
+          chrome.tabs.create({ url: user.userPortalUrl });
+        });
+        usernameP.appendChild(a);
+      }
     }
   })
   .catch((error) => {

src/serviceWorker.ts

@@ -9,6 +9,7 @@ import {
 import { loggedIn, loggedOut, unhealthy } from './shared';
 import {
   defaultAllowlist,
+  getGeneralPortalUrl,
   getStorageItem,
   initializeStorageWithDefaults,
   setStorageItem,
@@ -36,8 +37,15 @@ chrome.runtime.onInstalled.addListener(async () => {
     // Inline the code for openAuthTab() because we can't invoke sendMessage.
     const uuid = uuidv4();
     authStates.push(uuid);
+    const portalUrl = await (async (): Promise<string> => {
+      const url = await getGeneralPortalUrl();
+      if (url === undefined) {
+        return 'https://www.codeium.com';
+      }
+      return url;
+    })();
     await chrome.tabs.create({
-      url: `https://www.codeium.com/profile?redirect_uri=chrome-extension://${chrome.runtime.id}&state=${uuid}`,
+      url: `${portalUrl}/profile?redirect_uri=chrome-extension://${chrome.runtime.id}&state=${uuid}`,
     });
   } else {
     await loggedIn();
@@ -81,7 +89,6 @@ chrome.runtime.onMessageExternal.addListener(async (message, sender, sendRespons
     return;
   }
   authStates.splice(stateIndex, 1);
-  console.log('Obtained token');
   await login(typedMessage.token);
 });
 
@@ -93,22 +100,27 @@ chrome.runtime.onStartup.addListener(async () => {
   }
 });
 
-chrome.runtime.onMessage.addListener(async (message) => {
+chrome.runtime.onMessage.addListener((message) => {
   // TODO(prem): Strongly type this.
   if (message.type === 'state') {
     const payload = message.payload as { state: string };
     authStates.push(payload.state);
   } else if (message.type === 'manual') {
-    await login(message.token);
+    login(message.token);
   } else {
     console.log('Unrecognized message:', message);
   }
 });
 
 const clientMap = new Map<string, LanguageServerServiceWorkerClient>();
 
+// TODO(prem): Is it safe to make this listener async to simplify the LanguageServerServiceWorkerClient constructor?
 chrome.runtime.onConnectExternal.addListener((port) => {
-  clientMap.set(port.name, new LanguageServerServiceWorkerClient(port.name));
+  // TODO(prem): Technically this URL isn't synchronized with the user/API key.
+  clientMap.set(
+    port.name,
+    new LanguageServerServiceWorkerClient(getLanguageServerUrl(), port.name)
+  );
   port.onDisconnect.addListener((port) => {
     clientMap.delete(port.name);
   });
@@ -134,8 +146,13 @@ chrome.runtime.onConnectExternal.addListener((port) => {
 
 async function login(token: string) {
   try {
-    const user = await registerUser(token);
-    await setStorageItem('user', { apiKey: user.api_key, name: user.name });
+    const portalUrl = await getGeneralPortalUrl();
+    const user = await registerUser(token, portalUrl);
+    await setStorageItem('user', {
+      apiKey: user.api_key,
+      name: user.name,
+      userPortalUrl: portalUrl,
+    });
     await loggedIn();
     // TODO(prem): Open popup.
     // https://github.com/GoogleChrome/developer.chrome.com/issues/2602
@@ -144,3 +161,12 @@ async function login(token: string) {
     console.log(error);
   }
 }
+
+async function getLanguageServerUrl(): Promise<string> {
+  const user = await getStorageItem('user');
+  const userPortalUrl = user?.userPortalUrl;
+  if (userPortalUrl === undefined || userPortalUrl === '') {
+    return 'https://server.codeium.com';
+  }
+  return `${userPortalUrl}/_route/language_server`;
+}