PoC Extension

extension/app/background.ts

@@ -1,3 +1,4 @@
+import { ChromeStorageService } from "@extension/chrome/storage";
 import type { PendingUpdate } from "@extension/lib/storage";
 import { getStoredUser, savePendingUpdate } from "@extension/lib/storage";
 
@@ -41,7 +42,7 @@ chrome.runtime.onUpdateAvailable.addListener(async (details) => {
     version: details.version,
     detectedAt: Date.now(),
   };
-  await savePendingUpdate(pendingUpdate);
+  await savePendingUpdate(new ChromeStorageService(), pendingUpdate);
 });
 
 /**
@@ -78,7 +79,7 @@ const shouldDisableContextMenuForDomain = async (
     return true;
   }
 
-  const user = await getStoredUser();
+  const user = await getStoredUser(new ChromeStorageService());
   if (!user || !user.selectedWorkspace) {
     return false;
   }
@@ -620,6 +621,7 @@ const exchangeCodeForTokens = async (
  * Logout the user from Auth0.
  */
 const logout = (sendResponse: (response: AuthBackgroundResponse) => void) => {
+  console.log("logout");
   const redirectUri = chrome.identity.getRedirectURL();
   const logoutUrl = `https://${AUTH0_CLIENT_DOMAIN}/v2/logout?client_id=${AUTH0_CLIENT_ID}&returnTo=${encodeURIComponent(redirectUri)}`;
 

extension/app/src/chrome/components/AttachButtons.tsx

@@ -0,0 +1,107 @@
+import { Button, CameraIcon, DocumentPlusIcon } from "@dust-tt/sparkle";
+import { useCurrentUrlAndDomain } from "@extension/hooks/useCurrentDomain";
+import type { AttachButtonProps } from "@extension/shared/services/platform";
+
+export const ChromeAttachButtons = ({
+  isBlinking,
+  isLoading,
+  fileUploaderService,
+  owner,
+}: AttachButtonProps) => {
+  // Blacklisting logic to disable share buttons.
+  const { currentDomain, currentUrl } = useCurrentUrlAndDomain();
+  const blacklistedConfig: string[] = owner.blacklistedDomains ?? [];
+
+  const isBlacklisted =
+    currentDomain === "chrome" ||
+    blacklistedConfig.some((d) =>
+      d.startsWith("http://") || d.startsWith("https://")
+        ? currentUrl.startsWith(d)
+        : currentDomain.endsWith(d)
+    );
+
+  return (
+    <>
+      <div className="block sm:hidden">
+        <Button
+          icon={DocumentPlusIcon}
+          tooltip={
+            !isBlacklisted
+              ? "Attach text from page"
+              : "Attachment disabled on this website"
+          }
+          variant="outline"
+          size="sm"
+          className={isBlinking ? "animate-[bgblink_200ms_3]" : ""}
+          onClick={() =>
+            fileUploaderService.uploadContentTab({
+              includeContent: true,
+              includeCapture: false,
+            })
+          }
+          disabled={isLoading || isBlacklisted}
+        />
+      </div>
+      <div className="block sm:hidden">
+        <Button
+          icon={CameraIcon}
+          tooltip={
+            !isBlacklisted
+              ? "Attach page screenshot"
+              : "Attachment disabled on this website"
+          }
+          variant="outline"
+          size="sm"
+          onClick={() =>
+            fileUploaderService.uploadContentTab({
+              includeContent: false,
+              includeCapture: true,
+            })
+          }
+          disabled={isLoading || isBlacklisted}
+        />
+      </div>
+      <div className="hidden sm:block">
+        <Button
+          icon={DocumentPlusIcon}
+          label="Add page text"
+          tooltip={
+            !isBlacklisted
+              ? "Attach text from page"
+              : "Attachment disabled on this website"
+          }
+          variant="outline"
+          size="sm"
+          className={isBlinking ? "animate-[bgblink_200ms_3]" : ""}
+          onClick={() =>
+            fileUploaderService.uploadContentTab({
+              includeContent: true,
+              includeCapture: false,
+            })
+          }
+          disabled={isLoading || isBlacklisted}
+        />
+      </div>
+      <div className="hidden sm:block">
+        <Button
+          icon={CameraIcon}
+          label="Add page screenshot"
+          tooltip={
+            !isBlacklisted
+              ? "Attach page screenshot"
+              : "Attachment disabled on this website"
+          }
+          variant="outline"
+          size="sm"
+          onClick={() =>
+            fileUploaderService.uploadContentTab({
+              includeContent: false,
+              includeCapture: true,
+            })
+          }
+          disabled={isLoading || isBlacklisted}
+        />
+      </div>
+    </>
+  );
+};

extension/app/src/chrome/platform.ts

@@ -0,0 +1,15 @@
+import { ChromeAttachButtons } from "@extension/chrome/components/AttachButtons";
+import { ChromeAuth } from "@extension/chrome/services/auth";
+import { ChromeCaptureService } from "@extension/chrome/services/capture";
+import { ChromeStorageService } from "@extension/chrome/storage";
+import type { PlatformService } from "@extension/shared/services/platform";
+
+export const chromePlatform: PlatformService = {
+  platform: "chrome",
+  auth: new ChromeAuth(),
+  storage: new ChromeStorageService(),
+  components: {
+    AttachButtons: ChromeAttachButtons,
+  },
+  capture: new ChromeCaptureService(),
+};

extension/app/src/chrome/services/auth.ts

@@ -0,0 +1,177 @@
+import type { Result } from "@dust-tt/client";
+import { Err, Ok } from "@dust-tt/client";
+import { ChromeStorageService } from "@extension/chrome/storage";
+import {
+  AUTH0_CLAIM_NAMESPACE,
+  DEFAULT_DUST_API_DOMAIN,
+  DUST_EU_URL,
+  DUST_US_URL,
+} from "@extension/lib/config";
+import {
+  sendAuthMessage,
+  sendRefreshTokenMessage,
+  sentLogoutMessage,
+} from "@extension/lib/messages";
+import type {
+  StoredTokens,
+  UserTypeWithExtensionWorkspaces,
+} from "@extension/lib/storage";
+import {
+  clearStoredData,
+  getStoredTokens,
+  saveTokens,
+  saveUser,
+} from "@extension/lib/storage";
+import type { AuthService } from "@extension/shared/services/auth";
+import {
+  AuthError,
+  makeEnterpriseConnectionName,
+} from "@extension/shared/services/auth";
+import { jwtDecode } from "jwt-decode";
+
+const REGIONS = ["europe-west1", "us-central1"] as const;
+export type RegionType = (typeof REGIONS)[number];
+
+const isRegionType = (region: string): region is RegionType =>
+  REGIONS.includes(region as RegionType);
+
+const REGION_CLAIM = `${AUTH0_CLAIM_NAMESPACE}region`;
+const CONNECTION_STRATEGY_CLAIM = `${AUTH0_CLAIM_NAMESPACE}connection.strategy`;
+const WORKSPACE_ID_CLAIM = `${AUTH0_CLAIM_NAMESPACE}workspaceId`;
+
+const DOMAIN_FOR_REGION: Record<RegionType, string> = {
+  "us-central1": DUST_US_URL,
+  "europe-west1": DUST_EU_URL,
+};
+
+const log = console.error;
+
+export class ChromeAuth implements AuthService {
+  constructor() {}
+
+  // Login sends a message to the background script to call the auth0 login endpoint.
+  // It saves the tokens in the extension and schedules a token refresh.
+  // Then it calls the /me route to get the user info.
+  async login(isForceLogin?: boolean, forcedConnection?: string) {
+    try {
+      const response = await sendAuthMessage(isForceLogin, forcedConnection);
+      if (!response.accessToken) {
+        throw new Error("No access token received.");
+      }
+      const tokens = await saveTokens(new ChromeStorageService(), response);
+
+      const claims = jwtDecode<Record<string, string>>(tokens.accessToken);
+
+      const dustDomain = getDustDomain(claims);
+      const connectionDetails = getConnectionDetails(claims);
+
+      const res = await fetchMe(tokens.accessToken, dustDomain);
+      if (res.isErr()) {
+        return res;
+      }
+      const workspaces = res.value.user.workspaces;
+      const user = await saveUser(new ChromeStorageService(), {
+        ...res.value.user,
+        ...connectionDetails,
+        dustDomain,
+        selectedWorkspace: workspaces.length === 1 ? workspaces[0].sId : null,
+      });
+      return new Ok({ tokens, user });
+    } catch (error) {
+      return new Err(new AuthError("not_authenticated", error?.toString()));
+    }
+  }
+
+  // Logout sends a message to the background script to call the auth0 logout endpoint.
+  // It also clears the stored tokens in the extension.
+  async logout(): Promise<boolean> {
+    try {
+      const response = await sentLogoutMessage();
+      if (!response?.success) {
+        throw new Error("No success response received.");
+      }
+      return true;
+    } catch (error) {
+      log("Logout failed: Unknown error.", error);
+      return false;
+    } finally {
+      await clearStoredData();
+    }
+  }
+
+  // Refresh token sends a message to the background script to call the auth0 refresh token endpoint.
+  // It updates the stored tokens with the new access token.
+  // If the refresh token is invalid, it will call handleLogout.
+  async refreshToken(
+    tokens?: StoredTokens | null
+  ): Promise<Result<StoredTokens, AuthError>> {
+    try {
+      tokens = tokens ?? (await getStoredTokens(new ChromeStorageService()));
+      if (!tokens) {
+        return new Err(new AuthError("not_authenticated", "No tokens found."));
+      }
+      const response = await sendRefreshTokenMessage(
+        new ChromeStorageService(),
+        tokens.refreshToken
+      );
+      if (!response?.accessToken) {
+        return new Err(
+          new AuthError("not_authenticated", "No access token received")
+        );
+      }
+      return new Ok(await saveTokens(new ChromeStorageService(), response));
+    } catch (error) {
+      log("Refresh token: unknown error.", error);
+      return new Err(new AuthError("not_authenticated", error?.toString()));
+    }
+  }
+
+  async getAccessToken(): Promise<string | null> {
+    let tokens = await getStoredTokens(new ChromeStorageService());
+    if (!tokens || !tokens.accessToken || tokens.expiresAt < Date.now()) {
+      const refreshRes = await this.refreshToken(tokens);
+      if (refreshRes.isOk()) {
+        tokens = refreshRes.value;
+      }
+    }
+
+    return tokens?.accessToken ?? null;
+  }
+}
+
+const getDustDomain = (claims: Record<string, string>) => {
+  const region = claims[REGION_CLAIM];
+
+  return (
+    (isRegionType(region) && DOMAIN_FOR_REGION[region]) ||
+    DEFAULT_DUST_API_DOMAIN
+  );
+};
+
+const getConnectionDetails = (claims: Record<string, string>) => {
+  const connectionStrategy = claims[CONNECTION_STRATEGY_CLAIM];
+  const ws = claims[WORKSPACE_ID_CLAIM];
+  return {
+    connectionStrategy,
+    connection: ws ? makeEnterpriseConnectionName(ws) : undefined,
+  };
+};
+
+// Fetch me sends a request to the /me route to get the user info.
+const fetchMe = async (
+  accessToken: string,
+  dustDomain: string
+): Promise<Result<{ user: UserTypeWithExtensionWorkspaces }, AuthError>> => {
+  const response = await fetch(`${dustDomain}/api/v1/me`, {
+    headers: {
+      Authorization: `Bearer ${accessToken}`,
+      "X-Request-Origin": "extension",
+    },
+  });
+  const me = await response.json();
+
+  if (!response.ok) {
+    return new Err(new AuthError(me.error.type, me.error.message));
+  }
+  return new Ok(me);
+};

extension/app/src/chrome/services/capture.ts

@@ -0,0 +1,26 @@
+import { assertNever } from "@dust-tt/client";
+import { getIncludeCurrentTab } from "@extension/lib/conversation";
+import type { GetActiveTabOptions } from "@extension/lib/messages";
+import type {
+  CaptureOperationId,
+  CaptureService,
+} from "@extension/shared/services/capture";
+
+export class ChromeCaptureService implements CaptureService {
+  isOperationSupported(id: CaptureOperationId) {
+    return ["capture-page-content"].includes(id);
+  }
+
+  async handleOperation(id: CaptureOperationId, options: GetActiveTabOptions) {
+    switch (id) {
+      case "capture-page-content":
+        return getIncludeCurrentTab(options);
+
+      case "capture-screenshot":
+        throw new Error("Not implemented");
+
+      default:
+        assertNever(id);
+    }
+  }
+}

extension/app/src/chrome/storage.ts

@@ -0,0 +1,21 @@
+import type { StorageService } from "@extension/shared/interfaces/storage";
+
+export class ChromeStorageService implements StorageService {
+  async get<T>(key: string | string[]): Promise<T | null> {
+    if (Array.isArray(key)) {
+      const result = await chrome.storage.local.get(key);
+      return result as T;
+    }
+
+    const result = await chrome.storage.local.get([key]);
+    return result[key] ?? null;
+  }
+
+  async set<T>(key: string, value: T): Promise<void> {
+    await chrome.storage.local.set({ [key]: value });
+  }
+
+  async remove(key: string): Promise<void> {
+    await chrome.storage.local.remove(key);
+  }
+}