diff --git a/apps/agent/entrypoints/background/index.ts b/apps/agent/entrypoints/background/index.ts index 3cdfb825..e032ebe2 100644 --- a/apps/agent/entrypoints/background/index.ts +++ b/apps/agent/entrypoints/background/index.ts @@ -18,6 +18,7 @@ import { syncScheduledJobs, } from '@/lib/schedules/scheduleStorage' import { searchActionsStorage } from '@/lib/search-actions/searchActionsStorage' +import { removeTabFromSharedSidepanelSession } from '@/lib/sidepanel/shared-sidepanel-session' import { stopAgentStorage } from '@/lib/stop-agent/stop-agent-storage' import { scheduledJobRuns } from './scheduledJobRuns' @@ -31,6 +32,10 @@ export default defineBackground(() => { scheduledJobRuns() + chrome.tabs.onRemoved.addListener((tabId) => { + removeTabFromSharedSidepanelSession(tabId).catch(() => null) + }) + chrome.action.onClicked.addListener(async (tab) => { if (tab.id) { await toggleSidePanel(tab.id) diff --git a/apps/agent/entrypoints/sidepanel/index/useChatSession.ts b/apps/agent/entrypoints/sidepanel/index/useChatSession.ts index c9d13e7a..691df60f 100644 --- a/apps/agent/entrypoints/sidepanel/index/useChatSession.ts +++ b/apps/agent/entrypoints/sidepanel/index/useChatSession.ts @@ -6,6 +6,7 @@ import { useSearchParams } from 'react-router' import useDeepCompareEffect from 'use-deep-compare-effect' import type { Provider } from '@/components/chat/chatComponentTypes' import { Capabilities, Feature } from '@/lib/browseros/capabilities' +import { isSidePanelOpen } from '@/lib/browseros/toggleSidePanel' import { useAgentServerUrl } from '@/lib/browseros/useBrowserOSProviders' import type { ChatAction } from '@/lib/chat-actions/types' import { @@ -26,6 +27,16 @@ import { useGraphqlQuery } from '@/lib/graphql/useGraphqlQuery' import { useLlmProviders } from '@/lib/llm-providers/useLlmProviders' import { track } from '@/lib/metrics/track' import { searchActionsStorage } from '@/lib/search-actions/searchActionsStorage' +import { + getSharedSidepanelConversation, + saveSharedSidepanelConversation, + watchSharedSidepanelConversations, +} from '@/lib/sidepanel/shared-sidepanel-conversation' +import { + ensureSharedSidepanelSession, + getSharedSidepanelSessionForTab, + watchSharedSidepanelSessionForTab, +} from '@/lib/sidepanel/shared-sidepanel-session' import { stopAgentStorage } from '@/lib/stop-agent/stop-agent-storage' import { selectedWorkspaceStorage } from '@/lib/workspace/workspace-storage' import type { ChatMode } from './chatTypes' @@ -73,6 +84,10 @@ export interface ChatSessionOptions { const NEWTAB_SYSTEM_PROMPT = `IMPORTANT: The user is chatting from the New Tab page. When performing browser actions, ALWAYS open content in a NEW TAB rather than navigating the current tab. The user's new tab page should remain accessible.` +function haveMessagesChanged(left: UIMessage[], right: UIMessage[]): boolean { + return JSON.stringify(left) !== JSON.stringify(right) +} + export const useChatSession = (options?: ChatSessionOptions) => { const { selectedLlmProviderRef, @@ -100,6 +115,7 @@ export const useChatSession = (options?: ChatSessionOptions) => { } = useRemoteConversationSave() const [searchParams, setSearchParams] = useSearchParams() const conversationIdParam = searchParams.get('conversationId') + const isSidepanelOrigin = options?.origin !== 'newtab' const agentUrlRef = useRef(agentServerUrl) @@ -121,11 +137,93 @@ export const useChatSession = (options?: ChatSessionOptions) => { const [disliked, setDisliked] = useState>({}) const [conversationId, setConversationId] = useState(crypto.randomUUID()) const conversationIdRef = useRef(conversationId) + const [currentTabId, setCurrentTabId] = useState(null) + const [sharedConversationId, setSharedConversationId] = useState< + string | null + >(null) + const [isResolvingSharedConversation, setIsResolvingSharedConversation] = + useState(isSidepanelOrigin) useEffect(() => { conversationIdRef.current = conversationId }, [conversationId]) + useEffect(() => { + if (!isSidepanelOrigin) { + setIsResolvingSharedConversation(false) + return + } + + let cancelled = false + + const syncCurrentTab = async () => { + const activeTabs = await chrome.tabs.query({ + active: true, + currentWindow: true, + }) + if (cancelled) return + setCurrentTabId(activeTabs[0]?.id ?? null) + } + + syncCurrentTab().catch(() => { + if (!cancelled) { + setCurrentTabId(null) + } + }) + + const handleActivated = () => { + syncCurrentTab().catch(() => { + if (!cancelled) { + setCurrentTabId(null) + } + }) + } + + chrome.tabs.onActivated.addListener(handleActivated) + + return () => { + cancelled = true + chrome.tabs.onActivated.removeListener(handleActivated) + } + }, [isSidepanelOrigin]) + + useEffect(() => { + if (!isSidepanelOrigin) return + + if (!currentTabId) { + setSharedConversationId(null) + setIsResolvingSharedConversation(false) + return + } + + let cancelled = false + setIsResolvingSharedConversation(true) + + getSharedSidepanelSessionForTab(currentTabId) + .then((session) => { + if (cancelled) return + setSharedConversationId(session?.conversationId ?? null) + setIsResolvingSharedConversation(false) + }) + .catch(() => { + if (cancelled) return + setSharedConversationId(null) + setIsResolvingSharedConversation(false) + }) + + const unwatch = watchSharedSidepanelSessionForTab( + currentTabId, + (session) => { + setSharedConversationId(session?.conversationId ?? null) + }, + ) + + return () => { + cancelled = true + unwatch() + } + }, [currentTabId, isSidepanelOrigin]) + const onClickLike = (messageId: string) => { const { responseText, queryText } = getResponseAndQueryFromMessageId( messages, @@ -333,17 +431,20 @@ export const useChatSession = (options?: ChatSessionOptions) => { useNotifyActiveTab({ messages, status, - conversationId: conversationIdRef.current, + conversationId, + hostTabId: currentTabId, }) + const conversationIdToRestore = conversationIdParam ?? sharedConversationId + const { data: remoteConversationData, isFetched: isRemoteConversationFetched, } = useGraphqlQuery( GetConversationWithMessagesDocument, - { conversationId: conversationIdParam ?? '' }, + { conversationId: conversationIdToRestore ?? '' }, { - enabled: !!conversationIdParam && isLoggedIn, + enabled: !!conversationIdToRestore && isLoggedIn, }, ) @@ -351,53 +452,122 @@ export const useChatSession = (options?: ChatSessionOptions) => { string | null >(null) - // biome-ignore lint/correctness/useExhaustiveDependencies: restore should only run when query data arrives or conversationIdParam changes useEffect(() => { - if (!conversationIdParam) return - if (restoredConversationId === conversationIdParam) return + if (isResolvingSharedConversation) return + if (!conversationIdToRestore) return + if (restoredConversationId === conversationIdToRestore) return - if (isLoggedIn) { - if (!isRemoteConversationFetched) return + let cancelled = false - if (remoteConversationData?.conversation) { - const restoredMessages = - remoteConversationData.conversation.conversationMessages.nodes - .filter((node): node is NonNullable => node !== null) - .map((node) => node.message as UIMessage) + const finishRestore = () => { + if (cancelled) return + setRestoredConversationId(conversationIdToRestore) + if (conversationIdParam) { + setSearchParams({}, { replace: true }) + } + } + const restoreFromLocalStores = async () => { + const sharedConversation = await getSharedSidepanelConversation( + conversationIdToRestore, + ) + if (sharedConversation) { setConversationId( - conversationIdParam as ReturnType, + conversationIdToRestore as ReturnType, ) - setMessages(restoredMessages) - markMessagesAsSaved(conversationIdParam, restoredMessages) - } - setRestoredConversationId(conversationIdParam) - setSearchParams({}, { replace: true }) - } else { - const restoreLocal = async () => { - const conversations = await conversationStorage.getValue() - const conversation = conversations?.find( - (c) => c.id === conversationIdParam, + setMessages(sharedConversation.messages) + markMessagesAsSaved( + conversationIdToRestore, + sharedConversation.messages, ) + finishRestore() + return true + } + + const conversations = await conversationStorage.getValue() + const localConversation = conversations?.find( + (conversation) => conversation.id === conversationIdToRestore, + ) + + if (!localConversation) return false + + setConversationId( + conversationIdToRestore as ReturnType, + ) + setMessages(localConversation.messages) + markMessagesAsSaved(conversationIdToRestore, localConversation.messages) + finishRestore() + return true + } + + const restoreConversation = async () => { + const restoredFromLocal = await restoreFromLocalStores() + if (restoredFromLocal || cancelled) return + + if (isLoggedIn) { + if (!isRemoteConversationFetched) return + + if (remoteConversationData?.conversation) { + const restoredMessages = + remoteConversationData.conversation.conversationMessages.nodes + .filter((node): node is NonNullable => node !== null) + .map((node) => node.message as UIMessage) - if (conversation) { setConversationId( - conversation.id as ReturnType, + conversationIdToRestore as ReturnType, ) - setMessages(conversation.messages) + setMessages(restoredMessages) + markMessagesAsSaved(conversationIdToRestore, restoredMessages) } - setRestoredConversationId(conversationIdParam) - setSearchParams({}, { replace: true }) + finishRestore() + return } - restoreLocal() + + finishRestore() } - }, [conversationIdParam, remoteConversationData, isLoggedIn]) + + restoreConversation() + + return () => { + cancelled = true + } + }, [ + conversationIdParam, + conversationIdToRestore, + isLoggedIn, + isRemoteConversationFetched, + isResolvingSharedConversation, + markMessagesAsSaved, + remoteConversationData, + restoredConversationId, + setSearchParams, + setMessages, + ]) + + useEffect(() => { + const unwatch = watchSharedSidepanelConversations((store) => { + const activeConversation = store[conversationIdRef.current] + if (!activeConversation) return + + const nextMessages = activeConversation.messages + if (!haveMessagesChanged(messagesRef.current, nextMessages)) return + + setMessages(nextMessages) + }) + + return () => unwatch() + }, [setMessages]) // biome-ignore lint/correctness/useExhaustiveDependencies: only need to run when messages change useEffect(() => { messagesRef.current = messages const messagesToSave = messages.filter((m) => m.parts?.length > 0) if (messagesToSave.length > 0) { + saveSharedSidepanelConversation( + conversationIdRef.current, + messagesToSave, + ).catch(() => null) + if (isLoggedIn) { if (status !== 'streaming') { saveRemoteConversation(conversationIdRef.current, messagesToSave) @@ -408,6 +578,36 @@ export const useChatSession = (options?: ChatSessionOptions) => { } }, [messages, isLoggedIn, status]) + useEffect(() => { + if (!isSidepanelOrigin || !currentTabId) return + if (isResolvingSharedConversation) return + if ( + conversationIdToRestore && + restoredConversationId !== conversationIdToRestore + ) { + return + } + + const syncSharedSession = async () => { + const panelOpen = await isSidePanelOpen(currentTabId) + if (!panelOpen) return + + await ensureSharedSidepanelSession({ + tabId: currentTabId, + conversationId, + }) + } + + syncSharedSession().catch(() => null) + }, [ + conversationId, + conversationIdToRestore, + currentTabId, + isResolvingSharedConversation, + isSidepanelOrigin, + restoredConversationId, + ]) + const sendMessage = (params: { text: string; action?: ChatAction }) => { track(MESSAGE_SENT_EVENT, { mode, @@ -470,6 +670,7 @@ export const useChatSession = (options?: ChatSessionOptions) => { stop() setConversationId(crypto.randomUUID()) setMessages([]) + setSharedConversationId(null) setTextToAction(new Map()) setLiked({}) setDisliked({}) @@ -478,7 +679,9 @@ export const useChatSession = (options?: ChatSessionOptions) => { } const isRestoringConversation = - !!conversationIdParam && restoredConversationId !== conversationIdParam + !isResolvingSharedConversation && + !!conversationIdToRestore && + restoredConversationId !== conversationIdToRestore return { mode, diff --git a/apps/agent/entrypoints/sidepanel/index/useNotifyActiveTab.tsx b/apps/agent/entrypoints/sidepanel/index/useNotifyActiveTab.tsx index 14483e2a..aa82fa2a 100644 --- a/apps/agent/entrypoints/sidepanel/index/useNotifyActiveTab.tsx +++ b/apps/agent/entrypoints/sidepanel/index/useNotifyActiveTab.tsx @@ -1,7 +1,9 @@ import type { ChatStatus, ToolUIPart, UIMessage } from 'ai' import { useEffect, useRef } from 'react' import type { GlowMessage } from '@/entrypoints/glow.content/GlowMessage' +import { openSidePanel } from '@/lib/browseros/toggleSidePanel' import { firstRunConfettiShownStorage } from '@/lib/onboarding/onboardingStorage' +import { linkTabToSharedSidepanelSession } from '@/lib/sidepanel/shared-sidepanel-session' function extractTabId(toolPart: ToolUIPart | null): number | undefined { if (!toolPart) return undefined @@ -19,25 +21,36 @@ function extractTabId(toolPart: ToolUIPart | null): number | undefined { return input?.tabId } +function extractTabIds(parts: UIMessage['parts'] | undefined): number[] { + if (!parts) return [] + + return [ + ...new Set( + parts + .filter((part): part is ToolUIPart => part.type?.startsWith('tool-')) + .map((part) => extractTabId(part)) + .filter((tabId): tabId is number => tabId !== undefined), + ), + ] +} + export const useNotifyActiveTab = ({ messages, status, conversationId, + hostTabId, }: { messages: UIMessage[] status: ChatStatus conversationId: string + hostTabId: number | null }) => { const lastTabIdRef = useRef(null) const lastMessage = messages?.[messages.length - 1] - - const latestTool = - lastMessage?.parts?.findLast((part) => part?.type?.startsWith('tool-')) ?? - null - - const hasToolCalls = !!latestTool - const toolTabId = extractTabId(latestTool as ToolUIPart | null) + const toolTabIds = extractTabIds(lastMessage?.parts) + const hasToolCalls = toolTabIds.length > 0 + const toolTabId = toolTabIds.at(-1) useEffect(() => { const isStreaming = status === 'streaming' @@ -73,15 +86,24 @@ export const useNotifyActiveTab = ({ let targetTabId = toolTabId if (!targetTabId) { - const tabs = await chrome.tabs.query({ - active: true, - currentWindow: true, - }) - targetTabId = tabs[0]?.id + targetTabId = hostTabId ?? undefined } if (cancelled || !targetTabId) return + if (hostTabId) { + for (const linkedTabId of toolTabIds) { + await linkTabToSharedSidepanelSession({ + sourceTabId: hostTabId, + targetTabId: linkedTabId, + conversationId, + }).catch(() => null) + await openSidePanel(linkedTabId).catch(() => null) + } + } + + await openSidePanel(targetTabId).catch(() => null) + if (previousTabId && previousTabId !== targetTabId) { const deactivateMessage: GlowMessage = { conversationId, @@ -105,7 +127,7 @@ export const useNotifyActiveTab = ({ return () => { cancelled = true } - }, [conversationId, status, hasToolCalls, toolTabId]) + }, [conversationId, hostTabId, status, hasToolCalls, toolTabId, toolTabIds]) return } diff --git a/apps/agent/lib/browseros/toggleSidePanel.ts b/apps/agent/lib/browseros/toggleSidePanel.ts index 5e56897f..2da94385 100644 --- a/apps/agent/lib/browseros/toggleSidePanel.ts +++ b/apps/agent/lib/browseros/toggleSidePanel.ts @@ -1,11 +1,18 @@ +/** + * @public + */ +export async function isSidePanelOpen(tabId: number): Promise { + // @ts-expect-error browserosIsOpen is a BrowserOS-specific API + return await chrome.sidePanel.browserosIsOpen({ tabId }) +} + /** * @public */ export async function openSidePanel( tabId: number, ): Promise<{ opened: boolean }> { - // @ts-expect-error browserosIsOpen is a BrowserOS-specific API - const isAlreadyOpen = await chrome.sidePanel.browserosIsOpen({ tabId }) + const isAlreadyOpen = await isSidePanelOpen(tabId) if (isAlreadyOpen) { return { opened: true } } diff --git a/apps/agent/lib/sidepanel/shared-sidepanel-conversation.ts b/apps/agent/lib/sidepanel/shared-sidepanel-conversation.ts new file mode 100644 index 00000000..4bc16754 --- /dev/null +++ b/apps/agent/lib/sidepanel/shared-sidepanel-conversation.ts @@ -0,0 +1,62 @@ +import { storage } from '@wxt-dev/storage' +import type { UIMessage } from 'ai' + +export interface SharedSidepanelConversation { + messages: UIMessage[] + updatedAt: number +} + +export type SharedSidepanelConversationStore = Record< + string, + SharedSidepanelConversation +> + +const EMPTY_SHARED_SIDEPANEL_CONVERSATIONS: SharedSidepanelConversationStore = + {} + +const sharedSidepanelConversationStorage = + storage.defineItem( + 'local:shared-sidepanel-conversations', + { + fallback: EMPTY_SHARED_SIDEPANEL_CONVERSATIONS, + }, + ) + +function haveMessagesChanged(left: UIMessage[], right: UIMessage[]): boolean { + return JSON.stringify(left) !== JSON.stringify(right) +} + +export async function getSharedSidepanelConversation( + conversationId: string, +): Promise { + const store = await sharedSidepanelConversationStorage.getValue() + return store[conversationId] ?? null +} + +export async function saveSharedSidepanelConversation( + conversationId: string, + messages: UIMessage[], +): Promise { + const store = await sharedSidepanelConversationStorage.getValue() + const existing = store[conversationId] + + if (existing && !haveMessagesChanged(existing.messages, messages)) { + return + } + + await sharedSidepanelConversationStorage.setValue({ + ...store, + [conversationId]: { + messages, + updatedAt: Date.now(), + }, + }) +} + +export function watchSharedSidepanelConversations( + callback: (store: SharedSidepanelConversationStore) => void, +): () => void { + return sharedSidepanelConversationStorage.watch((store) => { + callback(store ?? EMPTY_SHARED_SIDEPANEL_CONVERSATIONS) + }) +} diff --git a/apps/agent/lib/sidepanel/shared-sidepanel-session-state.test.ts b/apps/agent/lib/sidepanel/shared-sidepanel-session-state.test.ts new file mode 100644 index 00000000..55f16588 --- /dev/null +++ b/apps/agent/lib/sidepanel/shared-sidepanel-session-state.test.ts @@ -0,0 +1,137 @@ +import { describe, expect, it } from 'bun:test' +import { + EMPTY_SHARED_SIDEPANEL_STATE, + ensureSharedSidepanelSessionState, + getSharedSidepanelSessionForTabFromState, + linkTabToSharedSidepanelSessionState, + removeSharedSidepanelSessionTabState, +} from './shared-sidepanel-session-state' + +describe('shared sidepanel session state', () => { + it('creates a new shared session for the originating tab', () => { + const nextState = ensureSharedSidepanelSessionState( + EMPTY_SHARED_SIDEPANEL_STATE, + { + tabId: 11, + conversationId: 'conversation-1', + sessionId: 'session-1', + updatedAt: 100, + }, + ) + + expect(getSharedSidepanelSessionForTabFromState(nextState, 11)).toEqual({ + id: 'session-1', + rootTabId: 11, + conversationId: 'conversation-1', + tabIds: [11], + updatedAt: 100, + }) + }) + + it('links an agent-targeted tab into the source session', () => { + const initialState = ensureSharedSidepanelSessionState( + EMPTY_SHARED_SIDEPANEL_STATE, + { + tabId: 11, + conversationId: 'conversation-1', + sessionId: 'session-1', + updatedAt: 100, + }, + ) + + const nextState = linkTabToSharedSidepanelSessionState(initialState, { + sourceTabId: 11, + targetTabId: 42, + conversationId: 'conversation-1', + updatedAt: 200, + }) + + expect(getSharedSidepanelSessionForTabFromState(nextState, 42)).toEqual({ + id: 'session-1', + rootTabId: 11, + conversationId: 'conversation-1', + tabIds: [11, 42], + updatedAt: 200, + }) + }) + + it('updates the shared conversation for an existing session', () => { + const initialState = ensureSharedSidepanelSessionState( + EMPTY_SHARED_SIDEPANEL_STATE, + { + tabId: 11, + conversationId: 'conversation-1', + sessionId: 'session-1', + updatedAt: 100, + }, + ) + + const nextState = ensureSharedSidepanelSessionState(initialState, { + tabId: 11, + conversationId: 'conversation-2', + updatedAt: 300, + }) + + expect(getSharedSidepanelSessionForTabFromState(nextState, 11)).toEqual({ + id: 'session-1', + rootTabId: 11, + conversationId: 'conversation-2', + tabIds: [11], + updatedAt: 300, + }) + }) + + it('moves a tab out of its previous session when re-linked', () => { + const firstSession = ensureSharedSidepanelSessionState( + EMPTY_SHARED_SIDEPANEL_STATE, + { + tabId: 11, + conversationId: 'conversation-1', + sessionId: 'session-1', + updatedAt: 100, + }, + ) + + const secondSession = ensureSharedSidepanelSessionState(firstSession, { + tabId: 42, + conversationId: 'conversation-2', + sessionId: 'session-2', + updatedAt: 150, + }) + + const relinkedState = linkTabToSharedSidepanelSessionState(secondSession, { + sourceTabId: 11, + targetTabId: 42, + conversationId: 'conversation-1', + updatedAt: 200, + }) + + expect(getSharedSidepanelSessionForTabFromState(relinkedState, 42)).toEqual( + { + id: 'session-1', + rootTabId: 11, + conversationId: 'conversation-1', + tabIds: [11, 42], + updatedAt: 200, + }, + ) + expect(relinkedState.sessions['session-2']).toBeUndefined() + }) + + it('removes empty sessions when the last tab is cleared', () => { + const initialState = ensureSharedSidepanelSessionState( + EMPTY_SHARED_SIDEPANEL_STATE, + { + tabId: 11, + conversationId: 'conversation-1', + sessionId: 'session-1', + updatedAt: 100, + }, + ) + + const nextState = removeSharedSidepanelSessionTabState(initialState, 11) + + expect(getSharedSidepanelSessionForTabFromState(nextState, 11)).toBeNull() + expect(nextState.sessions['session-1']).toBeUndefined() + }) +}) diff --git a/apps/agent/lib/sidepanel/shared-sidepanel-session-state.ts b/apps/agent/lib/sidepanel/shared-sidepanel-session-state.ts new file mode 100644 index 00000000..84f03775 --- /dev/null +++ b/apps/agent/lib/sidepanel/shared-sidepanel-session-state.ts @@ -0,0 +1,233 @@ +export interface SharedSidepanelSession { + id: string + rootTabId: number + conversationId: string + tabIds: number[] + updatedAt: number +} + +interface SharedSidepanelTabLink { + sessionId: string +} + +export interface SharedSidepanelState { + tabs: Record + sessions: Record +} + +interface EnsureSharedSidepanelSessionParams { + tabId: number + conversationId: string + sessionId?: string + rootTabId?: number + updatedAt?: number +} + +interface LinkTabToSharedSidepanelSessionParams { + sourceTabId: number + targetTabId: number + conversationId: string + sessionId?: string + updatedAt?: number +} + +export const EMPTY_SHARED_SIDEPANEL_STATE: SharedSidepanelState = { + tabs: {}, + sessions: {}, +} + +function getTabKey(tabId: number): string { + return String(tabId) +} + +function uniqTabIds(tabIds: number[]): number[] { + return [...new Set(tabIds)] +} + +export function getSharedSidepanelSessionForTabFromState( + state: SharedSidepanelState, + tabId: number, +): SharedSidepanelSession | null { + const sessionId = state.tabs[getTabKey(tabId)]?.sessionId + if (!sessionId) return null + return state.sessions[sessionId] ?? null +} + +function removeTabLinkFromSession( + state: SharedSidepanelState, + tabId: number, +): SharedSidepanelState { + const tabKey = getTabKey(tabId) + const sessionId = state.tabs[tabKey]?.sessionId + if (!sessionId) return state + + const session = state.sessions[sessionId] + const nextTabs = { ...state.tabs } + delete nextTabs[tabKey] + + if (!session) { + return { + tabs: nextTabs, + sessions: state.sessions, + } + } + + const remainingTabIds = session.tabIds.filter((id) => id !== tabId) + const nextSessions = { ...state.sessions } + + if (remainingTabIds.length === 0) { + delete nextSessions[sessionId] + } else { + nextSessions[sessionId] = { + ...session, + tabIds: remainingTabIds, + updatedAt: Date.now(), + } + } + + return { + tabs: nextTabs, + sessions: nextSessions, + } +} + +function upsertTabIntoSession( + state: SharedSidepanelState, + { + tabId, + sessionId, + conversationId, + rootTabId, + updatedAt = Date.now(), + }: Required> & + Pick< + EnsureSharedSidepanelSessionParams, + 'conversationId' | 'rootTabId' | 'updatedAt' + >, +): SharedSidepanelState { + const tabKey = getTabKey(tabId) + const currentSessionId = state.tabs[tabKey]?.sessionId + const currentSession = currentSessionId + ? state.sessions[currentSessionId] + : undefined + const nextRootTabId = rootTabId ?? currentSession?.rootTabId ?? tabId + let nextState = state + + if (currentSessionId && currentSessionId !== sessionId) { + nextState = removeTabLinkFromSession(nextState, tabId) + } + + const targetSession = nextState.sessions[sessionId] + const nextTabIds = uniqTabIds([...(targetSession?.tabIds ?? []), tabId]) + const tabLinkChanged = nextState.tabs[tabKey]?.sessionId !== sessionId + const sessionChanged = + !targetSession || + targetSession.conversationId !== conversationId || + targetSession.rootTabId !== nextRootTabId || + targetSession.tabIds.length !== nextTabIds.length || + targetSession.tabIds.some( + (existingTabId, index) => existingTabId !== nextTabIds[index], + ) + + if (!tabLinkChanged && !sessionChanged) { + return nextState + } + + return { + tabs: tabLinkChanged + ? { + ...nextState.tabs, + [tabKey]: { sessionId }, + } + : nextState.tabs, + sessions: { + ...nextState.sessions, + [sessionId]: { + id: sessionId, + rootTabId: nextRootTabId, + conversationId, + tabIds: nextTabIds, + updatedAt, + }, + }, + } +} + +export function ensureSharedSidepanelSessionState( + state: SharedSidepanelState, + { + tabId, + conversationId, + sessionId, + rootTabId, + updatedAt = Date.now(), + }: EnsureSharedSidepanelSessionParams, +): SharedSidepanelState { + const existingSession = getSharedSidepanelSessionForTabFromState(state, tabId) + const resolvedSessionId = existingSession?.id ?? sessionId + + if (!resolvedSessionId) { + throw new Error('sessionId is required when creating a new shared session') + } + + return upsertTabIntoSession(state, { + tabId, + sessionId: resolvedSessionId, + conversationId, + rootTabId: existingSession?.rootTabId ?? rootTabId ?? tabId, + updatedAt, + }) +} + +export function linkTabToSharedSidepanelSessionState( + state: SharedSidepanelState, + { + sourceTabId, + targetTabId, + conversationId, + sessionId, + updatedAt = Date.now(), + }: LinkTabToSharedSidepanelSessionParams, +): SharedSidepanelState { + const sourceSession = getSharedSidepanelSessionForTabFromState( + state, + sourceTabId, + ) + const targetSession = getSharedSidepanelSessionForTabFromState( + state, + targetTabId, + ) + const resolvedSessionId = sourceSession?.id ?? targetSession?.id ?? sessionId + + if (!resolvedSessionId) { + throw new Error('sessionId is required when linking a new shared session') + } + + const rootTabId = + sourceSession?.rootTabId ?? targetSession?.rootTabId ?? sourceTabId + + let nextState = upsertTabIntoSession(state, { + tabId: sourceTabId, + sessionId: resolvedSessionId, + conversationId, + rootTabId, + updatedAt, + }) + + nextState = upsertTabIntoSession(nextState, { + tabId: targetTabId, + sessionId: resolvedSessionId, + conversationId, + rootTabId, + updatedAt, + }) + + return nextState +} + +export function removeSharedSidepanelSessionTabState( + state: SharedSidepanelState, + tabId: number, +): SharedSidepanelState { + return removeTabLinkFromSession(state, tabId) +} diff --git a/apps/agent/lib/sidepanel/shared-sidepanel-session.ts b/apps/agent/lib/sidepanel/shared-sidepanel-session.ts new file mode 100644 index 00000000..640b2086 --- /dev/null +++ b/apps/agent/lib/sidepanel/shared-sidepanel-session.ts @@ -0,0 +1,105 @@ +import { storage } from '@wxt-dev/storage' +import { + EMPTY_SHARED_SIDEPANEL_STATE, + ensureSharedSidepanelSessionState, + getSharedSidepanelSessionForTabFromState, + linkTabToSharedSidepanelSessionState, + removeSharedSidepanelSessionTabState, + type SharedSidepanelSession, + type SharedSidepanelState, +} from './shared-sidepanel-session-state' + +interface EnsureSharedSidepanelSessionParams { + tabId: number + conversationId: string + sessionId?: string + rootTabId?: number + updatedAt?: number +} + +interface LinkTabToSharedSidepanelSessionParams { + sourceTabId: number + targetTabId: number + conversationId: string + sessionId?: string + updatedAt?: number +} + +const sharedSidepanelSessionStorage = storage.defineItem( + 'local:shared-sidepanel-sessions', + { + fallback: EMPTY_SHARED_SIDEPANEL_STATE, + }, +) + +async function updateSharedSidepanelState( + updater: (state: SharedSidepanelState) => SharedSidepanelState, +): Promise { + const currentState = await sharedSidepanelSessionStorage.getValue() + const safeState = currentState ?? EMPTY_SHARED_SIDEPANEL_STATE + const nextState = updater(safeState) + + if (nextState !== safeState) { + await sharedSidepanelSessionStorage.setValue(nextState) + } + + return nextState +} + +export async function getSharedSidepanelSessionForTab( + tabId: number, +): Promise { + const state = await sharedSidepanelSessionStorage.getValue() + return getSharedSidepanelSessionForTabFromState( + state ?? EMPTY_SHARED_SIDEPANEL_STATE, + tabId, + ) +} + +export async function ensureSharedSidepanelSession( + params: Omit, +): Promise { + const nextState = await updateSharedSidepanelState((state) => + ensureSharedSidepanelSessionState(state, { + ...params, + sessionId: crypto.randomUUID(), + }), + ) + + return getSharedSidepanelSessionForTabFromState(nextState, params.tabId) +} + +export async function linkTabToSharedSidepanelSession( + params: Omit, +): Promise { + const nextState = await updateSharedSidepanelState((state) => + linkTabToSharedSidepanelSessionState(state, { + ...params, + sessionId: crypto.randomUUID(), + }), + ) + + return getSharedSidepanelSessionForTabFromState(nextState, params.targetTabId) +} + +export async function removeTabFromSharedSidepanelSession( + tabId: number, +): Promise { + await updateSharedSidepanelState((state) => + removeSharedSidepanelSessionTabState(state, tabId), + ) +} + +export function watchSharedSidepanelSessionForTab( + tabId: number, + callback: (session: SharedSidepanelSession | null) => void, +): () => void { + return sharedSidepanelSessionStorage.watch((state) => { + callback( + getSharedSidepanelSessionForTabFromState( + state ?? EMPTY_SHARED_SIDEPANEL_STATE, + tabId, + ), + ) + }) +}