MoHan Zhang
852f3e5290
Combines the original D302389 work (decouple CallContext + prompt loading
from openAIEngine) with the follow-up layering refactor from Tyler's
2026-06-03/04 design feedback. Drops the LLM call-wrapper class entirely;
restructures ChatConversation/ChatMessage into a layered inheritance with
a generic Conversation/Message base in models/ and chat-specific
subclasses in ui/modules/.
Layering
models/Message (new — generic base)
^
ui/modules/ChatMessage (chat-only fields)
models/Conversation (new — generic base; owns
engine, parameters, run)
^
ui/modules/ChatConversation (chat orchestration + UI;
composes EventEmitter)
The base does not extend EventEmitter and has no `this.emit` calls — chat
overrides handleChunk / receiveResponse / retryMessage / set messages /
getMessagesInChatCompletionsFormat with chat side effects (no no-op hooks).
Class moves / renames
- LLM (models/LLM.sys.mjs) — deleted
- MessageAccumulator (models/MessageAccumulator.sys.mjs) — replaced by
Conversation (models/Conversation.sys.mjs)
- openAIEngine pulled out of models/Utils.sys.mjs into its own file
models/openAIEngine.sys.mjs; Utils re-exports it for test back-compat
with a comment pointing new callers at the new module
- Remote Settings access (RS_AI_WINDOW_COLLECTION, getRemoteClient,
modelPrefObserver, _remoteClient cache) lives in models/Utils.sys.mjs;
openAIEngine is now strictly the LiteLLM-endpoint transport
- MESSAGE_ROLE canonical definition lives in models/Conversation.sys.mjs;
ui/modules/ChatEnums.sys.mjs re-exports it (was previously duplicated)
- PromptLoader.loadCallContext + buildLLM → buildConversation +
buildEngineForFeature helper returning {engine, parameters} for chat
(which keeps a persistent ChatConversation across turns and refreshes
engine per-turn). buildEngineForFeature resolves baseURL + apiKey via
openAIEngine.resolveEndpointConfig(modelChoiceId).
- toWireFormat → getMessagesInChatCompletionsFormat on the base; chat
override applies URL token substitution and splices userContext
messages just before the last user message. The pre-refactor
getMessagesInOpenAiFormat name is gone — all callers renamed.
What moves to the base
Message: id, createdDate, ordinal, role, content, turnIndex,
parentMessageId, modelId, params, usage, toolCallId, toolName
Conversation: id, createdDate, updatedDate, feature, engine, parameters,
#messages, _minNextOrdinal, currentTurnIndex, addMessage(role, content,
turnIndex, opts), add*Message variants, setSystemMessage (idempotent),
retryMessage (generic truncate), compactChatCompletions,
getMessagesInChatCompletionsFormat, handleChunk (returns Boolean for
whether anything was extracted), receiveResponse (drain + flush, no
chat post-stream phases), run / runWithGenerator, systemPromptVersion
getter
What stays on ChatConversation
Chat-only fields: title, description, pageUrl, pageMeta, status,
securityProperties, urlToToken / tokenToUrl / #baseTokenCounts /
seenUrls, activeBranchTipMessageId, transientStarterUrl/Starters,
memoriesToggled
UI methods: renderState, addUIToolToCurrentMessage, updateToolUI; the
composed #emitter and on/off/emit forwarders
Chat orchestration:
- loadSystemPrompt(opts) — idempotent upsert. Calls loadPrompt directly
and writes body + RS-record version onto the system message content.
Called at init AND on model change.
- injectRealTimeContext(message, opts) — leaf op. Owns fetch + render +
in-place mutation of message.content.userContext.realTimeContext.
Replaces static getRealTimeInfo.
- injectMemoriesContext(message, prompt) — same shape for memories.
Replaces getMemoriesContext.
- generatePrompt — tight sequencer: loadSystemPrompt → addUserMessage →
emit → injectRealTimeContext → injectMemoriesContext →
securityProperties.commit.
Chat overrides of base methods (call super + chat extras)
- _createMessage — factory hook; returns ChatMessage with convId set
- addAssistantMessage / addToolCallMessage — auto-populate modelId from
this.engine?.model when opts doesn't already provide it
- retryMessage — captures ephemeral system messages first, then super.
Refreshes #updateActiveBranchTipMessageId() since base splices in
place (bypasses the setter)
- handleChunk — calls consumeStreamChunk with this.tokenToUrl directly
(does not super through the base), then applies plainText/tokens and
emits message-update + ChatStore.updateConversation
- receiveResponse — super.receiveResponse(stream, currentMessage); then
memory-id resolve, URL token strip, ChatStore.updateConversation,
message-complete emit
- set messages — super + #updateActiveBranchTipMessageId
- getMessagesInChatCompletionsFormat — accepts {applyUrlTokens=true};
filters out empty-body assistant placeholders and legacy ephemeral
SYSTEM-role realtime/memories messages, splices userContext as USER
messages just before the last user message, resolves inline @mention
URLs, applies URL→token substitution via replaceUrlsWithTokens
Versioning + LLMaJ telemetry
RS-record `version` is NOT a Conversation field — D304293 (Bug 2044484)
landed the convention of storing version on the system message's
content (message.content.version). loadSystemPrompt writes it there.
Base Conversation exposes systemPromptVersion reading back from the
system message. TelemetryUtils.runLLMaJTelemetry loses its llm param
and reads conversation.engine?.model + conversation.systemPromptVersion
directly.
Call-site updates
- Chat.fetchWithHistory({conversation, browsingContext, mode, signal}) —
drops the llm param. Uses conversation.compactChatCompletions() +
conversation.runWithGenerator(opts). Uses conversation.engine?.model
for telemetry tool calls.
- ai-window.mjs per-turn: const {engine, parameters} = await
buildEngineForFeature(MODEL_FEATURES.CHAT, opts); assigns onto the
persistent ChatConversation. Model-switch path calls
conversation.loadSystemPrompt({modelChoiceIdOverride}).
- TitleGeneration / ConversationSuggestions / Memories / MemoriesManager
— use buildConversation + conversation.setSystemMessage /
addUserMessage / run(opts). Memory pipeline functions clearMessages()
between steps because they reuse one Conversation across
generation/dedup/filter.
- MemoriesManager.ensureLLMForGeneration / ensureLLMForUsage →
ensureConversationForGeneration / ensureConversationForUsage.
Bug fixes uncovered during refactor
- Chat.sys.mjs addToolCallMessage call sites were passing 3 args
(content, currentTurn, toolRoleOpts) where the chat signature only
takes (content, toolOpts). toolRoleOpts was silently dropped and
currentTurn was being spread into the Message constructor as opts.
Fixed at all 5 sites; modelId is now auto-populated.
- ChatConversation constructor now seeds
#updateActiveBranchTipMessageId() after super() so DB-restored
conversations have it set before the first turn (base constructor
assigned messages directly into the private array, bypassing the
chat setter).
Tests
- test_LLM.js / test_MessageAccumulator.js — deleted
- test_PromptLoader_buildLLM.js → test_PromptLoader_buildConversation.js
- test_Message.js (new), test_Conversation.js (new) — cover base classes
- test_ChatConversation.js — removed direct tests for static
getRealTimeInfo / instance getMemoriesContext; added equivalents for
injectRealTimeContext / injectMemoriesContext.
- test_ChatSwitchModel.js — updateSystemPromptForModel → loadSystemPrompt
- test_Chat.js / browser_conversation_stream.js — replaced LLM
construction with a setupConversationForChat helper that assigns
conversation.engine
- test_MemoriesManager.js — ensureLLMForUsage stubs renamed
- test_TelemetryUtils.js — fake conversation exposes
getMessagesInChatCompletionsFormat
- browser_smartwindow_prompts.js, browser_smartwindow_retry_context.js —
stub names + arg-index assertions updated for the new instance methods
- Tests that previously stubbed openAIEngine.getRemoteClient switch to a
_setRemoteClientForTesting / _clearRemoteClientForTesting test seam in
Utils.sys.mjs (openAIEngine no longer owns getRemoteClient).
- Browser-test fixture ui/test/browser/head.js — uses the test seam for
the RS client and _setLoadPromptForTesting for the chat system prompt.
Test status
xpcshell: 51/51 pass.
mochitest-browser: 2354/2357 pass; the 3 failures are the
browser_smartwindow_sanitize.js suite-order flake (Bug 2006444 —
passes solo, fails in the full suite). No regressions.
End-to-end verification through ml_driver: ran user_journey happy path
on gpt-oss-120b. 14 scenarios, 45/45 turns successful, 0 inference errors,
mean LLM-judge score 3.29.
Differential Revision: https://phabricator.services.mozilla.com/D302389
79 lines
2.1 KiB
JavaScript
79 lines
2.1 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 https://mozilla.org/MPL/2.0/. */
|
|
|
|
/**
|
|
* Generic LLM message — a single turn in a Conversation. Holds the wire-shape
|
|
* fields any consumer needs (role, content, ordinal, turnIndex) plus
|
|
* lightweight metadata for replay/telemetry (id, createdDate, parentMessageId,
|
|
* modelId, params, usage). Tool-call linkage (toolCallId, toolName) lives here
|
|
* too so the base can serialize tool messages in the OpenAI chat-completions
|
|
* format.
|
|
*/
|
|
export class Message {
|
|
id;
|
|
createdDate;
|
|
ordinal;
|
|
role;
|
|
content;
|
|
turnIndex;
|
|
parentMessageId;
|
|
modelId;
|
|
params;
|
|
usage;
|
|
toolCallId;
|
|
toolName;
|
|
|
|
/**
|
|
* @param {object} param
|
|
* @param {number} param.ordinal
|
|
* @param {string} param.role
|
|
* @param {*} param.content
|
|
* @param {number} param.turnIndex
|
|
* @param {string} [param.id]
|
|
* @param {number} [param.createdDate]
|
|
* @param {?string} [param.parentMessageId]
|
|
* @param {?string} [param.modelId]
|
|
* @param {?object} [param.params]
|
|
* @param {?object} [param.usage]
|
|
* @param {?string} [param.toolCallId]
|
|
* @param {?string} [param.toolName]
|
|
*/
|
|
constructor({
|
|
ordinal,
|
|
role,
|
|
content,
|
|
turnIndex,
|
|
id = crypto.randomUUID(),
|
|
createdDate = Date.now(),
|
|
parentMessageId = null,
|
|
modelId = null,
|
|
params = null,
|
|
usage = null,
|
|
toolCallId = null,
|
|
toolName = null,
|
|
} = {}) {
|
|
this.id = id;
|
|
this.createdDate = createdDate;
|
|
this.ordinal = ordinal;
|
|
this.role = role;
|
|
this.content = content;
|
|
this.turnIndex = turnIndex;
|
|
this.parentMessageId = parentMessageId;
|
|
this.modelId = modelId;
|
|
this.params = params;
|
|
this.usage = usage;
|
|
this.toolCallId = toolCallId;
|
|
this.toolName = toolName;
|
|
}
|
|
|
|
/**
|
|
* Hook for token-stream side effects (e.g., URL/search/memory tokens parsed
|
|
* out of the model output). Base does nothing; chat overrides on ChatMessage.
|
|
*
|
|
* @param {object} _tokens
|
|
*/
|
|
// eslint-disable-next-line no-unused-vars
|
|
addTokens(_tokens) {}
|
|
}
|