Skip to content

Commit 8820b14

Browse files
sashalavronsashalavron
committed
feat(omnibar): implement per-tab attachment persistence for files, images, and tabs
- Enhanced the omnibar to support per-tab persistence for file, image, and tab attachments, ensuring that attachments are scoped to their respective tabs. - Updated the `useFileAttachments`, `useImageAttachments`, and `useTabAttachments` hooks to utilize a new `useStateWithLocalPersistence` method, allowing attachments to survive tab switches. - Refactored the `Omnibar` and `OmnibarCustomized` components to integrate the new attachment providers, improving the overall structure and maintainability. - Added integration tests to verify that attachments persist correctly across tab switches and are cleared when a tab is closed. These changes significantly improve the user experience by maintaining context across tab interactions in the omnibar.
1 parent b17a5ee commit 8820b14

9 files changed

Lines changed: 227 additions & 19 deletions

File tree

‎special-pages/pages/new-tab/app/omnibar/components/Omnibar.js‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,7 @@ export function Omnibar({
170170
enableRecentAiChats={enableRecentAiChats}
171171
enableVoiceChatAccess={enableVoiceChatAccess}
172172
enableAttachTabs={enableAttachTabs}
173+
tabId={tabId}
173174
onChange={setQuery}
174175
onSubmit={handleSubmitChat}
175176
/>
@@ -189,6 +190,7 @@ export function Omnibar({
189190
* @param {boolean} props.enableRecentAiChats
190191
* @param {boolean} [props.enableVoiceChatAccess]
191192
* @param {boolean} [props.enableAttachTabs]
193+
* @param {string|null|undefined} [props.tabId]
192194
* @param {(query: string) => void} props.onChange
193195
* @param {(params: SubmitChatAction) => void} props.onSubmit
194196
*/
@@ -198,6 +200,7 @@ function AiChatContent({
198200
enableRecentAiChats,
199201
enableVoiceChatAccess = false,
200202
enableAttachTabs = false,
203+
tabId,
201204
onChange,
202205
onSubmit,
203206
}) {
@@ -210,7 +213,7 @@ function AiChatContent({
210213
const containerRef = useRef(/** @type {HTMLDivElement|null} */ (null));
211214
const hasVisibleImagesRef = useRef(false);
212215
const [imageWarning, setImageWarning] = useState(false);
213-
const imageState = useImageAttachments();
216+
const imageState = useImageAttachments(tabId);
214217

215218
const hasAttachedImages = imageState.attachedImages.length > 0;
216219
const imageGenerationPlaceholder = hasAttachedImages
@@ -219,11 +222,11 @@ function AiChatContent({
219222
const selectedModelSupportsImages = selectedModel?.supportsImageUpload ?? false;
220223
const canAttachImages = selectedModelSupportsImages || imageGenerationActive;
221224

222-
const fileState = useFileAttachments(selectedModel?.supportedFileTypes);
225+
const fileState = useFileAttachments(selectedModel?.supportedFileTypes, tabId);
223226
const canAttachFiles = !imageGenerationActive && (selectedModel?.supportedFileTypes?.length ?? 0) > 0;
224227

225228
const canAttachTabs = enableAttachTabs && !imageGenerationActive;
226-
const tabAttachments = useTabAttachments();
229+
const tabAttachments = useTabAttachments(tabId);
227230
const textareaRef = useRef(/** @type {HTMLTextAreaElement|null} */ (null));
228231
const mention = useMentionPicker({
229232
enabled: canAttachTabs,

‎special-pages/pages/new-tab/app/omnibar/components/OmnibarCustomized.js‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,17 @@ import { h } from 'preact';
66

77
import { OmnibarConsumer } from './OmnibarConsumer.js';
88
import { SearchIcon } from '../../components/Icons.js';
9-
import { PersistentModeProvider, PersistentTextInputProvider } from './PersistentOmnibarValuesProvider.js';
9+
import {
10+
FileAttachments,
11+
ImageAttachments,
12+
PersistentModeProvider,
13+
PersistentTextInputProvider,
14+
TabAttachments,
15+
} from './PersistentOmnibarValuesProvider.js';
16+
17+
const { Provider: TabAttachmentsProvider } = TabAttachments;
18+
const { Provider: FileAttachmentsProvider } = FileAttachments;
19+
const { Provider: ImageAttachmentsProvider } = ImageAttachments;
1020

1121
/**
1222
* @import enStrings from "../strings.json"
@@ -34,9 +44,15 @@ export function OmnibarCustomized() {
3444
return (
3545
<PersistentTextInputProvider>
3646
<PersistentModeProvider>
37-
<OmnibarProvider>
38-
<OmnibarConsumer />
39-
</OmnibarProvider>
47+
<TabAttachmentsProvider>
48+
<FileAttachmentsProvider>
49+
<ImageAttachmentsProvider>
50+
<OmnibarProvider>
51+
<OmnibarConsumer />
52+
</OmnibarProvider>
53+
</ImageAttachmentsProvider>
54+
</FileAttachmentsProvider>
55+
</TabAttachmentsProvider>
4056
</PersistentModeProvider>
4157
</PersistentTextInputProvider>
4258
);

‎special-pages/pages/new-tab/app/omnibar/components/PersistentOmnibarValuesProvider.js‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ import { PersistentValue } from '../../tabs/PersistentValue.js';
66

77
/**
88
* @typedef {import("../../../types/new-tab.js").OmnibarConfig["mode"]} Mode
9+
* @typedef {import("./chat-tools/tab-attachment/useTabAttachments.js").AttachedTab} AttachedTab
10+
* @typedef {import("./chat-tools/file-attachment/useFileAttachments.js").AttachedFile} AttachedFile
11+
* @typedef {import("./chat-tools/image-attachment/useImageAttachments.js").AttachedImage} AttachedImage
912
*/
1013

1114
const TextInputContext = createContext(/** @type {PersistentValue<string>|null} */ (null));
@@ -41,6 +44,57 @@ export function PersistentModeProvider({ children }) {
4144
return <ModeContext.Provider value={value}>{children}</ModeContext.Provider>;
4245
}
4346

47+
/**
48+
* @template T
49+
* @typedef {object} PersistentList
50+
* @property {(props: { children: import('preact').ComponentChildren }) => import('preact').VNode} Provider
51+
* @property {(tabId: string|null|undefined) => readonly [T[], (next: T[] | ((prev: T[]) => T[])) => void]} useStateWithLocalPersistence
52+
*/
53+
54+
/**
55+
* @template T
56+
* @returns {PersistentList<T>}
57+
*/
58+
function createPersistentList() {
59+
const Context = createContext(/** @type {PersistentValue<T[]>|null} */ (null));
60+
61+
/** @param {{ children: import('preact').ComponentChildren }} props */
62+
function Provider({ children }) {
63+
const [store] = useState(() => /** @type {PersistentValue<T[]>} */ (new PersistentValue()));
64+
const { all } = useTabState();
65+
useEffect(() => {
66+
return all.subscribe((tabIds) => store.prune({ preserve: tabIds }));
67+
}, [all, store]);
68+
69+
return <Context.Provider value={store}>{children}</Context.Provider>;
70+
}
71+
72+
/** @param {string|null|undefined} tabId */
73+
function useStateWithLocalPersistence(tabId) {
74+
const store = useContext(Context);
75+
const [items, setItems] = useState(() => store?.byId(tabId) ?? []);
76+
const setter = useCallback(
77+
/** @param {T[] | ((prev: T[]) => T[])} next */
78+
(next) => {
79+
setItems((prev) => {
80+
const value = typeof next === 'function' ? /** @type {(prev: T[]) => T[]} */ (next)(prev) : next;
81+
if (tabId) store?.update({ id: tabId, value });
82+
return value;
83+
});
84+
},
85+
[store, tabId],
86+
);
87+
return /** @type {const} */ ([items, setter]);
88+
}
89+
90+
return { Provider, useStateWithLocalPersistence };
91+
}
92+
93+
// Three independent attachment lists, each with its own Provider + `useStateWithLocalPersistence` hook.
94+
export const TabAttachments = /** @type {() => PersistentList<AttachedTab>} */ (createPersistentList)();
95+
export const FileAttachments = /** @type {() => PersistentList<AttachedFile>} */ (createPersistentList)();
96+
export const ImageAttachments = /** @type {() => PersistentList<AttachedImage>} */ (createPersistentList)();
97+
4498
/**
4599
* A normal set-state, but with values recorded. Must be used when the Omnibar Service is ready
46100
* @param {string|null|undefined} tabId

‎special-pages/pages/new-tab/app/omnibar/components/chat-tools/file-attachment/useFileAttachments.js‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
import { useState } from 'preact/hooks';
2+
import { FileAttachments } from '../../PersistentOmnibarValuesProvider';
3+
4+
const { useStateWithLocalPersistence } = FileAttachments;
25

36
/**
47
* @typedef {{ data: string, fileName: string, mimeType: string }} AttachedFile
@@ -9,9 +12,12 @@ const FILE_READ_TIMEOUT = 30000;
912

1013
/**
1114
* @param {string[] | undefined} supportedFileTypes — MIME types the active model accepts.
15+
* @param {string|null|undefined} [tabId] - The NTP tab these attachments belong to. Used to persist
16+
* them per-tab so they survive switching between browser tabs (see `PersistentAttachmentsProvider`).
1217
*/
13-
export function useFileAttachments(supportedFileTypes) {
14-
const [attachedFiles, setAttachedFiles] = useState(/** @type {AttachedFile[]} */ ([]));
18+
export function useFileAttachments(supportedFileTypes, tabId) {
19+
const [attachedFiles, setAttachedFiles] = useStateWithLocalPersistence(tabId);
20+
1521
const allowList = supportedFileTypes ?? [];
1622
const allowListKey = allowList.join('|');
1723

‎special-pages/pages/new-tab/app/omnibar/components/chat-tools/image-attachment/useImageAttachments.js‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
import { useState } from 'preact/hooks';
2+
import { ImageAttachments } from '../../PersistentOmnibarValuesProvider';
3+
4+
const { useStateWithLocalPersistence } = ImageAttachments;
25

36
/**
47
* @typedef {{ dataUrl: string, fileName: string, mimeType: string }} AttachedImage
@@ -78,8 +81,12 @@ function normaliseImage(srcDataUrl, targetMime) {
7881
});
7982
}
8083

81-
export function useImageAttachments() {
82-
const [attachedImages, setAttachedImages] = useState(/** @type {AttachedImage[]} */ ([]));
84+
/**
85+
* @param {string|null|undefined} [tabId] - The NTP tab these attachments belong to. Used to persist
86+
* them per-tab so they survive switching between browser tabs (see `PersistentAttachmentsProvider`).
87+
*/
88+
export function useImageAttachments(tabId) {
89+
const [attachedImages, setAttachedImages] = useStateWithLocalPersistence(tabId);
8390
const [imageError, setImageError] = useState(/** @type {ImageError|null} */ (null));
8491

8592
const imageLimitExceeded = attachedImages.length > MAX_IMAGES;

‎special-pages/pages/new-tab/app/omnibar/components/chat-tools/tab-attachment/fileChannels.js‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,23 @@
77

88
const IMAGE_ACCEPT = 'image/jpeg,image/png,image/webp';
99

10+
/**
11+
* @type {Record<string, string>}
12+
*/
13+
const FILE_EXTENSIONS = {
14+
'application/pdf': '.pdf',
15+
};
16+
17+
/**
18+
* Builds the file channel's `accept` fragment from the model's supported MIME
19+
* types (e.g. `['application/pdf']` → `.pdf`).
20+
* @param {string[]} mimeTypes
21+
* @returns {string}
22+
*/
23+
function buildFileAccept(mimeTypes) {
24+
return mimeTypes.map((mime) => FILE_EXTENSIONS[mime] ?? mime).join(',');
25+
}
26+
1027
/**
1128
* Collapses the optional image / file (PDF) channels into the config for the
1229
* single hidden `<input type="file">` rendered by {@link AttachMenu}:
@@ -24,7 +41,7 @@ const IMAGE_ACCEPT = 'image/jpeg,image/png,image/webp';
2441
export function resolveFileInput({ t, image, file }) {
2542
const label =
2643
image && file ? t('omnibar_attachImageOrFileLabel') : image ? t('omnibar_attachImageLabel') : t('omnibar_attachFileLabel');
27-
const accept = [image ? IMAGE_ACCEPT : '', file ? file.mimeTypes.join(',') : ''].filter(Boolean).join(',');
44+
const accept = [...(image ? [IMAGE_ACCEPT] : []), ...(file ? [buildFileAccept(file.mimeTypes)] : [])].filter(Boolean).join(',');
2845
const disabled = (image?.disabled ?? true) && (file?.disabled ?? true);
2946

3047
/** @param {Event} event */

‎special-pages/pages/new-tab/app/omnibar/components/chat-tools/tab-attachment/useTabAttachments.js‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
1-
import { useCallback, useContext, useMemo, useState } from 'preact/hooks';
1+
import { useCallback, useContext, useMemo } from 'preact/hooks';
22
import { OmnibarContext } from '../../OmnibarProvider';
3+
import { TabAttachments } from '../../PersistentOmnibarValuesProvider';
4+
5+
const { useStateWithLocalPersistence } = TabAttachments;
36

47
/**
58
* @typedef {import('../../../../../types/new-tab.js').TabMetadata} TabMetadata
@@ -14,10 +17,14 @@ import { OmnibarContext } from '../../OmnibarProvider';
1417
* @property {PageContext|null} pageContext — Native-extracted content, populated once `getTabContent` resolves. `null` when extraction failed.
1518
*/
1619

17-
export function useTabAttachments() {
20+
/**
21+
* @param {string|null|undefined} tabId — The NTP tab these attachments belong to. Used to persist
22+
* them per-tab so they survive switching between browser tabs (see `PersistentAttachmentsProvider`).
23+
*/
24+
export function useTabAttachments(tabId) {
1825
const { getTabContent } = useContext(OmnibarContext);
1926

20-
const [attachedTabs, setAttachedTabs] = useState(/** @type {AttachedTab[]} */ ([]));
27+
const [attachedTabs, setAttachedTabs] = useStateWithLocalPersistence(tabId);
2128

2229
const isAttached = useCallback(
2330
/** @param {string} tabId */
@@ -58,20 +65,20 @@ export function useTabAttachments() {
5865
setAttachedTabs((prev) => prev.filter((t) => t.tabId !== tabToAttach.tabId));
5966
}
6067
},
61-
[getTabContent],
68+
[getTabContent, setAttachedTabs],
6269
);
6370

6471
const removeTab = useCallback(
6572
/** @param {string} tabId */
6673
(tabId) => {
6774
setAttachedTabs((prev) => prev.filter((tab) => tab.tabId !== tabId));
6875
},
69-
[],
76+
[setAttachedTabs],
7077
);
7178

7279
const clearAttachedTabs = useCallback(() => {
7380
setAttachedTabs([]);
74-
}, []);
81+
}, [setAttachedTabs]);
7582

7683
/**
7784
* @returns {PageContext[] | null}

‎special-pages/pages/new-tab/app/omnibar/integration-tests/omnibar-attachments.spec.js‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ const PDF_BYTES = Buffer.from(
1919
'base64',
2020
);
2121

22+
/** A 1x1 PNG, base64-encoded, used to drive image `setInputFiles` without a fixture file. */
23+
const TINY_PNG = Buffer.from('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwADhQGAWjR9awAAAABJRU5ErkJggg==', 'base64');
24+
2225
/** @param {import('@playwright/test').Page} page @param {import('@playwright/test').TestInfo} workerInfo */
2326
function setup(page, workerInfo) {
2427
const ntp = NewtabPage.create(page, workerInfo);
@@ -266,3 +269,97 @@ test.describe('omnibar attachment coexistence', () => {
266269
expect(params.pageContext?.[0].tabId).toBe('tab-2');
267270
});
268271
});
272+
273+
/**
274+
* Attachments are stored per NTP tab (above the tab-keyed remount boundary), so they must
275+
* survive switching between browser tabs and be dropped when a tab closes — see the tech
276+
* designs' "chips persist across browser-tab switches and clear when the NTP tab closes".
277+
*/
278+
test.describe('omnibar attachment per-tab persistence', () => {
279+
/**
280+
* Multi-tab (`tabs`) so `didSwitchToTab` drives a real tabId change, AI mode with attach-tabs
281+
* on, and a model that accepts both images and PDFs so all three attachment routes are available.
282+
*/
283+
const config = {
284+
tabs: true,
285+
'omnibar.mode': 'ai',
286+
'omnibar.enableAttachTabs': 'true',
287+
'omnibar.enableAiChatTools': 'true',
288+
'omnibar.selectedModelId': 'claude-haiku-4-5',
289+
};
290+
291+
test('attached tabs survive a browser-tab switch and stay scoped to their tab', async ({ page }, workerInfo) => {
292+
const { ntp, omnibar } = setup(page, workerInfo);
293+
await ntp.reducedMotion();
294+
await ntp.openPage({ additional: config });
295+
await omnibar.ready();
296+
297+
// attach two tabs on tab 01
298+
await omnibar.attachTab('Starbucks Coffee Company');
299+
await omnibar.attachTab('MacBook Neo - Apple');
300+
await expect(omnibar.attachmentChips().locator('[data-status="ready"]')).toHaveCount(2);
301+
302+
// a freshly-focused tab 02 has its own (empty) attachment state
303+
await omnibar.didSwitchToTab('02', ['01', '02']);
304+
await expect(omnibar.attachmentChips()).toHaveCount(0);
305+
306+
// returning to 01 restores both chips
307+
await omnibar.didSwitchToTab('01', ['01', '02']);
308+
await expect(omnibar.attachmentChips().locator('[data-status="ready"]')).toHaveCount(2);
309+
});
310+
311+
test('attached files survive a browser-tab switch', async ({ page }, workerInfo) => {
312+
const { ntp, omnibar } = setup(page, workerInfo);
313+
await ntp.reducedMotion();
314+
await ntp.openPage({ additional: config });
315+
await omnibar.ready();
316+
317+
// wait for the attach entry point so the model (and thus the file route) has resolved
318+
await expect(omnibar.attachMenuButton()).toBeVisible();
319+
await omnibar.fileInput().setInputFiles({ name: 'q3-report.pdf', mimeType: 'application/pdf', buffer: PDF_BYTES });
320+
await expect(omnibar.fileChip()).toHaveCount(1);
321+
322+
await omnibar.didSwitchToTab('02', ['01', '02']);
323+
await expect(omnibar.fileChip()).toHaveCount(0);
324+
325+
await omnibar.didSwitchToTab('01', ['01', '02']);
326+
await expect(omnibar.fileChip()).toHaveCount(1);
327+
});
328+
329+
test('attached images survive a browser-tab switch', async ({ page }, workerInfo) => {
330+
const { ntp, omnibar } = setup(page, workerInfo);
331+
await ntp.reducedMotion();
332+
await ntp.openPage({ additional: config });
333+
await omnibar.ready();
334+
335+
// wait for the attach entry point so the model (and thus the image route) has resolved
336+
await expect(omnibar.attachMenuButton()).toBeVisible();
337+
await omnibar.fileInput().setInputFiles({ name: 'shot.png', mimeType: 'image/png', buffer: TINY_PNG });
338+
await expect(omnibar.attachmentChips().locator('[data-attachment-kind="image"]')).toHaveCount(1);
339+
340+
await omnibar.didSwitchToTab('02', ['01', '02']);
341+
await expect(omnibar.attachmentChips()).toHaveCount(0);
342+
343+
await omnibar.didSwitchToTab('01', ['01', '02']);
344+
await expect(omnibar.attachmentChips().locator('[data-attachment-kind="image"]')).toHaveCount(1);
345+
});
346+
347+
test('closing an NTP tab clears its persisted attachments', async ({ page }, workerInfo) => {
348+
const { ntp, omnibar } = setup(page, workerInfo);
349+
await ntp.reducedMotion();
350+
await ntp.openPage({ additional: config });
351+
await omnibar.ready();
352+
353+
// attach on tab 01, then open tab 02 (01 still present, so it's preserved)
354+
await omnibar.attachTab('Starbucks Coffee Company');
355+
await expect(omnibar.attachmentChips().locator('[data-status="ready"]')).toHaveCount(1);
356+
await omnibar.didSwitchToTab('02', ['01', '02']);
357+
358+
// tab 01 closes (drops out of the tab id list) → its stored attachments are pruned
359+
await omnibar.didSwitchToTab('02', ['02']);
360+
361+
// 01 comes back around with nothing restored
362+
await omnibar.didSwitchToTab('01', ['01', '02']);
363+
await expect(omnibar.attachmentChips()).toHaveCount(0);
364+
});
365+
});

‎special-pages/pages/new-tab/app/tabs/PersistentValue.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
/**
2-
* @template {string} T - the value to hold.
2+
* @template T - the value to hold. Keys are always tab ids (strings); the value
3+
* can be any type (e.g. a string query, a mode, or an array of attachments).
34
*/
45
export class PersistentValue {
56
/** @type {Map<string, T>} */

0 commit comments

Comments
 (0)