Gate Documents UI by RAG plan and add Drive connect UX (#81-#85)

Hide Documents without features.rag, polish upload/active toggle, add personal
and company-manager Drive connect sections, and show upgrade messaging when
chat RAG is denied by plan.

Companion to chat_backend#42
Closes #81 #82 #83 #84 #85

llm-fe/src/llm-fe/auth/sso.js

@@ -6,7 +6,7 @@
 /**
  * Absolute backend OAuth start URL.
  * @param {'google'|'microsoft'} provider
- * @param {'login'|'signup'} intent
+ * @param {'login'|'signup'|'link_drive'|'link_company_drive'} intent
  * @returns {string}
  */
 export function oauthStartUrl(provider, intent = 'login') {
@@ -18,7 +18,7 @@ export function oauthStartUrl(provider, intent = 'login') {
 /**
  * Begin browser redirect to IdP via backend.
  * @param {'google'|'microsoft'} provider
- * @param {'login'|'signup'} intent
+ * @param {'login'|'signup'|'link_drive'|'link_company_drive'} intent
  */
 export function startOAuth(provider, intent = 'login') {
   window.location.assign(oauthStartUrl(provider, intent));

llm-fe/src/llm-fe/components/ConversationDetailCard/ConversationDetailCard.test.tsx

@@ -1,5 +1,6 @@
 import React from 'react';
 import { render, screen } from '@testing-library/react';
+import { MemoryRouter } from 'react-router-dom';
 import { ThemeProvider } from 'styled-components';
 import ConversationDetailCard from './ConversationDetailCard';
 
@@ -19,9 +20,11 @@ const renderCard = (props: {
   user_created: boolean;
 }) =>
   render(
-    <ThemeProvider theme={theme as never}>
+    <MemoryRouter>
-      <ConversationDetailCard {...props} />
+      <ThemeProvider theme={theme as never}>
-    </ThemeProvider>
+        <ConversationDetailCard {...props} />
+      </ThemeProvider>
+    </MemoryRouter>
   );
 
 describe('ConversationDetailCard', () => {
@@ -43,4 +46,38 @@ describe('ConversationDetailCard', () => {
     expect(screen.getByText('Hi')).toBeInTheDocument();
     expect(screen.queryByTestId('message-token-usage')).not.toBeInTheDocument();
   });
+
+  it('shows a friendly upgrade bubble for RAG feature_not_allowed errors (#85)', () => {
+    renderCard({
+      message: JSON.stringify({
+        type: 'error',
+        code: 'feature_not_allowed',
+        content: 'Your plan does not include rag document search.',
+        details: { feature: 'rag' },
+      }),
+      user_created: false,
+    });
+
+    expect(
+      screen.getByText('Your plan does not include rag document search.')
+    ).toBeInTheDocument();
+    const upgradeLink = screen.getByRole('link', { name: /upgrade your plan/i });
+    expect(upgradeLink).toBeInTheDocument();
+    expect(upgradeLink).toHaveAttribute('href', '/account/');
+  });
+
+  it('renders unrelated feature_not_allowed errors as a normal inline error, not the upgrade bubble', () => {
+    renderCard({
+      message: JSON.stringify({
+        type: 'error',
+        code: 'feature_not_allowed',
+        content: 'Your plan does not include image generation.',
+        details: { feature: 'image_generation' },
+      }),
+      user_created: false,
+    });
+
+    expect(screen.queryByRole('link', { name: /upgrade your plan/i })).not.toBeInTheDocument();
+    expect(screen.getByText(/image generation/i)).toBeInTheDocument();
+  });
 });

llm-fe/src/llm-fe/components/ConversationDetailCard/ConversationDetailCard.tsx

@@ -1,6 +1,8 @@
 import React from "react";
 import Markdown from "markdown-to-jsx";
+import { Link } from "react-router-dom";
 import styled, { keyframes } from "styled-components";
+import { isRagFeatureNotAllowed, parseChatErrorPayload } from "../../utils/chatErrors";
 
 const fadeIn = keyframes`
   from { opacity: 0; transform: translateY(10px); }
@@ -91,6 +93,27 @@ const LoadingContainer = styled.div`
   padding: 0.5rem;
 `;
 
+const UpgradeNotice = styled.div`
+  display: flex;
+  flex-direction: column;
+  align-items: flex-start;
+  gap: 0.65rem;
+`;
+
+const UpgradeLink = styled(Link)`
+  background: ${(props) => props.theme.main};
+  color: #fff;
+  text-decoration: none;
+  font-weight: 600;
+  font-size: 0.85rem;
+  padding: 0.5rem 1rem;
+  border-radius: 0.5rem;
+
+  &:hover {
+    opacity: 0.9;
+  }
+`;
+
 type ConversationDetailCardProps = {
   message: string;
   user_created: boolean;
@@ -133,6 +156,23 @@ const ConversationDetailCard = ({
     );
   }
 
+  const errorPayload = parseChatErrorPayload(message);
+  if (errorPayload && isRagFeatureNotAllowed(errorPayload)) {
+    return (
+      <MessageContainer $isUser={false}>
+        <Bubble $isUser={false}>
+          <UpgradeNotice>
+            <span>
+              {errorPayload.content ||
+                "Document search (RAG) isn't included in your current plan."}
+            </span>
+            <UpgradeLink to="/account/">Upgrade your plan</UpgradeLink>
+          </UpgradeNotice>
+        </Bubble>
+      </MessageContainer>
+    );
+  }
+
   let contentToAdd = message;
   try {
     const parsedMessage = JSON.parse(message);

llm-fe/src/llm-fe/components/DriveConnectionsSection/DriveConnectionsSection.tsx

@@ -0,0 +1,379 @@
+import React, { useCallback, useEffect, useState } from "react";
+import styled from "styled-components";
+import {
+  DriveConnectIntent,
+  DriveConnectionKind,
+  DriveConnectionType,
+  DriveProvider,
+  connectDrive,
+  disconnectDriveConnection,
+  fetchDriveConnections,
+  parseResourceIdsInput,
+  saveDriveResourceSelection,
+  syncDriveConnection,
+} from "../../utils/drive";
+
+const Section = styled.div`
+  background: ${({ theme }) => theme.colors.cardBackground};
+  backdrop-filter: blur(10px);
+  border: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  border-radius: 1rem;
+  padding: 2rem;
+  width: 100%;
+  max-width: 1000px;
+  margin-bottom: 2rem;
+  box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.37);
+`;
+
+const SectionTitle = styled.h2`
+  font-size: 1.8rem;
+  margin-bottom: 0.5rem;
+  color: ${({ theme }) => theme.colors.text};
+  border-bottom: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  padding-bottom: 1rem;
+`;
+
+const SectionDescription = styled.p`
+  color: ${({ theme }) => theme.colors.text};
+  opacity: 0.7;
+  margin: 0.75rem 0 1.5rem 0;
+  line-height: 1.5;
+`;
+
+const ConnectButtonRow = styled.div`
+  display: flex;
+  flex-wrap: wrap;
+  gap: 1rem;
+  margin-bottom: 1.5rem;
+`;
+
+const ConnectButton = styled.button`
+  background: ${({ theme }) => theme.main};
+  border: none;
+  border-radius: 0.5rem;
+  color: #fff;
+  padding: 0.7rem 1.25rem;
+  font-weight: 600;
+  cursor: pointer;
+  transition: transform 0.2s ease, box-shadow 0.2s ease;
+
+  &:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 4px 12px ${({ theme }) => theme.main}66;
+  }
+`;
+
+const ConnectionList = styled.div`
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+`;
+
+const ConnectionCard = styled.div`
+  border: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  border-radius: 0.75rem;
+  padding: 1.25rem;
+`;
+
+const ConnectionHeader = styled.div`
+  display: flex;
+  justify-content: space-between;
+  align-items: flex-start;
+  gap: 1rem;
+  flex-wrap: wrap;
+`;
+
+const ConnectionTitle = styled.strong`
+  color: ${({ theme }) => theme.colors.text};
+  font-size: 1.05rem;
+`;
+
+const ConnectionMeta = styled.p`
+  margin: 0.35rem 0 0 0;
+  font-size: 0.85rem;
+  color: ${({ theme }) => theme.colors.text};
+  opacity: 0.65;
+`;
+
+const ConnectionActions = styled.div`
+  display: flex;
+  gap: 0.5rem;
+  flex-wrap: wrap;
+`;
+
+const SmallButton = styled.button`
+  background: transparent;
+  border: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  border-radius: 0.5rem;
+  color: ${({ theme }) => theme.colors.text};
+  padding: 0.45rem 0.9rem;
+  font-weight: 600;
+  font-size: 0.85rem;
+  cursor: pointer;
+
+  &:hover {
+    background: ${({ theme }) => theme.darkMode ? 'rgba(255, 255, 255, 0.1)' : 'rgba(0, 0, 0, 0.1)'};
+  }
+
+  &:disabled {
+    opacity: 0.5;
+    cursor: not-allowed;
+  }
+`;
+
+const DangerButton = styled(SmallButton)`
+  color: #ff6b6b;
+  border-color: rgba(255, 107, 107, 0.4);
+
+  &:hover {
+    background: rgba(255, 71, 87, 0.1);
+  }
+`;
+
+const ResourceForm = styled.div`
+  margin-top: 1rem;
+  display: flex;
+  gap: 0.5rem;
+  flex-wrap: wrap;
+`;
+
+const ResourceInput = styled.input`
+  flex: 1;
+  min-width: 200px;
+  background: ${({ theme }) => theme.darkMode ? 'rgba(255, 255, 255, 0.05)' : 'rgba(0, 0, 0, 0.05)'};
+  border: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  border-radius: 0.5rem;
+  padding: 0.6rem 0.8rem;
+  color: ${({ theme }) => theme.colors.text};
+
+  &::placeholder {
+    color: ${({ theme }) => theme.darkMode ? 'rgba(255, 255, 255, 0.4)' : 'rgba(0, 0, 0, 0.4)'};
+  }
+`;
+
+const ResourceTagList = styled.div`
+  margin-top: 0.75rem;
+  display: flex;
+  flex-wrap: wrap;
+  gap: 0.4rem;
+`;
+
+const ResourceTag = styled.span`
+  background: ${({ theme }) => theme.darkMode ? 'rgba(255, 255, 255, 0.1)' : 'rgba(0, 0, 0, 0.08)'};
+  color: ${({ theme }) => theme.colors.text};
+  border-radius: 999px;
+  padding: 0.25rem 0.75rem;
+  font-size: 0.8rem;
+`;
+
+const EmptyState = styled.p`
+  color: ${({ theme }) => theme.colors.text};
+  opacity: 0.6;
+`;
+
+const PROVIDER_LABELS: Record<DriveProvider, string> = {
+  google: "Google Drive",
+  microsoft: "OneDrive",
+};
+
+type PendingAction = 'sync' | 'save' | 'disconnect';
+
+type DriveConnectionsSectionProps = {
+  kind: DriveConnectionKind;
+  title: string;
+  description?: string;
+  connectIntent: DriveConnectIntent;
+};
+
+const DriveConnectionsSection = ({
+  kind,
+  title,
+  description,
+  connectIntent,
+}: DriveConnectionsSectionProps): JSX.Element => {
+  const [connections, setConnections] = useState<DriveConnectionType[]>([]);
+  const [loading, setLoading] = useState(true);
+  const [loadError, setLoadError] = useState(false);
+  const [connectError, setConnectError] = useState<string | null>(null);
+  const [resourceInputs, setResourceInputs] = useState<Record<number, string>>({});
+  const [pendingAction, setPendingAction] = useState<Record<number, PendingAction | undefined>>({});
+
+  const loadConnections = useCallback(async () => {
+    setLoading(true);
+    setLoadError(false);
+    try {
+      const all = await fetchDriveConnections();
+      setConnections(all.filter((conn) => (conn.kind || 'personal') === kind));
+    } catch {
+      setLoadError(true);
+    } finally {
+      setLoading(false);
+    }
+  }, [kind]);
+
+  useEffect(() => {
+    loadConnections();
+  }, [loadConnections]);
+
+  const setAction = (id: number, action?: PendingAction) => {
+    setPendingAction((prev) => ({ ...prev, [id]: action }));
+  };
+
+  const handleDisconnect = async (id: number) => {
+    setAction(id, 'disconnect');
+    try {
+      await disconnectDriveConnection(id);
+      setConnections((prev) => prev.filter((conn) => conn.id !== id));
+    } catch (err) {
+      console.log(err);
+    } finally {
+      setAction(id, undefined);
+    }
+  };
+
+  const handleSync = async (id: number) => {
+    setAction(id, 'sync');
+    try {
+      await syncDriveConnection(id);
+      await loadConnections();
+    } catch (err) {
+      console.log(err);
+    } finally {
+      setAction(id, undefined);
+    }
+  };
+
+  const handleSaveResources = async (id: number) => {
+    const resourceIds = parseResourceIdsInput(resourceInputs[id] || '');
+    if (resourceIds.length === 0) {
+      return;
+    }
+    setAction(id, 'save');
+    try {
+      const updated = await saveDriveResourceSelection(id, {
+        resource_ids: resourceIds,
+        resource_labels: resourceIds,
+      });
+      setConnections((prev) => prev.map((conn) => (conn.id === id ? { ...conn, ...updated } : conn)));
+      setResourceInputs((prev) => ({ ...prev, [id]: '' }));
+    } catch (err) {
+      console.log(err);
+    } finally {
+      setAction(id, undefined);
+    }
+  };
+
+  return (
+    <Section>
+      <SectionTitle>{title}</SectionTitle>
+      {description && <SectionDescription>{description}</SectionDescription>}
+
+      <ConnectButtonRow>
+        <ConnectButton
+          type="button"
+          onClick={() => {
+            setConnectError(null);
+            void connectDrive('google', connectIntent).catch((err) => {
+              console.error(err);
+              setConnectError('Could not start Google Drive connect. Try again.');
+            });
+          }}
+        >
+          Connect Google Drive
+        </ConnectButton>
+        <ConnectButton
+          type="button"
+          onClick={() => {
+            setConnectError(null);
+            void connectDrive('microsoft', connectIntent).catch((err) => {
+              console.error(err);
+              setConnectError('Could not start OneDrive connect. Try again.');
+            });
+          }}
+        >
+          Connect OneDrive
+        </ConnectButton>
+      </ConnectButtonRow>
+
+      {connectError && <EmptyState>{connectError}</EmptyState>}
+
+      {loading ? (
+        <EmptyState>Loading connections…</EmptyState>
+      ) : loadError ? (
+        <EmptyState>Could not load drive connections. Try again later.</EmptyState>
+      ) : connections.length === 0 ? (
+        <EmptyState>No drives connected yet.</EmptyState>
+      ) : (
+        <ConnectionList>
+          {connections.map((conn) => {
+            const action = pendingAction[conn.id];
+            const selectedLabels = conn.selected_resource_labels?.length
+              ? conn.selected_resource_labels
+              : conn.selected_resource_ids || [];
+
+            return (
+              <ConnectionCard key={conn.id}>
+                <ConnectionHeader>
+                  <div>
+                    <ConnectionTitle>{PROVIDER_LABELS[conn.provider] || conn.provider}</ConnectionTitle>
+                    <ConnectionMeta>
+                      {conn.external_account_email || 'Connected account'}
+                      {conn.last_sync_status ? ` · ${conn.last_sync_status}` : ''}
+                      {conn.last_sync_at
+                        ? ` · Last synced ${new Date(conn.last_sync_at).toLocaleString()}`
+                        : ''}
+                    </ConnectionMeta>
+                  </div>
+                  <ConnectionActions>
+                    <SmallButton
+                      type="button"
+                      onClick={() => handleSync(conn.id)}
+                      disabled={Boolean(action)}
+                    >
+                      {action === 'sync' ? 'Syncing…' : 'Sync now'}
+                    </SmallButton>
+                    <DangerButton
+                      type="button"
+                      onClick={() => handleDisconnect(conn.id)}
+                      disabled={Boolean(action)}
+                    >
+                      {action === 'disconnect' ? 'Disconnecting…' : 'Disconnect'}
+                    </DangerButton>
+                  </ConnectionActions>
+                </ConnectionHeader>
+
+                {selectedLabels.length > 0 && (
+                  <ResourceTagList>
+                    {selectedLabels.map((label, idx) => (
+                      <ResourceTag key={`${conn.id}-${idx}`}>{label}</ResourceTag>
+                    ))}
+                  </ResourceTagList>
+                )}
+
+                <ResourceForm>
+                  <ResourceInput
+                    type="text"
+                    placeholder="Folder or file IDs, comma separated"
+                    value={resourceInputs[conn.id] || ''}
+                    onChange={(e) =>
+                      setResourceInputs((prev) => ({ ...prev, [conn.id]: e.target.value }))
+                    }
+                  />
+                  <SmallButton
+                    type="button"
+                    onClick={() => handleSaveResources(conn.id)}
+                    disabled={action === 'save' || !(resourceInputs[conn.id] || '').trim()}
+                  >
+                    {action === 'save' ? 'Saving…' : 'Save folders'}
+                  </SmallButton>
+                </ResourceForm>
+              </ConnectionCard>
+            );
+          })}
+        </ConnectionList>
+      )}
+    </Section>
+  );
+};
+
+export default DriveConnectionsSection;

llm-fe/src/llm-fe/components/Header2/Header2.test.tsx

@@ -0,0 +1,97 @@
+import React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import { MemoryRouter } from 'react-router-dom';
+import { ThemeProvider } from 'styled-components';
+import Header2 from './Header2';
+import { AuthContext } from '../../contexts/AuthContext';
+import { AccountContext } from '../../contexts/AccountContext';
+import { resetSubscriptionCache } from '../../hooks/useSubscription';
+
+const mockGet = jest.fn();
+const mockPost = jest.fn();
+
+jest.mock('../../../axiosApi', () => ({
+  axiosInstance: {
+    get: (...args: unknown[]) => mockGet(...args),
+    post: (...args: unknown[]) => mockPost(...args),
+    defaults: { headers: { common: {} as Record<string, string | null> } },
+  },
+  applyAccessToken: jest.fn(),
+}));
+
+const theme = {
+  main: '#4a90e2',
+  focus: '#224466',
+  darkMode: true,
+  colors: {
+    text: '#ffffff',
+    cardBorder: 'rgba(255,255,255,0.1)',
+  },
+};
+
+const renderHeader = () =>
+  render(
+    <MemoryRouter>
+      <ThemeProvider theme={theme as never}>
+        <AuthContext.Provider
+          value={{
+            authenticated: true,
+            setAuthentication: jest.fn(),
+            needsNewPassword: false,
+            setNeedsNewPassword: jest.fn(),
+            loading: false,
+          }}
+        >
+          <AccountContext.Provider value={{ account: undefined, setAccount: jest.fn() }}>
+            <Header2 />
+          </AccountContext.Provider>
+        </AuthContext.Provider>
+      </ThemeProvider>
+    </MemoryRouter>
+  );
+
+const subscriptionWithRag = (rag: boolean) => ({
+  data: {
+    plan: {
+      slug: rag ? 'pro' : 'standard',
+      name: rag ? 'Pro' : 'Standard',
+      features: {
+        text_generation: true,
+        image_generation: rag,
+        rag,
+        all_future_features: false,
+      },
+    },
+    status: 'active',
+    source: 'stripe',
+    needs_checkout: false,
+    stripe_subscription_id: 'sub_1',
+    usage: {},
+  },
+});
+
+describe('Header2 (#81 subscription-aware Documents nav link)', () => {
+  beforeEach(() => {
+    mockGet.mockReset();
+    mockPost.mockReset();
+    resetSubscriptionCache();
+  });
+
+  it('hides the Documents link while the plan does not include RAG', async () => {
+    mockGet.mockResolvedValue(subscriptionWithRag(false));
+
+    renderHeader();
+
+    await waitFor(() => expect(mockGet).toHaveBeenCalledWith('/finance/subscription/'));
+    expect(screen.queryByText('Documents')).not.toBeInTheDocument();
+  });
+
+  it('shows the Documents link once the plan includes RAG', async () => {
+    mockGet.mockResolvedValue(subscriptionWithRag(true));
+
+    renderHeader();
+
+    // Desktop nav + mobile dropdown both render a "Documents" link.
+    expect(await screen.findAllByText('Documents')).toHaveLength(2);
+  });
+});

llm-fe/src/llm-fe/components/Header2/Header2.tsx

@@ -5,6 +5,7 @@ import { AuthContext } from '../../contexts/AuthContext';
 import { AccountContext } from '../../contexts/AccountContext';
 import { applyAccessToken, axiosInstance } from '../../../axiosApi';
 import { clearTokens, getRefreshToken } from '../../auth/tokenStorage';
+import { useSubscription } from '../../hooks/useSubscription';
 import hesychiaMark from '../../assets/brand/hesychia-mark.png';
 
 const HeaderContainer = styled.header`
@@ -157,6 +158,7 @@ const Header2 = ({ absolute = false, light = false, isMini = false }: Header2Pro
   const navigate = useNavigate();
   const [isMenuOpen, setIsMenuOpen] = useState(false);
   const theme = useTheme();
+  const { hasRag } = useSubscription();
 
   const handleSignOut = async () => {
     try {
@@ -189,7 +191,7 @@ const Header2 = ({ absolute = false, light = false, isMini = false }: Header2Pro
       <Nav>
         <NavLink onClick={() => navigate('/')}>Dashboard</NavLink>
         <NavLink onClick={() => navigate('/account/')}>Account</NavLink>
-        <NavLink onClick={() => navigate('/document_storage/')}>Documents</NavLink>
+        {hasRag && <NavLink onClick={() => navigate('/document_storage/')}>Documents</NavLink>}
         <NavLink onClick={() => navigate('/analytics/')}>Analytics</NavLink>
         <NavLink onClick={() => navigate('/feedback/')}>Feedback</NavLink>
         <SignOutButton onClick={handleSignOut}>Sign Out</SignOutButton>
@@ -204,7 +206,7 @@ const Header2 = ({ absolute = false, light = false, isMini = false }: Header2Pro
       <MobileMenuDropdown isOpen={isMenuOpen}>
         <NavLink onClick={() => handleNavClick('/')}>Dashboard</NavLink>
         <NavLink onClick={() => handleNavClick('/account/')}>Account</NavLink>
-        <NavLink onClick={() => handleNavClick('/document_storage/')}>Documents</NavLink>
+        {hasRag && <NavLink onClick={() => handleNavClick('/document_storage/')}>Documents</NavLink>}
         <NavLink onClick={() => handleNavClick('/analytics/')}>Analytics</NavLink>
         <NavLink onClick={() => handleNavClick('/feedback/')}>Feedback</NavLink>
         <SignOutButton onClick={() => { handleSignOut(); setIsMenuOpen(false); }}>Sign Out</SignOutButton>

llm-fe/src/llm-fe/hooks/useSubscription.ts

@@ -0,0 +1,93 @@
+import { useCallback, useContext, useEffect, useRef, useState } from 'react';
+import { AxiosResponse } from 'axios';
+import { axiosInstance } from '../../axiosApi';
+import { AuthContext } from '../contexts/AuthContext';
+import { planAllowsRag, SubscriptionMe } from '../utils/finance';
+
+/**
+ * Module-level cache so every component that calls useSubscription() during the
+ * same session shares one /finance/subscription/ request instead of each
+ * mounting its own (e.g. Header2 + DocumentStoragePage on the same page).
+ */
+let cachedSubscription: SubscriptionMe | null = null;
+let inFlightRequest: Promise<SubscriptionMe | null> | null = null;
+
+async function loadSubscription(): Promise<SubscriptionMe | null> {
+  if (cachedSubscription) {
+    return cachedSubscription;
+  }
+  if (!inFlightRequest) {
+    inFlightRequest = axiosInstance
+      .get<SubscriptionMe>('/finance/subscription/')
+      .then((response: AxiosResponse<SubscriptionMe>) => {
+        cachedSubscription = response.data || null;
+        return cachedSubscription;
+      })
+      .catch(() => null)
+      .finally(() => {
+        inFlightRequest = null;
+      });
+  }
+  return inFlightRequest;
+}
+
+/** Drops the cached subscription so the next useSubscription() call refetches (e.g. after sign-out or a plan change). */
+export function resetSubscriptionCache(): void {
+  cachedSubscription = null;
+  inFlightRequest = null;
+}
+
+export type UseSubscriptionResult = {
+  subscription: SubscriptionMe | null;
+  loading: boolean;
+  hasRag: boolean;
+  refresh: () => Promise<void>;
+};
+
+/** Fetches /finance/subscription/ once per session (shared across callers) and exposes plan-derived flags. */
+export function useSubscription(): UseSubscriptionResult {
+  const { authenticated } = useContext(AuthContext);
+  const [subscription, setSubscription] = useState<SubscriptionMe | null>(cachedSubscription);
+  const [loading, setLoading] = useState(!cachedSubscription);
+  const mountedRef = useRef(true);
+
+  const load = useCallback(async () => {
+    if (!authenticated) {
+      setSubscription(null);
+      setLoading(false);
+      return;
+    }
+    setLoading((prev) => prev || !cachedSubscription);
+    const result = await loadSubscription();
+    if (mountedRef.current) {
+      setSubscription(result);
+      setLoading(false);
+    }
+  }, [authenticated]);
+
+  useEffect(() => {
+    mountedRef.current = true;
+    load();
+    return () => {
+      mountedRef.current = false;
+    };
+  }, [load]);
+
+  useEffect(() => {
+    if (!authenticated) {
+      resetSubscriptionCache();
+    }
+  }, [authenticated]);
+
+  const refresh = useCallback(async () => {
+    resetSubscriptionCache();
+    await load();
+  }, [load]);
+
+  return {
+    subscription,
+    loading,
+    hasRag: planAllowsRag(subscription),
+    refresh,
+  };
+}

llm-fe/src/llm-fe/pages/AuthCallback/AuthCallback.test.tsx

@@ -43,6 +43,7 @@ const renderCallback = (query: string) => {
             <Route path="/" element={<div>Home</div>} />
             <Route path="/terms_of_service/" element={<div>TOS</div>} />
             <Route path="/signin/" element={<div>Sign In Page</div>} />
+            <Route path="/document_storage/" element={<div>Document Storage Page</div>} />
           </Routes>
         </AccountContext.Provider>
       </AuthContext.Provider>
@@ -136,4 +137,13 @@ describe('AuthCallback', () => {
       expect(assignMock).toHaveBeenCalledWith('https://checkout.stripe.test/session');
     });
   });
+
+  it('redirects drive-link callbacks straight to Documents (#83/#84)', async () => {
+    renderCallback('?drive_connected=1');
+
+    await waitFor(() => {
+      expect(screen.getByText('Document Storage Page')).toBeInTheDocument();
+    });
+    expect(mockGet).not.toHaveBeenCalled();
+  });
 });

llm-fe/src/llm-fe/pages/AuthCallback/AuthCallback.tsx

@@ -111,6 +111,13 @@ const AuthCallback = (): JSX.Element => {
         return;
       }
 
+      // Drive-link flows (#83/#84) redirect here already authenticated — just
+      // bounce to Documents with a flag so it can show a success banner.
+      if (searchParams.get('drive_connected') === '1') {
+        navigate('/document_storage/?drive_connected=1', { replace: true });
+        return;
+      }
+
       const access = searchParams.get('access');
       const refresh = searchParams.get('refresh');
       const needsCheckout = searchParams.get('needs_checkout') === '1';

llm-fe/src/llm-fe/pages/DocumentStoragePage/DocumentStoragePage.test.tsx

@@ -0,0 +1,139 @@
+import React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import { MemoryRouter } from 'react-router-dom';
+import { ThemeProvider } from 'styled-components';
+import DocumentStoragePage from './DocumentStoragePage';
+import { AuthContext } from '../../contexts/AuthContext';
+import { AccountContext } from '../../contexts/AccountContext';
+import { Account } from '../../data';
+import { resetSubscriptionCache } from '../../hooks/useSubscription';
+
+const mockGet = jest.fn();
+const mockPost = jest.fn();
+const mockPatch = jest.fn();
+const mockDelete = jest.fn();
+
+jest.mock('../../../axiosApi', () => ({
+  axiosInstance: {
+    get: (...args: unknown[]) => mockGet(...args),
+    post: (...args: unknown[]) => mockPost(...args),
+    patch: (...args: unknown[]) => mockPatch(...args),
+    delete: (...args: unknown[]) => mockDelete(...args),
+    defaults: { headers: { common: {} as Record<string, string | null> } },
+  },
+  applyAccessToken: jest.fn(),
+}));
+
+const theme = {
+  main: '#4a90e2',
+  focus: '#224466',
+  darkMode: true,
+  colors: {
+    text: '#ffffff',
+    cardBackground: 'rgba(0,0,0,0.3)',
+    cardBorder: 'rgba(255,255,255,0.1)',
+  },
+};
+
+const subscriptionResponse = (rag: boolean) => ({
+  data: {
+    plan: {
+      slug: rag ? 'pro' : 'standard',
+      name: rag ? 'Pro' : 'Standard',
+      features: {
+        text_generation: true,
+        image_generation: rag,
+        rag,
+        all_future_features: false,
+      },
+    },
+    status: 'active',
+    source: 'stripe',
+    needs_checkout: false,
+    stripe_subscription_id: 'sub_1',
+    usage: {},
+  },
+});
+
+const mockGetByUrl = (overrides: { rag: boolean }) => {
+  mockGet.mockImplementation((url: string) => {
+    if (url === '/finance/subscription/') {
+      return Promise.resolve(subscriptionResponse(overrides.rag));
+    }
+    if (url === '/documents/') {
+      return Promise.resolve({ data: [] });
+    }
+    if (url === '/drive/connections/') {
+      return Promise.resolve({ data: [] });
+    }
+    return Promise.resolve({ data: [] });
+  });
+};
+
+const renderPage = (isCompanyManager = false) => {
+  const account = new Account({ email: 'user@example.com', is_company_manager: isCompanyManager });
+  return render(
+    <MemoryRouter initialEntries={['/document_storage/']}>
+      <ThemeProvider theme={theme as never}>
+        <AuthContext.Provider
+          value={{
+            authenticated: true,
+            setAuthentication: jest.fn(),
+            needsNewPassword: false,
+            setNeedsNewPassword: jest.fn(),
+            loading: false,
+          }}
+        >
+          <AccountContext.Provider value={{ account, setAccount: jest.fn() }}>
+            <DocumentStoragePage />
+          </AccountContext.Provider>
+        </AuthContext.Provider>
+      </ThemeProvider>
+    </MemoryRouter>
+  );
+};
+
+describe('DocumentStoragePage (#81/#82/#83/#84)', () => {
+  beforeEach(() => {
+    mockGet.mockReset();
+    mockPost.mockReset();
+    mockPatch.mockReset();
+    mockDelete.mockReset();
+    resetSubscriptionCache();
+  });
+
+  it('shows an upgrade card instead of upload tables when the plan has no RAG', async () => {
+    mockGetByUrl({ rag: false });
+
+    renderPage();
+
+    expect(await screen.findByText('Unlock document storage')).toBeInTheDocument();
+    expect(screen.getByRole('link', { name: /upgrade in billing/i })).toHaveAttribute(
+      'href',
+      '/account/'
+    );
+    expect(screen.queryByText('Company documents')).not.toBeInTheDocument();
+    expect(screen.queryByText('Upload a Document')).not.toBeInTheDocument();
+  });
+
+  it('shows the company documents table, upload card, and personal Cloud drives section when RAG is allowed', async () => {
+    mockGetByUrl({ rag: true });
+
+    renderPage(false);
+
+    expect(await screen.findByText('Company documents')).toBeInTheDocument();
+    expect(screen.getByText('Upload a Document')).toBeInTheDocument();
+    expect(screen.getByText('Cloud drives')).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: /connect google drive/i })).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: /connect onedrive/i })).toBeInTheDocument();
+    expect(screen.queryByText('Company knowledge sources')).not.toBeInTheDocument();
+  });
+
+  it('also shows the company knowledge sources section for company managers (#84)', async () => {
+    mockGetByUrl({ rag: true });
+
+    renderPage(true);
+
+    expect(await screen.findByText('Company knowledge sources')).toBeInTheDocument();
+  });
+});

llm-fe/src/llm-fe/pages/DocumentStoragePage/DocumentStoragePage.tsx

@@ -1,9 +1,13 @@
-import { useEffect, useState } from "react";
+import { useCallback, useContext, useEffect, useState } from "react";
+import { Link, useSearchParams } from "react-router-dom";
 import { Document, DocumentType } from "../../data";
 import { AxiosResponse } from "axios";
 import { axiosInstance } from "../../../axiosApi";
 import Header2 from "../../components/Header2/Header2";
 import ParticleBackground from "../../components/ParticleBackground/ParticleBackground";
+import DriveConnectionsSection from "../../components/DriveConnectionsSection/DriveConnectionsSection";
+import { AccountContext } from "../../contexts/AccountContext";
+import { useSubscription } from "../../hooks/useSubscription";
 import styled from "styled-components";
 
 // Styled Components
@@ -106,6 +110,26 @@ const StyledButton = styled.button`
   }
 `;
 
+const StyledButtonLink = styled(Link)`
+  background: ${({ theme }) => theme.main};
+  border: none;
+  border-radius: 0.5rem;
+  color: #fff;
+  padding: 0.8rem 1.5rem;
+  font-weight: 600;
+  cursor: pointer;
+  transition: transform 0.2s ease, box-shadow 0.2s ease;
+  display: inline-flex;
+  align-items: center;
+  gap: 0.5rem;
+  text-decoration: none;
+
+  &:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 4px 12px ${({ theme }) => theme.main}66;
+  }
+`;
+
 const FileInputLabel = styled.label`
   background: ${({ theme }) => theme.main};
   border: none;
@@ -181,38 +205,61 @@ const Checkbox = styled.input`
   }
 `;
 
+const HelperNote = styled.p`
+  margin: 1rem 0 0 0;
+  padding-top: 1rem;
+  border-top: 1px solid ${({ theme }) => theme.colors.cardBorder};
+  color: ${({ theme }) => theme.colors.text};
+  opacity: 0.65;
+  font-size: 0.9rem;
+  line-height: 1.5;
+`;
+
+const UpgradeText = styled.p`
+  color: ${({ theme }) => theme.colors.text};
+  opacity: 0.8;
+  line-height: 1.6;
+  margin-bottom: 1.5rem;
+`;
+
+const SuccessBanner = styled.div`
+  width: 100%;
+  max-width: 1000px;
+  margin-bottom: 1.5rem;
+  padding: 1rem 1.5rem;
+  border-radius: 0.75rem;
+  background: rgba(76, 175, 80, 0.15);
+  border: 1px solid rgba(76, 175, 80, 0.4);
+  color: ${({ theme }) => theme.colors.text};
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+  gap: 1rem;
+`;
+
+const DismissButton = styled.button`
+  background: transparent;
+  border: none;
+  color: ${({ theme }) => theme.colors.text};
+  font-size: 1.1rem;
+  cursor: pointer;
+  line-height: 1;
+  opacity: 0.7;
+
+  &:hover {
+    opacity: 1;
+  }
+`;
 
 type DocumentTableCardProps = {
-  documents: Document[],
+  documents: Document[];
-  setDocuments: React.Dispatch<React.SetStateAction<Document[]>>
+  onToggleActive: (id: number, active: boolean) => void;
 }
 
-const CompanyDocumentStorageTableCard = ({ documents, setDocuments }: DocumentTableCardProps): JSX.Element => {
+const CompanyDocumentStorageTableCard = ({ documents, onToggleActive }: DocumentTableCardProps): JSX.Element => {
-
-  useEffect(() => {
-    async function getUploadedDocuments() {
-      try {
-        const { data, }: AxiosResponse<DocumentType[]> = await axiosInstance.get(`/documents/`);
-        setDocuments(data.map((item) => new Document({
-
-          id: item.id,
-          name: item.file.replace(/^.*[\\/]/, ''),
-          date_uploaded: item.created.substring(0, 10),
-          active: item.active,
-          processed: item.processed,
-
-        })))
-
-
-      } catch (error) {
-        console.log(error)
-      }
-    }
-    getUploadedDocuments();
-  }, [setDocuments])
   return (
     <GlassCard>
-      <CardTitle>Your documents in the company workspace</CardTitle>
+      <CardTitle>Company documents</CardTitle>
       <div style={{ overflowX: 'auto' }}>
         <StyledTable>
           <thead>
@@ -240,7 +287,7 @@ const CompanyDocumentStorageTableCard = ({ documents, setDocuments }: DocumentTa
                     <Checkbox
                       type="checkbox"
                       checked={doc.active}
-                      disabled={true}
+                      onChange={(event) => onToggleActive(doc.id, event.target.checked)}
                     />
                     <Slider />
                   </ToggleSwitch>
@@ -250,41 +297,43 @@ const CompanyDocumentStorageTableCard = ({ documents, setDocuments }: DocumentTa
           </tbody>
         </StyledTable>
       </div>
+      <HelperNote>
+        These documents are shared with your whole company workspace. Want to search your own
+        files privately? Connect a personal cloud drive in the Cloud drives section below.
+      </HelperNote>
     </GlassCard>
   )
 }
 
-const UserDocumentStorageTableCard = (): JSX.Element => {
+type DocumentUploadCardProps = {
-  return (
+  onUploaded: () => void | Promise<void>;
-    <GlassCard>
-      <CardTitle>Your documents in your personal workspace</CardTitle>
-      <p style={{ color: 'rgba(255,255,255,0.7)' }}>This will become available shortly</p>
-    </GlassCard>
-  )
 }
 
-const DocumentUploadCard = (): JSX.Element => {
+const DocumentUploadCard = ({ onUploaded }: DocumentUploadCardProps): JSX.Element => {
   const [selectedFile, setSelectedFile] = useState<File | null>(null);
+  const [uploading, setUploading] = useState<boolean>(false);
 
   const handleDocumentUpload = async (): Promise<void> => {
+    if (!selectedFile) {
+      return;
+    }
+    setUploading(true);
+    try {
+      await axiosInstance.post('/documents/', {
+        file: selectedFile
+      },
+        {
+          headers: {
+            'Content-Type': 'multipart/form-data',
+          },
+        })
 
-    console.log(selectedFile)
+      setSelectedFile(null);
-    if (selectedFile) {
+      await onUploaded();
-      try {
+    } catch (error) {
-        await axiosInstance.post('/documents/', {
+      console.log(error)
-          file: selectedFile
+    } finally {
-        },
+      setUploading(false);
-          {
-            headers: {
-              'Content-Type': 'multipart/form-data',
-            },
-          })
-
-        // TODO set the documents here
-      }
-      finally {
-
-      }
     }
   }
 
@@ -299,14 +348,14 @@ const DocumentUploadCard = (): JSX.Element => {
       <div style={{ display: 'flex', alignItems: 'center', flexWrap: 'wrap', gap: '1rem' }}>
         <FileInputLabel>
           Select File
-          <input type="file" hidden onChange={handleFileChange} />
+          <input type="file" hidden onChange={handleFileChange} disabled={uploading} />
         </FileInputLabel>
 
         {selectedFile && (
           <>
             <span style={{ color: '#fff' }}>{selectedFile.name}</span>
-            <StyledButton onClick={handleDocumentUpload}>
+            <StyledButton onClick={handleDocumentUpload} disabled={uploading}>
-              Upload
+              {uploading ? 'Uploading…' : 'Upload'}
             </StyledButton>
           </>
         )}
@@ -317,22 +366,113 @@ const DocumentUploadCard = (): JSX.Element => {
 
 const DocumentStoragePageInner = (): JSX.Element => {
   const [documents, setDocuments] = useState<Document[]>([]);
+  const { account } = useContext(AccountContext);
+  const [searchParams, setSearchParams] = useSearchParams();
+  const [showDriveConnectedBanner, setShowDriveConnectedBanner] = useState(false);
+
+  const fetchDocuments = useCallback(async () => {
+    try {
+      const { data, }: AxiosResponse<DocumentType[]> = await axiosInstance.get(`/documents/`);
+      setDocuments(data.map((item) => new Document({
+
+        id: item.id,
+        name: item.file.replace(/^.*[\\/]/, ''),
+        date_uploaded: item.created.substring(0, 10),
+        active: item.active,
+        processed: item.processed,
+
+      })))
+    } catch (error) {
+      console.log(error)
+    }
+  }, [])
+
+  useEffect(() => {
+    fetchDocuments();
+  }, [fetchDocuments])
+
+  useEffect(() => {
+    if (searchParams.get('drive_connected') === '1') {
+      setShowDriveConnectedBanner(true);
+      const next = new URLSearchParams(searchParams);
+      next.delete('drive_connected');
+      setSearchParams(next, { replace: true });
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
+  const handleToggleActive = useCallback(async (id: number, active: boolean) => {
+    setDocuments((prev) => prev.map((doc) => (doc.id === id ? { ...doc, active } as Document : doc)));
+    try {
+      await axiosInstance.patch(`documents_details/${id}`, { active });
+    } catch (error) {
+      console.log(error);
+      setDocuments((prev) => prev.map((doc) => (doc.id === id ? { ...doc, active: !active } as Document : doc)));
+    }
+  }, []);
+
   return (
     <>
-      <CompanyDocumentStorageTableCard documents={documents} setDocuments={setDocuments} />
+      {showDriveConnectedBanner && (
-      <UserDocumentStorageTableCard />
+        <SuccessBanner>
-      <DocumentUploadCard />
+          <span>Drive connected. Choose folders below to include them in your knowledge base.</span>
+          <DismissButton
+            type="button"
+            aria-label="Dismiss"
+            onClick={() => setShowDriveConnectedBanner(false)}
+          >
+            ×
+          </DismissButton>
+        </SuccessBanner>
+      )}
+
+      <CompanyDocumentStorageTableCard documents={documents} onToggleActive={handleToggleActive} />
+      <DocumentUploadCard onUploaded={fetchDocuments} />
+
+      <DriveConnectionsSection
+        kind="personal"
+        title="Cloud drives"
+        description="Connect your personal Google Drive or OneDrive so its files can be searched in chat."
+        connectIntent="link_drive"
+      />
+
+      {account?.is_company_manager && (
+        <DriveConnectionsSection
+          kind="company"
+          title="Company knowledge sources"
+          description="Connect a shared Google Shared Drive or Microsoft 365 site so your whole team can search these files in chat."
+          connectIntent="link_company_drive"
+        />
+      )}
     </>
   )
 }
 
 const DocumentStoragePage = (): JSX.Element => {
+  const { hasRag, loading } = useSubscription();
+
   return (
     <PageContainer>
       <ParticleBackground />
       <Header2 />
       <ContentWrapper>
-        <DocumentStoragePageInner />
+        {loading ? (
+          <GlassCard>
+            <CardTitle>Loading…</CardTitle>
+          </GlassCard>
+        ) : hasRag ? (
+          <DocumentStoragePageInner />
+        ) : (
+          <GlassCard>
+            <CardTitle>Unlock document storage</CardTitle>
+            <UpgradeText>
+              Document uploads and cloud drive connections are available on plans that include
+              RAG document search. Upgrade your plan to start uploading files and connecting
+              Google Drive or OneDrive.
+            </UpgradeText>
+            <StyledButtonLink to="/account/">Upgrade in Billing</StyledButtonLink>
+          </GlassCard>
+        )}
       </ContentWrapper>
     </PageContainer>
   )

llm-fe/src/llm-fe/utils/chatErrors.test.ts

@@ -0,0 +1,60 @@
+import { isRagFeatureNotAllowed, parseChatErrorPayload } from './chatErrors';
+
+describe('parseChatErrorPayload', () => {
+  it('parses an error-type websocket payload', () => {
+    const raw = JSON.stringify({ type: 'error', code: 'feature_not_allowed', content: 'no rag' });
+    expect(parseChatErrorPayload(raw)).toEqual({
+      type: 'error',
+      code: 'feature_not_allowed',
+      content: 'no rag',
+    });
+  });
+
+  it('returns null for non-error JSON payloads', () => {
+    expect(parseChatErrorPayload(JSON.stringify({ type: 'text', content: 'hi' }))).toBeNull();
+  });
+
+  it('returns null for plain streamed text chunks', () => {
+    expect(parseChatErrorPayload('just a plain chunk of text')).toBeNull();
+  });
+});
+
+describe('isRagFeatureNotAllowed', () => {
+  it('is false for unrelated error codes', () => {
+    expect(isRagFeatureNotAllowed({ code: 'prompt_quota_exceeded', content: 'slow down' })).toBe(false);
+  });
+
+  it('is false when feature_not_allowed is about a different feature', () => {
+    expect(
+      isRagFeatureNotAllowed({
+        code: 'feature_not_allowed',
+        content: 'Your plan does not include image generation.',
+        details: { feature: 'image_generation' },
+      })
+    ).toBe(false);
+  });
+
+  it('is true when details.feature mentions rag', () => {
+    expect(
+      isRagFeatureNotAllowed({
+        code: 'feature_not_allowed',
+        content: 'Blocked',
+        details: { feature: 'rag' },
+      })
+    ).toBe(true);
+  });
+
+  it('is true when the message text mentions rag without structured details', () => {
+    expect(
+      isRagFeatureNotAllowed({
+        code: 'feature_not_allowed',
+        content: 'Your plan does not include RAG document search.',
+      })
+    ).toBe(true);
+  });
+
+  it('is false for a null/undefined payload', () => {
+    expect(isRagFeatureNotAllowed(null)).toBe(false);
+    expect(isRagFeatureNotAllowed(undefined)).toBe(false);
+  });
+});

llm-fe/src/llm-fe/utils/chatErrors.ts

@@ -0,0 +1,37 @@
+/**
+ * Helpers for the websocket chat error payloads sent by the backend
+ * (see chat_backend/consumers.py — `{"type": "error", "code": ..., "content": ..., "details": {...}}`).
+ */
+export type ChatErrorPayload = {
+  type?: string;
+  code?: string;
+  content?: string;
+  message?: string;
+  details?: Record<string, unknown> | null;
+};
+
+/** Parses a raw websocket text chunk into an error payload, or null when it isn't one. */
+export function parseChatErrorPayload(message: string): ChatErrorPayload | null {
+  try {
+    const parsed = JSON.parse(message);
+    if (parsed && typeof parsed === 'object' && parsed.type === 'error') {
+      return parsed as ChatErrorPayload;
+    }
+  } catch {
+    /* not JSON — plain streamed text chunk */
+  }
+  return null;
+}
+
+/** True when a `feature_not_allowed` error payload is about the RAG (document search) feature. */
+export function isRagFeatureNotAllowed(payload: ChatErrorPayload | null | undefined): boolean {
+  if (!payload || payload.code !== 'feature_not_allowed') {
+    return false;
+  }
+  const feature = payload.details?.feature;
+  if (typeof feature === 'string' && feature.toLowerCase().includes('rag')) {
+    return true;
+  }
+  const text = `${payload.content ?? ''} ${payload.message ?? ''}`.toLowerCase();
+  return text.includes('rag');
+}

llm-fe/src/llm-fe/utils/drive.test.ts

@@ -0,0 +1,35 @@
+import { driveConnectUrl, parseResourceIdsInput } from './drive';
+
+describe('parseResourceIdsInput', () => {
+  it('splits comma and newline separated ids and trims whitespace', () => {
+    expect(parseResourceIdsInput('abc, def\nghi ,, ')).toEqual(['abc', 'def', 'ghi']);
+  });
+
+  it('returns an empty array for blank input', () => {
+    expect(parseResourceIdsInput('   ')).toEqual([]);
+  });
+});
+
+describe('driveConnectUrl', () => {
+  const originalEnv = process.env.REACT_APP_BACKEND_REST_API_BASE_URL;
+
+  beforeAll(() => {
+    process.env.REACT_APP_BACKEND_REST_API_BASE_URL = 'https://api.example.com/';
+  });
+
+  afterAll(() => {
+    process.env.REACT_APP_BACKEND_REST_API_BASE_URL = originalEnv;
+  });
+
+  it('builds a personal drive-link url', () => {
+    expect(driveConnectUrl('google', 'link_drive')).toBe(
+      'https://api.example.com/auth/oauth/google/start/?intent=link_drive'
+    );
+  });
+
+  it('builds a company drive-link url', () => {
+    expect(driveConnectUrl('microsoft', 'link_company_drive')).toBe(
+      'https://api.example.com/auth/oauth/microsoft/start/?intent=link_company_drive'
+    );
+  });
+});

llm-fe/src/llm-fe/utils/drive.ts

@@ -0,0 +1,90 @@
+import { axiosInstance } from '../../axiosApi';
+import { oauthStartUrl } from '../auth/sso';
+
+export type DriveProvider = 'google' | 'microsoft';
+
+export type DriveConnectionKind = 'personal' | 'company';
+
+/** intent query param sent to /auth/oauth/:provider/start/ for drive-linking flows (#83/#84). */
+export type DriveConnectIntent = 'link_drive' | 'link_company_drive';
+
+export type DriveConnectionType = {
+  id: number;
+  provider: DriveProvider;
+  kind?: DriveConnectionKind;
+  is_active?: boolean;
+  external_account_email?: string | null;
+  selected_resource_ids?: string[];
+  selected_resource_labels?: string[];
+  last_sync_at?: string | null;
+  last_sync_status?: string | null;
+  last_sync_error?: string;
+  created?: string;
+};
+
+export async function fetchDriveConnections(): Promise<DriveConnectionType[]> {
+  const { data } = await axiosInstance.get<DriveConnectionType[]>('/drive/connections/');
+  return Array.isArray(data) ? data : [];
+}
+
+export async function disconnectDriveConnection(connectionId: number): Promise<void> {
+  await axiosInstance.delete(`/drive/connections/${connectionId}/`);
+}
+
+export type DriveResourceSelection = {
+  resource_ids: string[];
+  resource_labels: string[];
+};
+
+export async function saveDriveResourceSelection(
+  connectionId: number,
+  selection: DriveResourceSelection
+): Promise<DriveConnectionType> {
+  const { data } = await axiosInstance.post<DriveConnectionType>(
+    `/drive/connections/${connectionId}/resources/`,
+    selection
+  );
+  return data;
+}
+
+export async function syncDriveConnection(connectionId: number): Promise<DriveConnectionType | null> {
+  const { data } = await axiosInstance.post<{
+    connection?: DriveConnectionType;
+  } & DriveConnectionType>(`/drive/connections/${connectionId}/sync/`);
+  if (!data) return null;
+  return data.connection || data;
+}
+
+/** Absolute backend OAuth start URL for a drive-link flow (personal or company). */
+export function driveConnectUrl(provider: DriveProvider, intent: DriveConnectIntent): string {
+  return oauthStartUrl(provider, intent);
+}
+
+/**
+ * Begin Drive link: call start with Bearer JWT, then navigate to IdP authorize URL.
+ * Full-page assign alone cannot send Authorization, so the backend returns JSON.
+ */
+export async function connectDrive(
+  provider: DriveProvider,
+  intent: DriveConnectIntent
+): Promise<void> {
+  const { data } = await axiosInstance.get<{ authorize_url: string }>(
+    `/auth/oauth/${provider}/start/`,
+    {
+      params: { intent, response: 'json' },
+      headers: { Accept: 'application/json' },
+    }
+  );
+  if (!data?.authorize_url) {
+    throw new Error('Drive connect did not return an authorize URL.');
+  }
+  window.location.assign(data.authorize_url);
+}
+
+/** Splits a comma/newline separated textarea/input value into trimmed, non-empty resource ids. */
+export function parseResourceIdsInput(raw: string): string[] {
+  return raw
+    .split(/[,\n]/)
+    .map((value) => value.trim())
+    .filter(Boolean);
+}

llm-fe/src/llm-fe/utils/finance.test.ts

@@ -4,8 +4,46 @@ import {
   formatTokenCount,
   humanizeStatus,
   pickPrimaryInvoice,
+  planAllowsRag,
 } from './finance';
-import type { FinanceInvoice } from './finance';
+import type { FinanceInvoice, SubscriptionMe, SubscriptionPlanInfo } from './finance';
+
+const basePlan = (overrides: Partial<SubscriptionPlanInfo> = {}): SubscriptionPlanInfo => ({
+  slug: 'standard',
+  name: 'Standard',
+  description: '',
+  price_cents: 999,
+  currency: 'usd',
+  interval: 'month',
+  is_public: true,
+  is_selectable: true,
+  features: {
+    text_generation: true,
+    image_generation: false,
+    rag: false,
+    all_future_features: false,
+  },
+  prompt_quota_per_window: 100,
+  prompt_window_hours: 6,
+  monthly_token_quota: null,
+  sort_order: 1,
+  ...overrides,
+});
+
+const baseUsage: SubscriptionMe['usage'] = {
+  prompts_in_window: 0,
+  prompt_quota: null,
+  prompts_remaining: null,
+  window_hours: 6,
+  tokens_in_period: null,
+  tokens_out_period: null,
+  tokens_total_period: null,
+  turns_missing_token_usage: 0,
+  monthly_token_quota: null,
+  tokens_remaining: null,
+  period_start: null,
+  period_end: null,
+};
 
 const baseInvoice = (overrides: Partial<FinanceInvoice> = {}): FinanceInvoice => ({
   id: 1,
@@ -59,3 +97,50 @@ describe('finance helpers', () => {
     expect(canOpenBillingPortal([baseInvoice({ status: 'paid' })])).toBe(true);
   });
 });
+
+describe('planAllowsRag', () => {
+  it('is false when there is no plan/subscription', () => {
+    expect(planAllowsRag(null)).toBe(false);
+    expect(planAllowsRag(undefined)).toBe(false);
+  });
+
+  it('is false when the rag feature flag is off', () => {
+    expect(planAllowsRag(basePlan())).toBe(false);
+  });
+
+  it('is true when the rag feature flag is on', () => {
+    expect(
+      planAllowsRag(basePlan({ features: { text_generation: true, image_generation: false, rag: true, all_future_features: false } }))
+    ).toBe(true);
+  });
+
+  it('is true when all_future_features unlocks it', () => {
+    expect(
+      planAllowsRag(basePlan({ features: { text_generation: true, image_generation: true, rag: false, all_future_features: true } }))
+    ).toBe(true);
+  });
+
+  it('accepts a SubscriptionMe wrapper and reads its plan', () => {
+    const subscription: SubscriptionMe = {
+      plan: basePlan({ features: { text_generation: true, image_generation: false, rag: true, all_future_features: false } }),
+      status: 'active',
+      source: 'stripe',
+      needs_checkout: false,
+      stripe_subscription_id: 'sub_1',
+      usage: baseUsage,
+    };
+    expect(planAllowsRag(subscription)).toBe(true);
+  });
+
+  it('is false when a subscription has no plan yet', () => {
+    const subscription: SubscriptionMe = {
+      plan: null,
+      status: 'none',
+      source: 'none',
+      needs_checkout: true,
+      stripe_subscription_id: '',
+      usage: baseUsage,
+    };
+    expect(planAllowsRag(subscription)).toBe(false);
+  });
+});

llm-fe/src/llm-fe/utils/finance.ts

@@ -34,6 +34,7 @@ export type FinancePayment = {
 export type PlanFeatures = {
   text_generation: boolean;
   image_generation: boolean;
+  rag: boolean;
   all_future_features: boolean;
 };
 
@@ -105,6 +106,17 @@ export function higherSelectablePlans(
   );
 }
 
+/** True if the plan (or a subscription's plan) includes RAG document search, or unlocks all future features. */
+export function planAllowsRag(
+  input: SubscriptionPlanInfo | SubscriptionMe | null | undefined
+): boolean {
+  if (!input) return false;
+  const plan: SubscriptionPlanInfo | null | undefined =
+    'features' in input ? input : input.plan;
+  if (!plan) return false;
+  return Boolean(plan.features?.rag || plan.features?.all_future_features);
+}
+
 export function otherSelectablePlans(
   plans: SubscriptionPlanInfo[],
   currentSlug: string | null | undefined