909 lines
27 KiB
JavaScript
909 lines
27 KiB
JavaScript
/**
|
|
* This Source Code Form is subject to the terms of the Mozilla Public
|
|
* License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
* file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
*/
|
|
|
|
/**
|
|
* @typedef {object} TabSelectionData
|
|
* @property {string} linkedPanel - ID of the linked panel (e.g., "panel-3-1")
|
|
* @property {string} url - URL of the tab
|
|
* @property {string} title - Display title of the tab
|
|
* @property {string} [iconSrc] - URL for the tab's favicon (optional)
|
|
* @property {boolean} [checked] - Whether the tab is selected in UI (optional)
|
|
*/
|
|
|
|
/**
|
|
* @typedef {object} ToolUpdateData
|
|
* @property {Array<TabSelectionData>} [selectedTabs] - Array of selected tabs
|
|
* @property {string} [operationId] - Operation ID for undo operations
|
|
* @property {boolean} [wasRestored] - Flag indicating tabs were restored
|
|
* @property {number} [restoredCount] - Number of tabs restored
|
|
* @property {Array<TabSelectionData>} [originalClosedTabs] - Original tabs that were closed
|
|
*/
|
|
|
|
/**
|
|
* @typedef {object} HandlerContext
|
|
* @property {object} message - Message containing the tool UI
|
|
* @property {string} toolCallId - ID of the tool call
|
|
* @property {ToolUpdateData} updateData - Update data for the handler
|
|
* @property {object} conversation - Conversation object
|
|
* @property {ChromeWindow} window - Browser window object
|
|
* @property {object} originalData - Original update data passed to handleUpdate
|
|
*/
|
|
|
|
const lazy = {};
|
|
|
|
ChromeUtils.defineESModuleGetters(lazy, {
|
|
tabManagementService:
|
|
"moz-src:///browser/components/aiwindow/ui/modules/TabManagementService.sys.mjs",
|
|
ToolUITelemetry:
|
|
"moz-src:///browser/components/aiwindow/ui/modules/ToolUITelemetry.sys.mjs",
|
|
MESSAGE_ROLE:
|
|
"moz-src:///browser/components/aiwindow/ui/modules/ChatEnums.sys.mjs",
|
|
});
|
|
|
|
ChromeUtils.defineLazyGetter(lazy, "console", function () {
|
|
return console.createInstance({
|
|
prefix: "ToolUI",
|
|
});
|
|
});
|
|
|
|
/**
|
|
* UI labels for tool results and follow-ups.
|
|
*/
|
|
export const UI_TYPES = {
|
|
WEBSITE_CONFIRMATION: "website-confirmation",
|
|
TAB_GROUP_CONFIRMATION: "tab-group-confirmation",
|
|
AI_ACTION_RESULT: "ai-action-result",
|
|
CANCELLED_COMPONENT: "cancelled-component",
|
|
RETRY_COMPONENT: "retry-component",
|
|
};
|
|
|
|
/**
|
|
* UI update types for communicating user interactions with tool UIs back to the actor.
|
|
*/
|
|
export const UI_UPDATE_TYPES = {
|
|
CONFIRMATION_TAB_SELECTION: "confirmation-tab-selection",
|
|
CANCEL_TAB_SELECTION: "cancel-tab-selection",
|
|
CONFIRM_TAB_GROUP_SELECTION: "confirm-tab-group-selection",
|
|
UNDO_TAB_CLOSE: "undo-tab-close",
|
|
UNDO_TAB_GROUP: "undo-tab-group",
|
|
RETRY_PROMPT: "retry-prompt",
|
|
};
|
|
|
|
/**
|
|
* UI types that represent confirmation dialogs with stateful behavior
|
|
* These UIs can be cancelled and need special handling when restored from DB
|
|
*/
|
|
export const CONFIRMATION_UI_TYPES = [
|
|
UI_TYPES.WEBSITE_CONFIRMATION,
|
|
UI_TYPES.TAB_GROUP_CONFIRMATION,
|
|
];
|
|
|
|
/**
|
|
* Manages the Tool UI updates and orchestrates state changes for tool UI components
|
|
*/
|
|
export class ToolUI {
|
|
/**
|
|
* Per-confirmation map of selection token to browser permanentKey, keyed by
|
|
* toolCallId. Populated when a tab confirmation is shown and read back when
|
|
* the user confirms, so each selected tab can be resolved back to a stable
|
|
* identity
|
|
*
|
|
* @type {Map<string, Map<string, object>>} toolCallId -> (token -> permanentKey)
|
|
*/
|
|
static #tabKeysByToolCall = new Map();
|
|
|
|
/**
|
|
* Register the token to permanentKey map for a close tabs confirmation.
|
|
*
|
|
* @param {string} toolCallId
|
|
* @param {Map<string, object>} tokenToKey
|
|
*/
|
|
static registerTabKeys(toolCallId, tokenToKey) {
|
|
if (toolCallId && tokenToKey?.size) {
|
|
this.#tabKeysByToolCall.set(toolCallId, tokenToKey);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Drop a confirmation's token map once it resolves or is cancelled
|
|
*
|
|
* @param {string} toolCallId
|
|
*/
|
|
static clearTabKeys(toolCallId) {
|
|
this.#tabKeysByToolCall.delete(toolCallId);
|
|
}
|
|
|
|
static #getConfirmationReason(tabs) {
|
|
if (tabs.some(t => t.pinned)) {
|
|
return "pinned_tab";
|
|
}
|
|
if (tabs.some(t => t.selected)) {
|
|
return "active_tab";
|
|
}
|
|
if (tabs.length === 1) {
|
|
return "last_tab";
|
|
}
|
|
return "user_action";
|
|
}
|
|
|
|
/**
|
|
* Resolve selected tabs to live tab objects in the given window by
|
|
* permanentKey
|
|
*
|
|
* @param {Array<TabSelectionData>} selectedTabs - Selected tabs
|
|
* @param {Map<string, object>} tokenToKey - token -> permanentKey for this operation
|
|
* @param {ChromeWindow} win - The browser window object
|
|
* @returns {Array<Tab>|null} Verified tab objects or null if none valid
|
|
* @private
|
|
*/
|
|
static #verifyAndCollectTabs(selectedTabs = [], tokenToKey = null, win) {
|
|
if (!win) {
|
|
lazy.console.error("No browser window provided");
|
|
return null;
|
|
}
|
|
if (!tokenToKey?.size) {
|
|
lazy.console.warn("No tab-key map for this operation");
|
|
return null;
|
|
}
|
|
|
|
const claimedTabs = new Set();
|
|
const verifiedTabObjects = [];
|
|
|
|
for (const selectedTab of selectedTabs) {
|
|
const permanentKey =
|
|
selectedTab.token && tokenToKey.get(selectedTab.token);
|
|
const tab =
|
|
permanentKey &&
|
|
win.gBrowser.tabs.find(
|
|
t => !claimedTabs.has(t) && t.permanentKey === permanentKey
|
|
);
|
|
|
|
if (!tab) {
|
|
lazy.console.warn(
|
|
`No live tab for selection ${selectedTab.url} (token ${selectedTab.token})`
|
|
);
|
|
continue;
|
|
}
|
|
|
|
claimedTabs.add(tab);
|
|
verifiedTabObjects.push(tab);
|
|
}
|
|
|
|
if (verifiedTabObjects.length === 0) {
|
|
lazy.console.warn("No valid tabs after verification");
|
|
return null;
|
|
}
|
|
|
|
return verifiedTabObjects;
|
|
}
|
|
|
|
/**
|
|
* Close the selected tabs after resolving them by permanentKey
|
|
*
|
|
* @param {Array<TabSelectionData>} selectedTabs - Selected tabs
|
|
* @param {Map<string, object>} tokenToKey - token -> permanentKey for this operation
|
|
* @param {ChromeWindow} win - The browser window object
|
|
* @returns {Promise<{operationId: string, closedTabs: Array, failedTabs: Array}|null>}
|
|
*/
|
|
static async closeSelectedTabs(selectedTabs = [], tokenToKey, win) {
|
|
const verifiedTabObjects = this.#verifyAndCollectTabs(
|
|
selectedTabs,
|
|
tokenToKey,
|
|
win
|
|
);
|
|
if (!verifiedTabObjects) {
|
|
return null;
|
|
}
|
|
|
|
const activeTab = verifiedTabObjects.find(
|
|
tab => tab === win.gBrowser.selectedTab
|
|
);
|
|
if (activeTab) {
|
|
activeTab.smartWindowActionSource = "close_current_tab";
|
|
}
|
|
|
|
const result = await lazy.tabManagementService.closeTabs({
|
|
tabs: verifiedTabObjects,
|
|
window: win,
|
|
});
|
|
|
|
return result;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Tool UI Update Handlers
|
|
* ======================================================================== */
|
|
|
|
/**
|
|
* Handler for tab selection confirmation
|
|
*
|
|
* @param {HandlerContext} context - Handler context
|
|
* @returns {Promise<boolean>} True if successful
|
|
* @private
|
|
*/
|
|
static async #handleConfirmationTabSelection(context) {
|
|
const {
|
|
updateData,
|
|
message,
|
|
conversation,
|
|
window,
|
|
originalData,
|
|
mode,
|
|
toolCallId,
|
|
} = context;
|
|
const { selectedTabs = [] } = updateData ?? {};
|
|
|
|
const tokenToKey = this.#tabKeysByToolCall.get(toolCallId);
|
|
const result = await this.closeSelectedTabs(
|
|
selectedTabs,
|
|
tokenToKey,
|
|
window
|
|
);
|
|
this.clearTabKeys(toolCallId);
|
|
if (!result) {
|
|
return false;
|
|
}
|
|
|
|
// Record telemetry for browser action prompt response
|
|
lazy.ToolUITelemetry.recordBrowserActionPromptResponse({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "close_tabs",
|
|
prompt_type: "safety_confirmation",
|
|
response: "confirm",
|
|
selected: selectedTabs.length,
|
|
reason: "user_action",
|
|
});
|
|
|
|
// Include the operationId in the update data for potential undo
|
|
const enhancedData = {
|
|
...originalData,
|
|
updateData: {
|
|
...updateData,
|
|
operationId: result.operationId,
|
|
actionTimestamp: Date.now(),
|
|
actionType: "close_tabs",
|
|
},
|
|
};
|
|
|
|
conversation.updateToolUI(message, enhancedData, UI_TYPES.AI_ACTION_RESULT);
|
|
|
|
const confirmationMessage = {
|
|
description:
|
|
"User confirmed the requested action. selectedTabs contains the tabs that were acted upon.",
|
|
selectedTabs: selectedTabs.map(({ url, title }) => ({ url, title })),
|
|
};
|
|
|
|
const pendingAction = conversation.messages.at(-1)?.content?.body?.action;
|
|
if (pendingAction) {
|
|
confirmationMessage.action = pendingAction;
|
|
}
|
|
conversation.resolvePendingToolConfirmation(
|
|
confirmationMessage,
|
|
toolCallId
|
|
);
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Handler for tab selection cancellation
|
|
*
|
|
* @param {HandlerContext} context - Handler context
|
|
* @returns {boolean} True if successful
|
|
* @private
|
|
*/
|
|
static #handleCancelTabSelection(context) {
|
|
const {
|
|
message,
|
|
conversation,
|
|
originalData,
|
|
mode,
|
|
updateData,
|
|
toolCallId,
|
|
} = context;
|
|
|
|
// Use the provided reason or default to user_action for manual cancellations
|
|
const reason = updateData?.reason || "user_action";
|
|
const actionType = updateData?.actionType;
|
|
|
|
// Record telemetry for browser action prompt response (cancellation)
|
|
lazy.ToolUITelemetry.recordBrowserActionPromptResponse({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: actionType,
|
|
prompt_type: "safety_confirmation",
|
|
response: "cancel",
|
|
selected: 0,
|
|
reason,
|
|
});
|
|
|
|
this.clearTabKeys(toolCallId);
|
|
|
|
conversation.updateToolUI(
|
|
message,
|
|
originalData,
|
|
UI_TYPES.CANCELLED_COMPONENT
|
|
);
|
|
conversation.resolvePendingToolConfirmation(
|
|
{ description: "User cancelled the tab action. No action was taken." },
|
|
toolCallId
|
|
);
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Handler for tab group confirmation
|
|
*
|
|
* @param {HandlerContext} context - Handler context
|
|
* @returns {Promise<boolean>} True if successful
|
|
* @private
|
|
*/
|
|
static async #handleConfirmTabGroupSelection(context) {
|
|
const {
|
|
updateData,
|
|
message,
|
|
conversation,
|
|
window,
|
|
originalData,
|
|
mode,
|
|
toolCallId,
|
|
} = context;
|
|
const { selectedTabs = [], tabGroupLabel = "Tab Group" } = updateData ?? {};
|
|
|
|
const tokenToKey = this.#tabKeysByToolCall.get(toolCallId);
|
|
const result = await this.createTabGroup({
|
|
tabs: selectedTabs,
|
|
tokenToKey,
|
|
window,
|
|
label: tabGroupLabel,
|
|
});
|
|
this.clearTabKeys(toolCallId);
|
|
if (!result?.success) {
|
|
return false;
|
|
}
|
|
|
|
// Record telemetry for browser action prompt response
|
|
lazy.ToolUITelemetry.recordBrowserActionPromptResponse({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "group_tabs",
|
|
prompt_type: "safety_confirmation",
|
|
response: "confirm",
|
|
selected: selectedTabs.length,
|
|
reason: "user_action",
|
|
});
|
|
|
|
// Include the group data in the update data
|
|
const enhancedData = {
|
|
...originalData,
|
|
updateData: {
|
|
...updateData,
|
|
operationId: result.group?.id,
|
|
actionTimestamp: Date.now(),
|
|
actionType: "group_tabs",
|
|
group: result.group,
|
|
},
|
|
};
|
|
|
|
conversation.updateToolUI(message, enhancedData, UI_TYPES.AI_ACTION_RESULT);
|
|
|
|
const confirmationMessage = {
|
|
description:
|
|
"User confirmed the requested action. selectedTabs contains the tabs that were acted upon.",
|
|
selectedTabs: selectedTabs.map(({ url, title }) => ({ url, title })),
|
|
};
|
|
|
|
const pendingAction = conversation.messages.at(-1)?.content?.body?.action;
|
|
if (pendingAction) {
|
|
confirmationMessage.action = pendingAction;
|
|
}
|
|
conversation.resolvePendingToolConfirmation(
|
|
confirmationMessage,
|
|
toolCallId
|
|
);
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Handler for undoing tab group operation
|
|
*
|
|
* @param {HandlerContext} context - Handler context
|
|
* @returns {Promise<boolean>} True if successful
|
|
* @private
|
|
*/
|
|
static async #handleUndoTabGroup(context) {
|
|
const { updateData, message, conversation, window, originalData, mode } =
|
|
context;
|
|
const { operationId, actionTimestamp } = updateData ?? {};
|
|
const undoStartTime = Date.now();
|
|
|
|
if (!operationId) {
|
|
lazy.console.error("ToolUI: No operation ID provided for undo tab group");
|
|
return false;
|
|
}
|
|
|
|
// The operationId is the group ID for tab groups
|
|
const groupId = operationId;
|
|
|
|
// Attempt to ungroup the tabs
|
|
const result = await lazy.tabManagementService.ungroupTabs({
|
|
groupId,
|
|
window,
|
|
});
|
|
|
|
if (!result?.success) {
|
|
lazy.console.error(
|
|
"ToolUI: Failed to undo tab group:",
|
|
result?.error || "Unknown error"
|
|
);
|
|
|
|
const timeDelta = actionTimestamp ? undoStartTime - actionTimestamp : 0;
|
|
|
|
lazy.ToolUITelemetry.recordBrowserActionUndo({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "group_tabs",
|
|
tabs_restored: 0,
|
|
time_delta: Math.max(0, timeDelta),
|
|
result: "error",
|
|
error: result?.error || "ungroup_failed",
|
|
});
|
|
|
|
return false;
|
|
}
|
|
|
|
// Calculate time delta from when action completed to when undo was clicked
|
|
const timeDelta = actionTimestamp ? undoStartTime - actionTimestamp : 0;
|
|
|
|
// Record telemetry for browser action undo
|
|
lazy.ToolUITelemetry.recordBrowserActionUndo({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "group_tabs",
|
|
tabs_restored: result.ungroupedTabs.length,
|
|
time_delta: Math.max(0, timeDelta),
|
|
result: "success",
|
|
error: "",
|
|
});
|
|
|
|
// Update the UI to show the undo was successful
|
|
const enhancedData = {
|
|
...originalData,
|
|
updateData: {
|
|
...updateData,
|
|
wasRestored: true,
|
|
originalGroupedTabs: result.ungroupedTabs,
|
|
actionType: "group_tabs", // Preserve the action type
|
|
},
|
|
};
|
|
|
|
conversation.updateToolUI(message, enhancedData, UI_TYPES.AI_ACTION_RESULT);
|
|
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Handler for undoing tab close operation
|
|
*
|
|
* @param {HandlerContext} context - Handler context
|
|
* @returns {Promise<boolean>} True if successful
|
|
* @private
|
|
*/
|
|
static async #handleUndoTabClose(context) {
|
|
const { updateData, message, conversation, window, originalData, mode } =
|
|
context;
|
|
const {
|
|
operationId,
|
|
selectedTabs = [],
|
|
actionTimestamp,
|
|
} = updateData ?? {};
|
|
const undoStartTime = Date.now();
|
|
|
|
if (!operationId) {
|
|
lazy.console.error("ToolUI: No operationId provided for undo");
|
|
return false;
|
|
}
|
|
|
|
if (!window) {
|
|
lazy.console.error("ToolUI: No window provided for undo");
|
|
return false;
|
|
}
|
|
|
|
try {
|
|
const { restoredCount, requestedCount, failedTabs } =
|
|
await lazy.tabManagementService.restoreTabs({
|
|
operationId,
|
|
window,
|
|
});
|
|
|
|
lazy.console.log(`Restored ${restoredCount} of ${requestedCount} tabs`);
|
|
|
|
// Calculate time delta from when action completed to when undo was clicked
|
|
const timeDelta = actionTimestamp ? undoStartTime - actionTimestamp : 0;
|
|
|
|
let undoResult = "success";
|
|
let errorCode = "";
|
|
|
|
if (failedTabs && failedTabs > 0) {
|
|
errorCode = "one_or_more_tabs_failed_to_restore";
|
|
undoResult = restoredCount > 0 ? "partial_success" : "error";
|
|
}
|
|
|
|
// Record telemetry for browser action undo
|
|
lazy.ToolUITelemetry.recordBrowserActionUndo({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "close_tabs",
|
|
tabs_restored: restoredCount,
|
|
time_delta: Math.max(0, timeDelta),
|
|
result: undoResult,
|
|
error: errorCode,
|
|
});
|
|
|
|
// Update the UI to show the undo was successful
|
|
const enhancedData = {
|
|
...originalData,
|
|
updateData: {
|
|
...updateData,
|
|
wasRestored: true,
|
|
restoredCount,
|
|
originalClosedTabs: selectedTabs,
|
|
actionType: "close_tabs",
|
|
},
|
|
};
|
|
|
|
conversation.updateToolUI(
|
|
message,
|
|
enhancedData,
|
|
UI_TYPES.AI_ACTION_RESULT
|
|
);
|
|
return true;
|
|
} catch (error) {
|
|
// This will only catch catastrophic errors like invalid window
|
|
// since TabManagementService has its own try/catch
|
|
lazy.console.error("Failed to restore tabs:", error);
|
|
|
|
// Calculate time delta for error case
|
|
const timeDelta = actionTimestamp ? undoStartTime - actionTimestamp : 0;
|
|
|
|
// Record telemetry for catastrophic failure
|
|
lazy.ToolUITelemetry.recordBrowserActionUndo({
|
|
location: mode,
|
|
chat_id: conversation?.id || "",
|
|
message_seq: conversation?.messages?.length || 0,
|
|
action_type: "close_tabs",
|
|
tabs_restored: 0,
|
|
time_delta: Math.max(0, timeDelta),
|
|
result: "error",
|
|
error: error?.name || "invalid_window",
|
|
});
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Handle retry prompt from the UI (clears the current tool UI)
|
|
*
|
|
* @param {object} context - The handler context
|
|
* @returns {Promise<boolean>} True if successful
|
|
* @private
|
|
*/
|
|
static async #handleRetryPrompt(context) {
|
|
const { message, conversation } = context;
|
|
await conversation.updateToolUI(message, null, null);
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Finds the last assistant text message in a conversation
|
|
*
|
|
* @param {Array} messages - The conversation messages array
|
|
* @returns {object|null} The last assistant text message or null if not found
|
|
* @private
|
|
*/
|
|
static #findLastAssistantTextMessage(messages) {
|
|
return (
|
|
messages.findLast(
|
|
message =>
|
|
message.role === lazy.MESSAGE_ROLE.ASSISTANT &&
|
|
message.content?.type === "text"
|
|
) ?? null
|
|
);
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Handler Mapping and Public API
|
|
* ======================================================================== */
|
|
|
|
/**
|
|
* Creates a tab group from selected tabs after verification.
|
|
*
|
|
* @param {object} options - Options for creating the tab group
|
|
* @param {Array<TabSelectionData>} options.tabs - Selected tabs (each carrying a `token`)
|
|
* @param {Map<string, object>} options.tokenToKey - token -> permanentKey for this operation
|
|
* @param {ChromeWindow} options.window - The browser window object
|
|
* @param {string} [options.label="Tab Group"] - Label for the tab group
|
|
* @returns {Promise<{
|
|
* success: boolean,
|
|
* group: {
|
|
* id: string,
|
|
* label: string,
|
|
* color: string,
|
|
* tabCount: number
|
|
* } | null,
|
|
* failedTabs: Array<{
|
|
* tab: Tab,
|
|
* reason: string
|
|
* }>,
|
|
* error?: string
|
|
* } | null>} Creation result with group details and success status, or null if no valid tabs
|
|
*/
|
|
static async createTabGroup({
|
|
tabs = [],
|
|
tokenToKey = null,
|
|
window: win,
|
|
label,
|
|
}) {
|
|
const verifiedTabObjects = this.#verifyAndCollectTabs(
|
|
tabs,
|
|
tokenToKey,
|
|
win
|
|
);
|
|
if (!verifiedTabObjects) {
|
|
return null;
|
|
}
|
|
|
|
const result = await lazy.tabManagementService.createTabGroup({
|
|
tabs: verifiedTabObjects,
|
|
window: win,
|
|
label,
|
|
});
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* Finds the original user prompt that led to the given assistant message
|
|
* by traversing the message chain backwards using parentMessageId
|
|
*
|
|
* @param {Array} messages - All conversation messages
|
|
* @param {object} assistantMessage - The assistant message with tool UI data
|
|
* @returns {string|null} The original user prompt text, or null if not found
|
|
*/
|
|
static findOriginalUserPrompt(messages, assistantMessage) {
|
|
// Use parentMessageId to trace back to the user message
|
|
let nextMessageId = assistantMessage.parentMessageId;
|
|
// To prevent potential infinite loops, we set a maximum depth for traversal
|
|
const maxDepth = 5;
|
|
let depth = 0;
|
|
|
|
// Follow the chain backwards to find the user message
|
|
while (nextMessageId && depth < maxDepth) {
|
|
const nextMessage = messages.find(m => m.id === nextMessageId);
|
|
|
|
if (!nextMessage) {
|
|
break;
|
|
}
|
|
|
|
if (
|
|
nextMessage.role === lazy.MESSAGE_ROLE.USER &&
|
|
nextMessage.content?.type === "text"
|
|
) {
|
|
return nextMessage.content.body;
|
|
}
|
|
|
|
// Continue up the chain
|
|
nextMessageId = nextMessage.parentMessageId;
|
|
depth++;
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
static handleUIDisplayTelemetry(toolUIData, telemetryData) {
|
|
if (toolUIData.uiType !== UI_TYPES.WEBSITE_CONFIRMATION) {
|
|
return;
|
|
}
|
|
|
|
const tabs = toolUIData.properties?.tabs ?? [];
|
|
const reason = this.#getConfirmationReason(tabs);
|
|
|
|
lazy.ToolUITelemetry.recordBrowserActionPrompt({
|
|
...telemetryData,
|
|
action_type: "close_tabs",
|
|
prompt_type: "safety_confirmation",
|
|
reason,
|
|
candidates: tabs.length,
|
|
preselected: 0, // we currently don't preselect any tabs
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Map of update type strings to their handler functions
|
|
*
|
|
* @private
|
|
*/
|
|
static #UPDATE_TYPE_HANDLERS = {
|
|
[UI_UPDATE_TYPES.CONFIRMATION_TAB_SELECTION]:
|
|
this.#handleConfirmationTabSelection.bind(this),
|
|
[UI_UPDATE_TYPES.CANCEL_TAB_SELECTION]:
|
|
this.#handleCancelTabSelection.bind(this),
|
|
[UI_UPDATE_TYPES.CONFIRM_TAB_GROUP_SELECTION]:
|
|
this.#handleConfirmTabGroupSelection.bind(this),
|
|
[UI_UPDATE_TYPES.UNDO_TAB_CLOSE]: this.#handleUndoTabClose.bind(this),
|
|
[UI_UPDATE_TYPES.UNDO_TAB_GROUP]: this.#handleUndoTabGroup.bind(this),
|
|
[UI_UPDATE_TYPES.RETRY_PROMPT]: this.#handleRetryPrompt.bind(this),
|
|
};
|
|
|
|
/**
|
|
* Checks if the conversation has an active website confirmation and auto-cancels it.
|
|
* This should be called when starting a new prompt to clean up pending confirmations.
|
|
*
|
|
* @param {object} conversation - The conversation object containing messages
|
|
* @param {ChromeWindow} window - The browser window object
|
|
* @param {string} mode - The mode of the AI Window (e.g., "sidebar", "popup")
|
|
* @returns {Promise<boolean>} True if a confirmation was cancelled, false otherwise
|
|
*/
|
|
static async autoCancelActiveConfirmation(conversation, window, mode) {
|
|
if (!conversation?.messages?.length) {
|
|
lazy.console.log("ToolUI: No conversation messages to check");
|
|
return false;
|
|
}
|
|
|
|
const lastAssistantTextMessage = this.#findLastAssistantTextMessage(
|
|
conversation.messages
|
|
);
|
|
|
|
// Early return if no confirmation to cancel (CONFIRMATION_UI_TYPES)
|
|
const uiType = lastAssistantTextMessage?.toolUIData?.uiType;
|
|
const isActiveConfirmation = CONFIRMATION_UI_TYPES.includes(uiType);
|
|
|
|
if (!isActiveConfirmation) {
|
|
lazy.console.log("ToolUI: No active confirmation to cancel");
|
|
return false;
|
|
}
|
|
|
|
lazy.console.log(`ToolUI: Found active ${uiType} to cancel`);
|
|
|
|
// Get the original user prompt from the existing toolUIData
|
|
// This was already added when the website confirmation was created
|
|
const originalUserPrompt =
|
|
lastAssistantTextMessage.toolUIData.properties?.originalUserPrompt;
|
|
|
|
const cancelData = {
|
|
messageId: lastAssistantTextMessage.id,
|
|
toolCallId: lastAssistantTextMessage.toolUIData.toolCallId,
|
|
updateType: UI_UPDATE_TYPES.CANCEL_TAB_SELECTION,
|
|
updateData: {
|
|
reason: "auto_cancel",
|
|
},
|
|
};
|
|
|
|
// Set pending retry state BEFORE the await to avoid race condition
|
|
// We'll clear it if cancellation fails - avoiding setting upstream methods to async for now
|
|
if (originalUserPrompt) {
|
|
conversation.pendingRetry = {
|
|
originalUserPrompt,
|
|
cancelledMessageId: lastAssistantTextMessage.id,
|
|
cancelledToolCallId: lastAssistantTextMessage.toolUIData.toolCallId,
|
|
cancelledUiType: lastAssistantTextMessage.toolUIData.uiType,
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
const cancelled = await this.handleUpdate(
|
|
cancelData,
|
|
conversation,
|
|
window,
|
|
mode
|
|
);
|
|
|
|
// Clear pendingRetry if cancellation failed
|
|
if (!cancelled && originalUserPrompt) {
|
|
conversation.pendingRetry = null;
|
|
}
|
|
|
|
return cancelled;
|
|
}
|
|
|
|
/**
|
|
* Inject retry toolUIData into a message if there's a pending retry state
|
|
*
|
|
* @param {object} msg - The message object to potentially modify
|
|
* @param {object} conversation - The conversation object that may have pendingRetry
|
|
* @returns {boolean} True if retry toolUIData was injected, false otherwise
|
|
*/
|
|
static injectRetryToolUIDataIfNeeded(msg, conversation) {
|
|
lazy.console.log("ToolUI: Checking if retry injection needed", {
|
|
hasPendingRetry: !!conversation?.pendingRetry,
|
|
msgRole: msg?.role,
|
|
contentType: msg?.content?.type,
|
|
});
|
|
|
|
if (
|
|
!conversation?.pendingRetry ||
|
|
msg?.role !== lazy.MESSAGE_ROLE.ASSISTANT ||
|
|
msg?.content?.type !== "text"
|
|
) {
|
|
return false;
|
|
}
|
|
|
|
lazy.console.log("ToolUI: Injecting retry component");
|
|
|
|
// Create the retry toolUIData with the original prompt
|
|
const retryToolUIData = {
|
|
uiType: UI_TYPES.RETRY_COMPONENT,
|
|
// Generate a unique synthetic ID for UI update handling (required by handleUpdate)
|
|
toolCallId: `retry-${crypto.randomUUID()}`,
|
|
properties: {
|
|
originalUserPrompt: conversation.pendingRetry.originalUserPrompt,
|
|
cancelledUiType: conversation.pendingRetry.cancelledUiType,
|
|
},
|
|
};
|
|
|
|
// Inject the retry toolUIData into the message itself
|
|
msg.toolUIData = retryToolUIData;
|
|
|
|
// Clear the pending retry state
|
|
conversation.pendingRetry = null;
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Handle updates to tool UI components from user interactions
|
|
*
|
|
* @param {object} data - The update data
|
|
* @param {string} data.messageId - ID of the message containing the tool UI
|
|
* @param {string} data.toolCallId - ID of the specific tool call
|
|
* @param {string} data.updateType - Type of update (confirmation, cancellation, etc.)
|
|
* @param {ToolUpdateData} data.updateData - Additional data for the update
|
|
* @param {object} conversation - The conversation object containing messages
|
|
* @param {ChromeWindow} window - The browser window object
|
|
* @param {string} [mode] - The mode of the AI Window (e.g., "sidebar", "popup") for context
|
|
* @returns {Promise<boolean>} True if update was successful, false otherwise
|
|
*/
|
|
static async handleUpdate(data, conversation, window, mode) {
|
|
const { messageId, toolCallId, updateType, updateData } = data ?? {};
|
|
|
|
if (!messageId || !toolCallId) {
|
|
return false;
|
|
}
|
|
|
|
// Find the message in the conversation
|
|
const message = conversation?.messages?.find(m => m.id === messageId);
|
|
|
|
// Check if the message exists and has matching toolUIData
|
|
if (message?.toolUIData?.toolCallId !== toolCallId) {
|
|
return false;
|
|
}
|
|
|
|
// Get the handler for this update type
|
|
const handler = this.#UPDATE_TYPE_HANDLERS[updateType];
|
|
if (typeof handler !== "function") {
|
|
lazy.console.error(`ToolUI: Unknown updateType "${updateType}"`);
|
|
return false;
|
|
}
|
|
|
|
// Call the handler with all context, let it destructure what it needs
|
|
return handler({
|
|
message,
|
|
toolCallId,
|
|
updateData,
|
|
conversation,
|
|
window,
|
|
originalData: data,
|
|
mode,
|
|
});
|
|
}
|
|
}
|