@nodetool-ai/websocket 0.8.2 → 0.8.3

package/dist/session/chat-turn.js

@@ -1,6 +1,6 @@
 import { registerChatRunTrace, registerWorkflowRunTrace, withRegisteredRunTrace, settleRegisteredRunTrace } from "@nodetool-ai/execution";
 import { getRunTraceScope, suppressRunTraceContent } from "@nodetool-ai/runtime";
-import { TRACE_RESTRICTED_CAPABILITY_MODULES } from "@nodetool-ai/protocol";
+import { TRACE_RESTRICTED_CAPABILITY_MODULES, isFullResourceId } from "@nodetool-ai/protocol";
 import { trace, SpanStatusCode } from "@opentelemetry/api";
 import { settleRunTrace, registerRunTraceParent, registerRunTraceParents, threadHasTraceContentTools } from "@nodetool-ai/models";
 import { capabilityModuleOf, listCapabilitySpecs } from "@nodetool-ai/agents";
@@ -8,6 +8,7 @@ import { createLogger } from "@nodetool-ai/config";
 import { randomUUID } from "node:crypto";
 import { SUPERSEDED_TOOL_RESULT, repairOrphanedToolCalls } from "../chat-tool-call-repair.js";
 import { attachChatPredictionForwarder } from "../chat-prediction-forwarder.js";
+import { admitSpend } from "../credit-gate.js";
 import { isGoogleWorkspaceEnabled } from "@nodetool-ai/config";
 import { isFunctionValue, isNonEmptyString, isNumber, isObjectLike, isRecord, isString } from "../lib/wire-values.js";
 import { storeAssetWithThumbnail } from "../lib/thumbnail.js";
@@ -248,6 +249,27 @@ const MEMORY_BLOCK_LIMIT = 100;
  * full thread load.
  */
 const SESSION_PROBE_WINDOW = 50;
+/** The most rows of a thread a turn loads as its history. */
+const THREAD_HISTORY_LIMIT = 1000;
+/**
+ * The thread's newest rows, oldest first. An ascending page would return the
+ * oldest rows instead, and a thread past the limit would then lose its latest
+ * turns, the user message just sent among them.
+ *
+ * A window cut by the limit starts at its first user row, so it never opens on
+ * a tool result whose call fell outside it.
+ */
+export async function loadRecentThreadRows(threadId, limit = THREAD_HISTORY_LIMIT) {
+    const [newestFirst] = await Message.paginate(threadId, {
+        reverse: true,
+        limit
+    });
+    const rows = newestFirst.reverse();
+    if (rows.length < limit)
+        return rows;
+    const firstUser = rows.findIndex((row) => row.role === "user");
+    return firstUser > 0 ? rows.slice(firstUser) : rows;
+}
 /**
  * Output-token worst case assumed when reserving a chat turn. The chat loop
  * sends no `maxTokens`, so nothing bounds the answer from this side; the
@@ -505,6 +527,12 @@ export class ChatTurnHandler {
         const existing = await Thread.find(userId, threadId);
         if (existing)
             return existing.id;
+        // `create` upserts on the primary key, so creating under an id another
+        // user's thread already holds would reassign that thread (and its
+        // history) to this caller.
+        if (await Thread.get(threadId)) {
+            throw new Error("Thread not found");
+        }
         const thread = await Thread.create({
             id: threadId,
             user_id: userId,
@@ -518,14 +546,24 @@ export class ChatTurnHandler {
      * Save a message dict to the database.
      * Mirrors Python's _save_message_to_db_async: pops id, type, user_id before create.
      */
-    async saveMessageToDb(messageData) {
+    async saveMessageToDb(messageData, options = {}) {
         const data = { ...messageData };
+        const clientId = data.id;
         delete data.id;
         delete data.type;
         const threadId = isString(data.thread_id) ? data.thread_id : "";
         delete data.thread_id;
         const userId = this.session.requireUserId();
         delete data.user_id;
+        // The user's own turn keeps the id the client minted for it, so the client
+        // can name that turn later (Regenerate, edit a sent message) without a
+        // round-trip. Anything that is not a fresh 32-hex id is replaced.
+        if (options.keepClientId &&
+            isString(clientId) &&
+            isFullResourceId(clientId) &&
+            !(await Message.find(clientId))) {
+            data.id = clientId;
+        }
         const message = await Message.create({ thread_id: threadId, user_id: userId, ...data });
         const scope = getRunTraceScope();
         if (scope) {
@@ -1083,7 +1121,7 @@ export class ChatTurnHandler {
         const userId = this.session.requireUserId();
         log.debug("Chat message", { threadId, model, provider: providerId });
         // Save user message to DB — matches Python's _save_message_to_db_async(data)
-        const turnMessage = await this.saveMessageToDb(data);
+        const turnMessage = await this.saveMessageToDb(data, { keepClientId: true });
         const rootContext = createRuntimeContext({ jobId: turnMessage.id, threadId, workflowId, userId, workspace: null });
         await registerChatRunTrace(rootContext, { messageId: turnMessage.id, threadId });
         return withRegisteredRunTrace(rootContext, "chat.turn", async () => {
@@ -1117,6 +1155,21 @@ export class ChatTurnHandler {
             });
             return;
         }
+        // The managed provider spends NodeTool's own keys against the user's
+        // credits. Every other managed entry point admits against the balance and
+        // the operator's model whitelist first, and so does a chat turn.
+        if (providerId === "nodetool") {
+            const decision = await admitSpend(userId, 0, [model]);
+            if (!decision.allowed) {
+                await settleRegisteredRunTrace(rootContext, "failed", decision.reason);
+                await this.session.send({
+                    type: "error",
+                    message: decision.reason,
+                    thread_id: threadId
+                });
+                return;
+            }
+        }
         // Route to the workflow processor ONLY when the client explicitly opts in
         // via `workflow_target: "workflow"`. A bare `workflow_id` is context, not a
         // routing signal: the editor binds the open workflow so `ui_*` tools target
@@ -1127,7 +1180,7 @@ export class ChatTurnHandler {
             ? data.workflow_target
             : null;
         if (workflowTarget === "workflow") {
-            await this.handleWorkflowMessage(data, requestSeq, signal);
+            await this.handleWorkflowMessage(data, turnMessage.id, requestSeq, signal);
             return;
         }
         // Route to media generation when the client requests a text-to-image or
@@ -1285,7 +1338,7 @@ export class ChatTurnHandler {
                 sessionCheckpointOverride =
                     probeSession.checkpoint + 1 + newTurns.length;
                 loadFullHistory = async () => {
-                    const [rows] = await Message.paginate(threadId, { limit: 1000 });
+                    const rows = await loadRecentThreadRows(threadId);
                     const full = convertDbMessages(historySinceCompaction(rows));
                     full.unshift(systemChatMessage());
                     const scope = getRunTraceScope();
@@ -1306,7 +1359,7 @@ export class ChatTurnHandler {
                 // all (a far-back session still resumes via the slice path). A session
                 // older than the compaction cut is not one of them: searching only the
                 // compacted slice is what keeps the newest marker winning here too.
-                const [rows] = await Message.paginate(threadId, { limit: 1000 });
+                const rows = await loadRecentThreadRows(threadId);
                 const kept = historySinceCompaction(rows);
                 chatHistory = convertDbMessages(kept);
                 priorSession = lastMatchingProviderSession(kept, providerId, model);
@@ -2063,7 +2116,7 @@ export class ChatTurnHandler {
          * turn that cannot run at all.
          */
         const compactThread = async (reason) => {
-            const [rows] = await Message.paginate(threadId, { limit: 1000 });
+            const rows = await loadRecentThreadRows(threadId);
             const cut = chooseCompactionCut(historySinceCompaction(rows), compaction.keepUserTurns, {
                 // Leave room for the summary, turn context and subsequent replies.
                 maxTokens: compactionThresholdTokens * 0.75 -
@@ -3214,7 +3267,7 @@ export class ChatTurnHandler {
      *   6. Collect output_update results
      *   7. Send response message with typed content, then the terminal chunk
      */
-    async handleWorkflowMessage(data, requestSeq, signal) {
+    async handleWorkflowMessage(data, turnMessageId, requestSeq, signal) {
         const threadId = isString(data.thread_id) ? data.thread_id : "";
         const workflowId = isString(data.workflow_id) ? data.workflow_id : null;
         const providerId = isString(data.provider)
@@ -3256,7 +3309,7 @@ export class ChatTurnHandler {
                 messagesName ??
                 "messages";
             // Build chat history for params — matches Python
-            const [dbMessages] = await Message.paginate(threadId, { limit: 1000 });
+            const dbMessages = await loadRecentThreadRows(threadId);
             const inheritedTrace = getRunTraceScope();
             seedStoredToolProvenance(dbMessages);
             if (inheritedTrace) {
@@ -3280,7 +3333,12 @@ export class ChatTurnHandler {
             // Prepare params — matches Python's WorkflowMessageProcessor
             const params = {
                 [messageInputName]: currentMessage,
-                [messagesInputName]: [...chatHistorySerialized, currentMessage]
+                // The stored history already holds this turn's message; it is appended
+                // once, as `currentMessage`.
+                [messagesInputName]: [
+                    ...chatHistorySerialized.filter((_, i) => dbMessages[i].id !== turnMessageId),
+                    currentMessage
+                ]
             };
             if (isObjectLike(data.params)) {
                 Object.assign(params, data.params);