fix: defer recall injection to pre-send rewrite

This commit is contained in:
Youzini-afk
2026-04-01 15:24:42 +08:00
parent bf3a24261b
commit 83f8f24f1c
4 changed files with 705 additions and 96 deletions

View File

@@ -121,6 +121,9 @@ export function registerCoreEventHooksController(runtime) {
if (eventTypes.MESSAGE_SENT) {
bind(eventTypes.MESSAGE_SENT, handlers.onMessageSent);
}
if (eventTypes.GENERATION_STARTED) {
bind(eventTypes.GENERATION_STARTED, handlers.onGenerationStarted);
}
const beforeCombineCleanup = runtime.registerBeforeCombinePrompts(
handlers.onBeforeCombinePrompts,
@@ -194,6 +197,40 @@ export function onMessageSentController(runtime, messageId) {
runtime.refreshPersistedRecallMessageUi?.();
}
export function onGenerationStartedController(
runtime,
type,
params = {},
dryRun = false,
) {
if (dryRun) return null;
if (params?.automatic_trigger || params?.quiet_prompt) return null;
const generationType = String(type || "normal").trim() || "normal";
if (generationType !== "normal") return null;
const pendingSendIntent = runtime.getPendingRecallSendIntent?.();
const pendingIntentText = runtime.isFreshRecallInputRecord?.(
pendingSendIntent,
)
? pendingSendIntent.text
: "";
const textareaText =
typeof runtime.getSendTextareaValue === "function"
? runtime.getSendTextareaValue()
: "";
const snapshotText =
runtime.normalizeRecallInputText?.(pendingIntentText || textareaText) || "";
if (!snapshotText) return null;
return runtime.freezeHostGenerationInputSnapshot(
snapshotText,
pendingIntentText
? "generation-started-send-intent"
: "generation-started-textarea",
);
}
export function onMessageDeletedController(
runtime,
chatLengthOrMessageId,
@@ -252,37 +289,70 @@ export async function onGenerationAfterCommandsController(
generationType,
recallOptions,
});
if (!recallContext.shouldRun) {
if (!recallContext.shouldRun && !recallContext.transaction) {
return;
}
const runtimeRecallOptions =
recallContext.recallOptions || recallOptions || {};
runtime.markGenerationRecallTransactionHookState(
const deliveryMode =
runtime.resolveGenerationRecallDeliveryMode?.(
recallContext.hookName,
recallContext.generationType,
runtimeRecallOptions,
) || "immediate";
let recallResult = runtime.getGenerationRecallTransactionResult?.(
recallContext.transaction,
recallContext.hookName,
"running",
);
const recallResult = await runtime.runRecall({
...runtimeRecallOptions,
recallKey: recallContext.recallKey,
hookName: recallContext.hookName,
signal: params?.signal,
});
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
runtime.getGenerationRecallHookStateFromResult(recallResult),
);
runtime.applyFinalRecallInjectionForGeneration({
if (recallContext.shouldRun) {
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
"running",
);
if (deliveryMode === "deferred") {
runtime.clearLiveRecallInjectionPromptForRewrite?.();
}
recallResult = await runtime.runRecall({
...runtimeRecallOptions,
deliveryMode,
recallKey: recallContext.recallKey,
hookName: recallContext.hookName,
signal: params?.signal,
});
runtime.storeGenerationRecallTransactionResult?.(
recallContext.transaction,
recallResult,
{
hookName: recallContext.hookName,
deliveryMode,
},
);
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
runtime.getGenerationRecallHookStateFromResult(recallResult),
);
}
if (deliveryMode === "deferred") {
return recallResult;
}
return runtime.applyFinalRecallInjectionForGeneration({
generationType: recallContext.generationType,
freshRecallResult: recallResult,
transaction: recallContext.transaction,
hookName: recallContext.hookName,
});
}
export async function onBeforeCombinePromptsController(runtime) {
export async function onBeforeCombinePromptsController(
runtime,
promptData = null,
) {
const frozenInputSnapshot =
runtime.consumeHostGenerationInputSnapshot?.() ||
runtime.getPendingHostGenerationInputSnapshot?.() ||
@@ -301,31 +371,58 @@ export async function onBeforeCombinePromptsController(runtime) {
generationType: "normal",
recallOptions,
});
if (!recallContext.shouldRun) {
if (!recallContext.shouldRun && !recallContext.transaction) {
return;
}
const runtimeRecallOptions =
recallContext.recallOptions || recallOptions || {};
runtime.markGenerationRecallTransactionHookState(
const deliveryMode =
runtime.resolveGenerationRecallDeliveryMode?.(
recallContext.hookName,
recallContext.generationType,
runtimeRecallOptions,
) || "deferred";
let recallResult = runtime.getGenerationRecallTransactionResult?.(
recallContext.transaction,
recallContext.hookName,
"running",
);
const recallResult = await runtime.runRecall({
...runtimeRecallOptions,
recallKey: recallContext.recallKey,
hookName: recallContext.hookName,
});
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
runtime.getGenerationRecallHookStateFromResult(recallResult),
);
runtime.applyFinalRecallInjectionForGeneration({
if (recallContext.shouldRun) {
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
"running",
);
if (deliveryMode === "deferred") {
runtime.clearLiveRecallInjectionPromptForRewrite?.();
}
recallResult = await runtime.runRecall({
...runtimeRecallOptions,
deliveryMode,
recallKey: recallContext.recallKey,
hookName: recallContext.hookName,
});
runtime.storeGenerationRecallTransactionResult?.(
recallContext.transaction,
recallResult,
{
hookName: recallContext.hookName,
deliveryMode,
},
);
runtime.markGenerationRecallTransactionHookState(
recallContext.transaction,
recallContext.hookName,
runtime.getGenerationRecallHookStateFromResult(recallResult),
);
}
return runtime.applyFinalRecallInjectionForGeneration({
generationType: recallContext.generationType,
freshRecallResult: recallResult,
transaction: recallContext.transaction,
promptData,
hookName: recallContext.hookName,
});
}

450
index.js
View File

@@ -45,6 +45,7 @@ import {
onChatChangedController,
onChatLoadedController,
onGenerationAfterCommandsController,
onGenerationStartedController,
onMessageDeletedController,
onMessageEditedController,
onMessageReceivedController,
@@ -1330,17 +1331,179 @@ function editMessageRecallRecord(messageIndex, nextInjectionText) {
return edited;
}
function rewriteRecallPayloadWithInjection(
promptData = null,
injectionText = "",
) {
const normalizedInjectionText = normalizeRecallInputText(injectionText);
if (!normalizedInjectionText) {
return {
applied: false,
path: "",
field: "",
reason: "empty-injection-text",
};
}
const finalMesSend = Array.isArray(promptData?.finalMesSend)
? promptData.finalMesSend
: null;
if (Array.isArray(finalMesSend) && finalMesSend.length > 0) {
for (let index = finalMesSend.length - 1; index >= 0; index--) {
const entry = finalMesSend[index];
if (!entry || typeof entry !== "object") continue;
if (entry.injected === true) continue;
const messageText = normalizeRecallInputText(
entry.message || entry.mes || entry.content || "",
);
if (!messageText) continue;
entry.extensionPrompts = Array.isArray(entry.extensionPrompts)
? entry.extensionPrompts
: [];
const alreadyPresent = entry.extensionPrompts.some((chunk) =>
String(chunk || "").includes(normalizedInjectionText),
);
if (!alreadyPresent) {
entry.extensionPrompts.push(`${normalizedInjectionText}\n`);
}
return {
applied: true,
path: "finalMesSend",
field: `finalMesSend[${index}].extensionPrompts`,
reason: alreadyPresent
? "rewrite-already-present"
: "finalMesSend-extensionPrompt-appended",
targetIndex: index,
};
}
return {
applied: false,
path: "finalMesSend",
field: "",
reason: "no-rewritable-finalMesSend-entry",
};
}
if (
typeof promptData?.combinedPrompt === "string" &&
promptData.combinedPrompt.trim()
) {
if (!promptData.combinedPrompt.includes(normalizedInjectionText)) {
promptData.combinedPrompt = `${normalizedInjectionText}\n\n${promptData.combinedPrompt}`;
}
return {
applied: true,
path: "combinedPrompt",
field: "combinedPrompt",
reason: "combinedPrompt-prefixed",
};
}
if (typeof promptData?.prompt === "string" && promptData.prompt.trim()) {
if (!promptData.prompt.includes(normalizedInjectionText)) {
promptData.prompt = `${normalizedInjectionText}\n\n${promptData.prompt}`;
}
return {
applied: true,
path: "prompt",
field: "prompt",
reason: "prompt-prefixed",
};
}
return {
applied: false,
path: "",
field: "",
reason: "prompt-payload-unavailable",
};
}
function readGenerationRecallTransactionFinalResolution(transaction) {
return transaction?.finalResolution || null;
}
function storeGenerationRecallTransactionFinalResolution(
transaction,
finalResolution = null,
) {
if (!transaction?.id) return transaction;
transaction.finalResolution = finalResolution ? { ...finalResolution } : null;
transaction.updatedAt = Date.now();
generationRecallTransactions.set(transaction.id, transaction);
return transaction;
}
function applyFinalRecallInjectionForGeneration({
generationType = "normal",
freshRecallResult = null,
transaction = null,
promptData = null,
hookName = "",
} = {}) {
const chat = getContext()?.chat;
if (!Array.isArray(chat)) {
applyModuleInjectionPrompt("", getSettings());
return { source: "none", targetUserMessageIndex: null, usedText: "" };
const existingFinalResolution =
readGenerationRecallTransactionFinalResolution(transaction);
if (existingFinalResolution) {
return existingFinalResolution;
}
let targetUserMessageIndex = resolveGenerationTargetUserMessageIndex(chat, {
const recallResult =
freshRecallResult ||
getGenerationRecallTransactionResult(transaction) ||
null;
const deliveryMode =
String(
recallResult?.deliveryMode ||
transaction?.lastDeliveryMode ||
resolveGenerationRecallDeliveryMode(
hookName,
generationType,
transaction?.frozenRecallOptions || {},
),
).trim() || "immediate";
const chat = getContext()?.chat;
let transport = {
applied: false,
source: "none",
mode: "none",
};
let targetUserMessageIndex = null;
let resolved = {
source: "none",
injectionText: "",
record: null,
};
const rewrite = {
applied: false,
path: "",
field: "",
reason: "no-recall-source",
};
let applicationMode = "none";
if (!Array.isArray(chat)) {
transport = applyModuleInjectionPrompt("", getSettings()) || transport;
const emptyResolution = {
source: "none",
isFallback: false,
targetUserMessageIndex: null,
usedText: "",
deliveryMode,
applicationMode: "none",
rewrite,
transport,
};
storeGenerationRecallTransactionFinalResolution(
transaction,
emptyResolution,
);
return emptyResolution;
}
targetUserMessageIndex = resolveGenerationTargetUserMessageIndex(chat, {
generationType,
});
if (
@@ -1354,15 +1517,71 @@ function applyFinalRecallInjectionForGeneration({
const persistedRecord = Number.isFinite(targetUserMessageIndex)
? readPersistedRecallFromUserMessage(chat, targetUserMessageIndex)
: null;
const resolved = resolveFinalRecallInjectionSource({
freshRecallResult,
resolved = resolveFinalRecallInjectionSource({
freshRecallResult: recallResult,
persistedRecord,
});
if (resolved.source === "persisted") {
applyModuleInjectionPrompt(resolved.injectionText || "", getSettings());
} else if (resolved.source === "none") {
applyModuleInjectionPrompt("", getSettings());
if (resolved.source === "fresh" && deliveryMode === "deferred") {
const rewriteResult = rewriteRecallPayloadWithInjection(
promptData,
resolved.injectionText || "",
);
Object.assign(rewrite, rewriteResult);
lastInjectionContent = resolved.injectionText || "";
if (rewriteResult.applied) {
applicationMode = "rewrite";
transport = clearLiveRecallInjectionPromptForRewrite() || {
applied: false,
source: "rewrite-cleared",
mode: "rewrite-cleared",
};
runtimeStatus = createUiStatus(
"召回已改写",
`本轮发送载荷已 rewrite · ${rewriteResult.path || rewriteResult.field || "payload"}`,
"success",
);
} else {
applicationMode = "fallback-injection";
transport =
applyModuleInjectionPrompt(
resolved.injectionText || "",
getSettings(),
) || transport;
runtimeStatus = createUiStatus(
"召回回退",
`rewrite 未命中,已回退注入 · ${rewriteResult.reason}`,
"warning",
);
}
} else if (resolved.source === "fresh") {
applicationMode = "injection";
transport =
applyModuleInjectionPrompt(resolved.injectionText || "", getSettings()) ||
transport;
lastInjectionContent = resolved.injectionText || "";
rewrite.reason = "immediate-injection";
runtimeStatus = createUiStatus(
"召回已注入",
"本轮已使用最新召回结果",
"success",
);
} else if (resolved.source === "persisted") {
applicationMode = "persisted-injection";
transport =
applyModuleInjectionPrompt(resolved.injectionText || "", getSettings()) ||
transport;
lastInjectionContent = resolved.injectionText || "";
rewrite.reason = "persisted-record-fallback";
runtimeStatus = createUiStatus(
"召回回退",
"已使用消息楼层持久化注入",
"info",
);
} else {
transport = applyModuleInjectionPrompt("", getSettings()) || transport;
lastInjectionContent = "";
runtimeStatus = createUiStatus("待命", "当前无有效注入内容", "idle");
}
if (
@@ -1373,32 +1592,87 @@ function applyFinalRecallInjectionForGeneration({
triggerChatMetadataSave(getContext(), { immediate: false });
}
if (resolved.source === "fresh") {
runtimeStatus = createUiStatus(
"召回已注入",
"本轮已使用最新召回结果",
"success",
);
} else if (resolved.source === "persisted") {
lastInjectionContent = resolved.injectionText || "";
runtimeStatus = createUiStatus(
"召回回退",
"已使用消息楼层持久化注入",
"info",
);
} else {
lastInjectionContent = "";
runtimeStatus = createUiStatus("待命", "当前无有效注入内容", "idle");
}
recordInjectionSnapshot("recall", {
taskType: "recall",
source:
String(
recallResult?.source ||
transaction?.frozenRecallOptions?.lockedSource ||
transaction?.frozenRecallOptions?.overrideSource ||
"",
).trim() || "unknown",
sourceLabel:
String(
recallResult?.sourceLabel ||
transaction?.frozenRecallOptions?.lockedSourceLabel ||
transaction?.frozenRecallOptions?.overrideSourceLabel ||
"",
).trim() || "未知",
reason:
String(
recallResult?.reason ||
transaction?.frozenRecallOptions?.lockedReason ||
transaction?.frozenRecallOptions?.overrideReason ||
"",
).trim() || "final-application",
sourceCandidates: Array.isArray(recallResult?.sourceCandidates)
? recallResult.sourceCandidates.map((candidate) => ({ ...candidate }))
: Array.isArray(transaction?.frozenRecallOptions?.sourceCandidates)
? transaction.frozenRecallOptions.sourceCandidates.map((candidate) => ({
...candidate,
}))
: [],
hookName: String(hookName || recallResult?.hookName || "").trim(),
selectedNodeIds: recallResult?.selectedNodeIds || [],
retrievalMeta: recallResult?.retrievalMeta || {},
llmMeta: recallResult?.llmMeta || {},
stats: recallResult?.stats || {},
injectionText: resolved.injectionText || "",
deliveryMode,
applicationMode,
transport,
rewrite,
targetUserMessageIndex,
sourceKind: resolved.source,
});
refreshPanelLiveState();
schedulePersistedRecallMessageUiRefresh();
return {
const finalResolution = {
source: resolved.source,
isFallback: resolved.source === "persisted",
isFallback:
resolved.source === "persisted" ||
applicationMode === "fallback-injection",
targetUserMessageIndex,
usedText: resolved.injectionText || "",
deliveryMode,
applicationMode,
rewrite,
transport,
};
storeGenerationRecallTransactionFinalResolution(transaction, finalResolution);
return finalResolution;
}
function clearLiveRecallInjectionPromptForRewrite() {
try {
return (
applyModuleInjectionPrompt("", getSettings()) || {
applied: false,
source: "rewrite-clear",
mode: "rewrite-clear",
}
);
} catch (error) {
console.warn("[ST-BME] 清理 rewrite 前旧注入失败:", error);
return {
applied: false,
source: "rewrite-clear-error",
mode: "rewrite-clear-error",
error: error instanceof Error ? error.message : String(error || ""),
};
}
}
function clearPersistedRecallMessageUiObserver() {
@@ -5187,6 +5461,28 @@ function normalizeGenerationRecallTransactionType(generationType = "normal") {
return normalized === "normal" ? "normal" : "history";
}
function resolveGenerationRecallDeliveryMode(
hookName,
generationType = "normal",
recallOptions = {},
) {
if (recallOptions?.forceImmediateDelivery === true) {
return "immediate";
}
const normalizedType = normalizeGenerationRecallTransactionType(
recallOptions?.generationType || generationType,
);
if (normalizedType !== "normal") {
return "immediate";
}
return hookName === "GENERATION_AFTER_COMMANDS" ||
hookName === "GENERATE_BEFORE_COMBINE_PROMPTS"
? "deferred"
: "immediate";
}
function freezeGenerationRecallOptionsForTransaction(
chat,
generationType = "normal",
@@ -5446,6 +5742,29 @@ function markGenerationRecallTransactionHookState(
return transaction;
}
function getGenerationRecallTransactionResult(transaction) {
return transaction?.lastRecallResult || null;
}
function storeGenerationRecallTransactionResult(
transaction,
recallResult = null,
meta = {},
) {
if (!transaction?.id) return transaction;
transaction.lastRecallResult = recallResult ? { ...recallResult } : null;
transaction.lastRecallMeta =
meta && typeof meta === "object" ? { ...meta } : {};
transaction.lastDeliveryMode =
String(meta?.deliveryMode || recallResult?.deliveryMode || "").trim() ||
transaction.lastDeliveryMode ||
"";
transaction.finalResolution = null;
transaction.updatedAt = Date.now();
generationRecallTransactions.set(transaction.id, transaction);
return transaction;
}
function clearGenerationRecallTransactionsForChat(
chatId = getCurrentChatId(),
{ clearAll = false } = {},
@@ -6944,17 +7263,14 @@ function onMessageSwiped(messageId, meta = null) {
);
}
async function onGenerationAfterCommands(type, params = {}, dryRun = false) {
return await onGenerationAfterCommandsController(
function onGenerationStarted(type, params = {}, dryRun = false) {
return onGenerationStartedController(
{
applyFinalRecallInjectionForGeneration,
buildGenerationAfterCommandsRecallInput,
consumeHostGenerationInputSnapshot,
createGenerationRecallContext,
getContext,
getGenerationRecallHookStateFromResult,
markGenerationRecallTransactionHookState,
runRecall,
freezeHostGenerationInputSnapshot,
getPendingRecallSendIntent: () => pendingRecallSendIntent,
getSendTextareaValue,
isFreshRecallInputRecord,
normalizeRecallInputText,
},
type,
params,
@@ -6962,18 +7278,47 @@ async function onGenerationAfterCommands(type, params = {}, dryRun = false) {
);
}
async function onBeforeCombinePrompts() {
return await onBeforeCombinePromptsController({
applyFinalRecallInjectionForGeneration,
buildHistoryGenerationRecallInput,
buildNormalGenerationRecallInput,
consumeHostGenerationInputSnapshot,
createGenerationRecallContext,
getContext,
getGenerationRecallHookStateFromResult,
markGenerationRecallTransactionHookState,
runRecall,
});
async function onGenerationAfterCommands(type, params = {}, dryRun = false) {
return await onGenerationAfterCommandsController(
{
applyFinalRecallInjectionForGeneration,
buildGenerationAfterCommandsRecallInput,
clearLiveRecallInjectionPromptForRewrite,
consumeHostGenerationInputSnapshot,
createGenerationRecallContext,
getContext,
getGenerationRecallHookStateFromResult,
getGenerationRecallTransactionResult,
markGenerationRecallTransactionHookState,
resolveGenerationRecallDeliveryMode,
runRecall,
storeGenerationRecallTransactionResult,
},
type,
params,
dryRun,
);
}
async function onBeforeCombinePrompts(promptData = null) {
return await onBeforeCombinePromptsController(
{
applyFinalRecallInjectionForGeneration,
buildHistoryGenerationRecallInput,
buildNormalGenerationRecallInput,
clearLiveRecallInjectionPromptForRewrite,
consumeHostGenerationInputSnapshot,
createGenerationRecallContext,
getContext,
getGenerationRecallHookStateFromResult,
getGenerationRecallTransactionResult,
markGenerationRecallTransactionHookState,
resolveGenerationRecallDeliveryMode,
runRecall,
storeGenerationRecallTransactionResult,
},
promptData,
);
}
function onMessageReceived() {
@@ -7290,6 +7635,7 @@ async function onReembedDirect() {
onChatChanged,
onChatLoaded,
onGenerationAfterCommands,
onGenerationStarted,
onMessageDeleted,
onMessageEdited,
onMessageReceived,

View File

@@ -169,6 +169,8 @@ export function applyRecallInjectionController(
reason: settings.recallEnableLLM ? "未提供 LLM 状态" : "LLM 精排已关闭",
candidatePool: 0,
};
const deliveryMode =
String(recallInput?.deliveryMode || "immediate").trim() || "immediate";
if (injectionText) {
const tokens = runtime.estimateTokens(injectionText);
@@ -183,10 +185,16 @@ export function applyRecallInjectionController(
});
}
const injectionTransport = runtime.applyModuleInjectionPrompt(
injectionText,
settings,
);
let injectionTransport = {
applied: false,
source: "deferred",
mode: "deferred",
};
if (deliveryMode === "immediate") {
injectionTransport =
runtime.applyModuleInjectionPrompt(injectionText, settings) ||
injectionTransport;
}
runtime.recordInjectionSnapshot("recall", {
taskType: "recall",
source: recallInput.source,
@@ -202,6 +210,18 @@ export function applyRecallInjectionController(
llmMeta,
stats: result.stats || {},
injectionText,
deliveryMode,
applicationMode:
deliveryMode === "immediate" ? "injection" : "pending-rewrite",
rewrite: {
applied: false,
path: "",
field: "",
reason:
deliveryMode === "immediate"
? "immediate-injection"
: "awaiting-generation-payload-rewrite",
},
transport: injectionTransport,
});
@@ -223,6 +243,7 @@ export function applyRecallInjectionController(
[
hookLabel,
recallInput.sourceLabel,
deliveryMode === "immediate" ? "即时注入" : "等待本轮 rewrite",
`ctx ${recentMessages.length}`,
`vector ${retrievalMeta.vectorHits ?? 0}`,
retrievalMeta.vectorMergedHits
@@ -256,7 +277,13 @@ export function applyRecallInjectionController(
}
}
return { injectionText, retrievalMeta, llmMeta };
return {
injectionText,
retrievalMeta,
llmMeta,
transport: injectionTransport,
deliveryMode,
};
}
export async function runRecallController(runtime, options = {}) {
@@ -366,6 +393,8 @@ export async function runRecallController(runtime, options = {}) {
}
recallInput.hookName = options.hookName || "";
recallInput.deliveryMode =
String(options.deliveryMode || "immediate").trim() || "immediate";
runtime.console.log("[ST-BME] 开始召回", {
source: recallInput.source,
@@ -425,6 +454,24 @@ export async function runRecallController(runtime, options = {}) {
reason: "召回完成",
selectedNodeIds: result.selectedNodeIds || [],
injectionText: applied?.injectionText || "",
retrievalMeta: applied?.retrievalMeta || {},
llmMeta: applied?.llmMeta || {},
transport: applied?.transport || {
applied: false,
source: "none",
mode: "none",
},
deliveryMode:
applied?.deliveryMode ||
String(recallInput?.deliveryMode || "immediate").trim() ||
"immediate",
source: recallInput?.source || "",
sourceLabel: recallInput?.sourceLabel || "",
hookName: recallInput?.hookName || "",
sourceCandidates: Array.isArray(recallInput?.sourceCandidates)
? recallInput.sourceCandidates.map((candidate) => ({ ...candidate }))
: [],
stats: result?.stats || {},
});
} catch (e) {
if (runtime.isAbortError(e)) {

View File

@@ -8,6 +8,7 @@ import { pruneProcessedMessageHashesFromFloor } from "../chat-history.js";
import {
onBeforeCombinePromptsController,
onGenerationAfterCommandsController,
onGenerationStartedController,
registerCoreEventHooksController,
} from "../event-binding.js";
import { onRerollController } from "../extraction-controller.js";
@@ -244,7 +245,8 @@ function createBatchStageHarness() {
});
}
function createGenerationRecallHarness() {
function createGenerationRecallHarness(options = {}) {
const { realApplyFinal = false } = options;
return fs.readFile(indexPath, "utf8").then((source) => {
const start = source.indexOf("const RECALL_INPUT_RECORD_TTL_MS = 60000;");
const end = source.indexOf("function onMessageReceived() {");
@@ -273,6 +275,33 @@ function createGenerationRecallHarness() {
},
result: null,
currentGraph: {},
_panelModule: null,
defaultSettings: {},
extension_settings: { [MODULE_NAME]: {} },
extension_prompt_types: {
NONE: 0,
BEFORE_PROMPT: 1,
IN_PROMPT: 2,
IN_CHAT: 3,
},
extension_prompt_roles: {
SYSTEM: 0,
USER: 1,
ASSISTANT: 2,
},
clampInt: (value, fallback = 0, min = 0, max = 9999) => {
const numeric = Number(value);
if (!Number.isFinite(numeric)) return fallback;
return Math.min(max, Math.max(min, Math.trunc(numeric)));
},
getHostAdapter: () => null,
migrateLegacyTaskProfiles: (settings = {}) => ({
taskProfilesVersion: settings?.taskProfilesVersion || 0,
taskProfiles: settings?.taskProfiles || {},
}),
refreshPanelLiveStateController: () => {
context.refreshPanelCalls += 1;
},
isRecalling: false,
getCurrentChatId: () => "chat-main",
normalizeRecallInputText: (text = "") => String(text || "").trim(),
@@ -298,6 +327,9 @@ function createGenerationRecallHarness() {
chat: [],
runRecallCalls: [],
applyFinalCalls: [],
moduleInjectionCalls: [],
recordedInjectionSnapshots: [],
refreshPanelCalls: 0,
createRecallInputRecord,
createRecallRunResult,
hashRecallInput,
@@ -318,6 +350,7 @@ function createGenerationRecallHarness() {
GRAPH_PERSISTENCE_META_KEY,
onBeforeCombinePromptsController,
onGenerationAfterCommandsController,
onGenerationStartedController,
readPersistedRecallFromUserMessage: () => null,
resolveFinalRecallInjectionSource: ({
freshRecallResult = null,
@@ -327,10 +360,24 @@ function createGenerationRecallHarness() {
record: null,
}),
bumpPersistedRecallGenerationCount: () => null,
applyModuleInjectionPrompt: () => ({}),
applyModuleInjectionPrompt: (text = "") => {
const normalizedText = String(text || "");
context.moduleInjectionCalls.push(normalizedText);
return {
applied: Boolean(normalizedText.trim()),
source: normalizedText.trim() ? "module-injection" : "rewrite-clear",
mode: normalizedText.trim() ? "module-injection" : "rewrite-clear",
};
},
getSettings: () => ({}),
triggerChatMetadataSave: () => "debounced",
refreshPanelLiveState: () => {},
refreshPanelLiveState: () => {
context.refreshPanelCalls += 1;
},
recordInjectionSnapshot: (_kind, snapshot = {}) => {
context.recordedInjectionSnapshots.push({ ...snapshot });
},
schedulePersistedRecallMessageUiRefresh: () => {},
resolveGenerationTargetUserMessageIndex: (
chat = [],
{ generationType } = {},
@@ -346,7 +393,7 @@ function createGenerationRecallHarness() {
};
vm.createContext(context);
vm.runInContext(
`${snippet}\nresult = { hashRecallInput, buildPreGenerationRecallKey, buildGenerationAfterCommandsRecallInput, cleanupGenerationRecallTransactions, buildGenerationRecallTransactionId, beginGenerationRecallTransaction, markGenerationRecallTransactionHookState, shouldRunRecallForTransaction, createGenerationRecallContext, onGenerationAfterCommands, onBeforeCombinePrompts, generationRecallTransactions, freezeHostGenerationInputSnapshot, consumeHostGenerationInputSnapshot, getPendingHostGenerationInputSnapshot, recordRecallSendIntent, recordRecallSentUserMessage, getPendingRecallSendIntent: () => pendingRecallSendIntent, getLastRecallSentUserMessage: () => lastRecallSentUserMessage };`,
`${snippet}\nresult = { hashRecallInput, buildPreGenerationRecallKey, buildGenerationAfterCommandsRecallInput, cleanupGenerationRecallTransactions, buildGenerationRecallTransactionId, beginGenerationRecallTransaction, markGenerationRecallTransactionHookState, shouldRunRecallForTransaction, createGenerationRecallContext, onGenerationStarted, onGenerationAfterCommands, onBeforeCombinePrompts, applyFinalRecallInjectionForGeneration, generationRecallTransactions, freezeHostGenerationInputSnapshot, consumeHostGenerationInputSnapshot, getPendingHostGenerationInputSnapshot, recordRecallSendIntent, recordRecallSentUserMessage, getPendingRecallSendIntent: () => pendingRecallSendIntent, getLastRecallSentUserMessage: () => lastRecallSentUserMessage };`,
context,
{ filename: indexPath },
);
@@ -377,8 +424,13 @@ function createGenerationRecallHarness() {
configurable: true,
},
});
const originalApplyFinalRecallInjectionForGeneration =
context.result.applyFinalRecallInjectionForGeneration;
context.applyFinalRecallInjectionForGeneration = (payload = {}) => {
context.applyFinalCalls.push({ ...payload });
if (realApplyFinal) {
return originalApplyFinalRecallInjectionForGeneration(payload);
}
return {
source: "fresh",
targetUserMessageIndex: null,
@@ -386,14 +438,37 @@ function createGenerationRecallHarness() {
};
context.runRecall = async (options = {}) => {
context.runRecallCalls.push({ ...options });
const overrideUserMessage = String(
options.overrideUserMessage || options.userMessage || "",
);
return {
status: "completed",
didRecall: true,
ok: true,
injectionText: `注入:${overrideUserMessage}`,
deliveryMode: String(options.deliveryMode || "immediate"),
source: options.overrideSource,
sourceLabel: options.overrideSourceLabel,
reason: options.overrideReason,
sourceCandidates: options.sourceCandidates,
sourceCandidates: Array.isArray(options.sourceCandidates)
? options.sourceCandidates.map((candidate) => ({ ...candidate }))
: [],
selectedNodeIds: ["node-test-1"],
retrievalMeta: {
vectorHits: 1,
vectorMergedHits: 0,
diffusionHits: 0,
candidatePoolAfterDpp: 1,
},
llmMeta: {
status: "disabled",
reason: "test-disabled",
candidatePool: 0,
},
stats: {
coreCount: 1,
recallCount: 1,
},
};
};
return context;
@@ -2670,6 +2745,7 @@ async function testRegisterCoreEventHooksIsIdempotent() {
CHAT_CHANGED: "chat-changed",
CHAT_LOADED: "chat-loaded",
MESSAGE_SENT: "message-sent",
GENERATION_STARTED: "generation-started",
MESSAGE_RECEIVED: "message-received",
MESSAGE_DELETED: "message-deleted",
MESSAGE_EDITED: "message-edited",
@@ -2680,6 +2756,7 @@ async function testRegisterCoreEventHooksIsIdempotent() {
onChatChanged() {},
onChatLoaded() {},
onMessageSent() {},
onGenerationStarted() {},
onGenerationAfterCommands() {},
onBeforeCombinePrompts() {},
onMessageReceived() {},
@@ -2709,7 +2786,7 @@ async function testRegisterCoreEventHooksIsIdempotent() {
registerCoreEventHooksController(runtime);
registerCoreEventHooksController(runtime);
assert.equal(eventRegistrations.length, 8);
assert.equal(eventRegistrations.length, 9);
assert.equal(makeFirstRegistrations.length, 2);
assert.equal(bindingState.registered, true);
}
@@ -2751,6 +2828,47 @@ async function testGenerationRecallAppliesFinalInjectionOncePerTransaction() {
assert.equal(harness.applyFinalCalls[0].generationType, "normal");
}
async function testGenerationRecallDeferredRewriteMutatesFinalMesSendPayload() {
const harness = await createGenerationRecallHarness({ realApplyFinal: true });
harness.chat = [{ is_user: false, mes: "assistant-tail" }];
harness.__sendTextareaValue = "发送前真实输入";
await harness.result.onGenerationStarted("normal", {}, false);
harness.__sendTextareaValue = "";
await harness.result.onGenerationAfterCommands("normal", {}, false);
const promptData = {
finalMesSend: [
{
injected: false,
message: "发送前真实输入",
extensionPrompts: [],
},
],
};
const resolution = await harness.result.onBeforeCombinePrompts(promptData);
assert.equal(harness.runRecallCalls.length, 1);
assert.equal(harness.applyFinalCalls.length, 1);
assert.equal(resolution.applicationMode, "rewrite");
assert.equal(resolution.deliveryMode, "deferred");
assert.equal(resolution.rewrite.applied, true);
assert.equal(resolution.rewrite.path, "finalMesSend");
assert.match(
promptData.finalMesSend[0].extensionPrompts.join("\n"),
/注入:发送前真实输入/,
);
assert.equal(
harness.moduleInjectionCalls.every((text) => text === ""),
true,
);
assert.equal(
harness.recordedInjectionSnapshots.at(-1)?.applicationMode,
"rewrite",
);
}
async function testGenerationRecallSendIntentBeatsChatTailAndStaysObservable() {
const harness = await createGenerationRecallHarness();
harness.chat = [{ is_user: true, mes: "旧的 chat tail" }];
@@ -3772,6 +3890,7 @@ await testGenerationRecallSentMessageClearsStaleTransactionForSameKey();
await testRegisterCoreEventHooksIsIdempotent();
await testRemoveNodeHandlesCyclicChildGraph();
await testGenerationRecallAppliesFinalInjectionOncePerTransaction();
await testGenerationRecallDeferredRewriteMutatesFinalMesSendPayload();
await testPersistentRecallDataLayerLifecycleAndCompatibility();
await testPersistentRecallSourceResolutionAndTargetRouting();
await testRecallCardMountsOnStandardUserMessageDom();