diff --git a/index.js b/index.js index 578a04f..faf7821 100644 --- a/index.js +++ b/index.js @@ -162,6 +162,7 @@ import { } from "./runtime/reroll-transaction-boundary.js"; import { createRecallInputState } from "./runtime/recall-input-state.js"; import { createRerollRecallInput } from "./runtime/reroll-recall-input.js"; +import { createGenerationContextTracker } from "./runtime/generation-context.js"; import { createGenerationRecallTransactions } from "./runtime/generation-recall-transactions.js"; import { createFinalRecallInjection } from "./runtime/final-recall-injection.js"; import { createAutoExtractionDefer } from "./runtime/auto-extraction-defer.js"; @@ -1739,6 +1740,10 @@ const generationRecallTransactionRuntime = createGenerationRecallTransactions({ }); const generationRecallTransactions = generationRecallTransactionRuntime.generationRecallTransactions; +const generationContextTracker = createGenerationContextTracker({ + getCurrentChatId, + ttlMs: GENERATION_RECALL_TRANSACTION_TTL_MS, +}); const finalRecallInjectionRuntime = createFinalRecallInjection({ applyModuleInjectionPrompt: (...args) => applyModuleInjectionPrompt(...args), areRecallNodeIdListsEqual: (...args) => areRecallNodeIdListsEqual(...args), @@ -15860,6 +15865,7 @@ async function runRecall(options = {}) { function onChatChanged() { isHostGenerationRunning = false; lastHostGenerationEndedAt = 0; + generationContextTracker.clear("chat-changed"); const { target, lightweightHostMode, adapter } = syncBmeHostRuntimeFlags(getContext()); updateGraphPersistenceState({ hostProfile: adapter.hostProfile, @@ -16042,6 +16048,7 @@ function onMessageUpdated(messageId, meta = null) { } async function onMessageSwiped(messageId, meta = null) { + generationContextTracker.noteSwipe(messageId, meta); const result = await onMessageSwipedController( { invalidateRecallAfterHistoryMutation, @@ -16155,6 +16162,10 @@ function onGenerationBeforeApiRequest(payload = {}) { function onGenerationStarted(type, params = {}, dryRun = false) { const generationType = String(type || "normal").trim() || "normal"; + generationContextTracker.begin(generationType, params, { + dryRun, + phase: "GENERATION_STARTED", + }); if ( !dryRun && !params?.automatic_trigger && @@ -16218,9 +16229,14 @@ function onGenerationEnded(_chatLength = null) { if (typeof scheduleMessageHideApply === "function") { scheduleMessageHideApply("generation-ended", 180); } + generationContextTracker.clear("generation-ended"); } async function onGenerationAfterCommands(type, params = {}, dryRun = false) { + generationContextTracker.update(type, params, { + dryRun, + phase: "GENERATION_AFTER_COMMANDS", + }); return await onGenerationAfterCommandsController( { applyFinalRecallInjectionForGeneration, diff --git a/package.json b/package.json index 41f9414..1b99485 100644 --- a/package.json +++ b/package.json @@ -5,6 +5,7 @@ "test:p0": "node tests/p0-regressions.mjs", "test:triviumdb-poc": "node tests/triviumdb-poc.mjs", "test:runtime-history": "node tests/runtime-history.mjs", + "test:generation-context": "node tests/generation-context.mjs", "test:graph-persistence": "node tests/graph-persistence.mjs", "test:index-slicing-ratchet": "node tests/index-slicing-ratchet.mjs", "test:runtime-deps": "node tests/runtime-deps-completeness.mjs", diff --git a/runtime/generation-context.js b/runtime/generation-context.js new file mode 100644 index 0000000..d55eaa3 --- /dev/null +++ b/runtime/generation-context.js @@ -0,0 +1,156 @@ +const DEFAULT_GENERATION_CONTEXT_TTL_MS = 60000; + +export function normalizeGenerationType(type = "normal") { + const normalized = String(type || "normal").trim(); + return normalized || "normal"; +} + +export function classifyGenerationKind(type = "normal", params = {}) { + const generationType = normalizeGenerationType(type); + if (params?.automatic_trigger || params?.quiet_prompt) { + return "skip"; + } + if (generationType === "quiet" || generationType === "impersonate") { + return "skip"; + } + if ( + generationType === "swipe" || + generationType === "regenerate" || + generationType === "continue" + ) { + return "no-new-user"; + } + return "fresh"; +} + +function clonePlain(value, fallback = null) { + if (!value || typeof value !== "object") return fallback; + try { + return structuredClone(value); + } catch (_error) { + try { + return JSON.parse(JSON.stringify(value)); + } catch (_jsonError) { + return fallback; + } + } +} + +export function createGenerationContextTracker(deps = {}) { + let current = null; + let pendingSwipe = null; + let sequence = 0; + + const now = () => + typeof deps.now === "function" ? Number(deps.now()) || Date.now() : Date.now(); + const ttlMs = () => + Number.isFinite(Number(deps.ttlMs)) + ? Math.max(1, Number(deps.ttlMs)) + : DEFAULT_GENERATION_CONTEXT_TTL_MS; + const getChatId = () => String(deps.getCurrentChatId?.() || "").trim(); + + function noteSwipe(messageId = null, meta = null) { + const parsed = Number(messageId); + pendingSwipe = { + assistantFloor: Number.isFinite(parsed) ? Math.floor(parsed) : null, + meta: clonePlain(meta, null), + chatId: getChatId(), + at: now(), + }; + return pendingSwipe; + } + + function begin(type = "normal", params = {}, { dryRun = false, phase = "" } = {}) { + if (dryRun) return null; + const at = now(); + const generationType = normalizeGenerationType(type); + const kind = classifyGenerationKind(generationType, params); + const context = { + id: `${at}:${++sequence}`, + type: generationType, + kind, + chatId: getChatId(), + params: clonePlain(params, {}), + dryRun: false, + startedAt: at, + updatedAt: at, + phase: String(phase || ""), + swipedAssistantFloor: + generationType === "swipe" && Number.isFinite(pendingSwipe?.assistantFloor) + ? pendingSwipe.assistantFloor + : null, + swipeMeta: generationType === "swipe" ? clonePlain(pendingSwipe?.meta, null) : null, + expectedMutation: "", + expectedMutationAt: 0, + }; + pendingSwipe = null; + current = context; + return { ...context }; + } + + function update(type = "normal", params = {}, { dryRun = false, phase = "" } = {}) { + if (dryRun) return null; + const at = now(); + const generationType = normalizeGenerationType(type); + const kind = classifyGenerationKind(generationType, params); + if (!current || current.type !== generationType || current.chatId !== getChatId()) { + return begin(generationType, params, { dryRun, phase }); + } + current = { + ...current, + type: generationType, + kind, + params: clonePlain(params, current.params || {}), + updatedAt: at, + afterCommandsAt: + String(phase || "") === "GENERATION_AFTER_COMMANDS" + ? at + : current.afterCommandsAt || 0, + phase: String(phase || current.phase || ""), + }; + return { ...current }; + } + + function get({ allowStale = false } = {}) { + if (!current) return null; + const age = now() - Number(current.updatedAt || current.startedAt || 0); + if (!allowStale && age > ttlMs()) { + current = null; + return null; + } + const activeChatId = getChatId(); + if (current.chatId && activeChatId && current.chatId !== activeChatId) { + current = null; + return null; + } + return { ...current }; + } + + function markExpectedMutation(kind = "", payload = {}) { + if (!current) return null; + current = { + ...current, + expectedMutation: String(kind || ""), + expectedMutationAt: now(), + expectedMutationPayload: clonePlain(payload, {}), + updatedAt: now(), + }; + return { ...current }; + } + + function clear(reason = "") { + const previous = current; + current = null; + pendingSwipe = null; + return previous ? { ...previous, clearReason: String(reason || "") } : null; + } + + return { + begin, + update, + get, + clear, + noteSwipe, + markExpectedMutation, + }; +} diff --git a/tests/generation-context.mjs b/tests/generation-context.mjs new file mode 100644 index 0000000..7f1d3ad --- /dev/null +++ b/tests/generation-context.mjs @@ -0,0 +1,102 @@ +import assert from "node:assert/strict"; +import { + classifyGenerationKind, + createGenerationContextTracker, +} from "../runtime/generation-context.js"; + +assert.equal(classifyGenerationKind("normal"), "fresh"); +assert.equal(classifyGenerationKind("swipe"), "no-new-user"); +assert.equal(classifyGenerationKind("regenerate"), "no-new-user"); +assert.equal(classifyGenerationKind("continue"), "no-new-user"); +assert.equal(classifyGenerationKind("quiet"), "skip"); +assert.equal(classifyGenerationKind("impersonate"), "skip"); +assert.equal(classifyGenerationKind("normal", { automatic_trigger: true }), "skip"); +assert.equal(classifyGenerationKind("normal", { quiet_prompt: true }), "skip"); + +{ + let chatId = "chat-swipe"; + let now = 1000; + const tracker = createGenerationContextTracker({ + getCurrentChatId: () => chatId, + now: () => now, + }); + + tracker.noteSwipe(7); + const context = tracker.begin("swipe"); + + assert.equal(context.type, "swipe"); + assert.equal(context.kind, "no-new-user"); + assert.equal(context.swipedAssistantFloor, 7); + assert.equal(context.chatId, chatId); +} + +{ + let chatId = "chat-dry-run"; + let now = 2000; + const tracker = createGenerationContextTracker({ + getCurrentChatId: () => chatId, + now: () => now, + }); + + const original = tracker.begin("normal", { existing: true }); + assert.equal(tracker.begin("swipe", {}, { dryRun: true }), null); + assert.deepEqual(tracker.get(), original); + + now += 1; + assert.equal(tracker.update("regenerate", {}, { dryRun: true }), null); + assert.deepEqual(tracker.get(), original); +} + +{ + let chatId = "chat-update"; + let now = 3000; + const tracker = createGenerationContextTracker({ + getCurrentChatId: () => chatId, + now: () => now, + }); + + tracker.begin("regenerate"); + now += 25; + const context = tracker.update( + "regenerate", + {}, + { phase: "GENERATION_AFTER_COMMANDS" }, + ); + + assert.equal(context.type, "regenerate"); + assert.equal(context.kind, "no-new-user"); + assert.equal(context.afterCommandsAt, now); +} + +{ + let chatId = "chat-ttl"; + let now = 4000; + const tracker = createGenerationContextTracker({ + getCurrentChatId: () => chatId, + now: () => now, + ttlMs: 10, + }); + + tracker.begin("normal"); + now += 11; + + assert.equal(tracker.get(), null); + assert.equal(tracker.get({ allowStale: true }), null); +} + +{ + let chatId = "chat-original"; + let now = 5000; + const tracker = createGenerationContextTracker({ + getCurrentChatId: () => chatId, + now: () => now, + }); + + tracker.begin("normal"); + chatId = "chat-current"; + + assert.equal(tracker.get(), null); + + chatId = "chat-original"; + assert.equal(tracker.get(), null); +}