@blokkli/editor 2.0.0-alpha.32 → 2.0.0-alpha.34

package/dist/modules/agent/runtime/server/Session.js

@@ -0,0 +1,734 @@
+import { buildSystemPrompt, buildSystemPromptText } from "./agentPrompt.js";
+import { provider, models } from "#blokkli-build/agent-server";
+import {
+  send,
+  KEEP_RECENT_TURNS,
+  resolveSkills,
+  transformText,
+  classifyError,
+  pruneMessages,
+  pruneForPersistence,
+  computeStateHash,
+  verifyStateHash,
+  validateMessages
+} from "./helpers.js";
+import { buildDefinition } from "./server-tools/index.js";
+import loadSkillTool from "./server-tools/load_skill/index.js";
+import loadToolsTool from "./server-tools/load_tools/index.js";
+import createPlanTool from "./server-tools/create_plan/index.js";
+import completePlanStepTool from "./server-tools/complete_plan_step/index.js";
+const serverTools = [
+  loadSkillTool,
+  loadToolsTool,
+  createPlanTool,
+  completePlanStepTool
+];
+export class Session {
+  messages = [];
+  pendingToolCalls = /* @__PURE__ */ new Map();
+  abortController = null;
+  isProcessing = false;
+  /** Eager tools sent to the LLM on every turn */
+  tools = [];
+  /** Lazy tools held back until activated via load_tools */
+  lazyTools = [];
+  /** Names of lazy tools that have been activated via load_tools */
+  activatedLazyTools = /* @__PURE__ */ new Set();
+  /** Names of skills that have been loaded via load_skill */
+  loadedSkills = /* @__PURE__ */ new Set();
+  /** Page context received from client on init */
+  pageContext;
+  /** Current plan (null when no plan is active) */
+  plan = null;
+  /** Pending plan approval promise resolver */
+  pendingPlanApproval = null;
+  /** Last tools payload sent to the provider (for transcript debugging) */
+  lastProviderTools = null;
+  // --------------------------------------------------------------------------
+  // Public methods
+  // --------------------------------------------------------------------------
+  init(tools, pageContext) {
+    this.tools = tools.filter((t) => !t.lazy);
+    this.lazyTools = tools.filter((t) => !!t.lazy);
+    this.activatedLazyTools = /* @__PURE__ */ new Set();
+    this.loadedSkills = /* @__PURE__ */ new Set();
+    this.pageContext = pageContext;
+  }
+  start(peer, prompt, apiKey, authSecret, selectedUuids) {
+    if (this.isProcessing) {
+      send(peer, {
+        type: "error",
+        errorType: "bad_request",
+        message: "Agent is already processing a request"
+      });
+      return;
+    }
+    this.runAgentLoop(peer, prompt, apiKey, authSecret, selectedUuids);
+  }
+  resolveToolResult(callId, result) {
+    const pending = this.pendingToolCalls.get(callId);
+    if (pending) {
+      pending.resolve(result);
+    }
+  }
+  cancel(peer) {
+    this.abortController?.abort();
+    for (const pending of this.pendingToolCalls.values()) {
+      pending.reject(new Error("Cancelled"));
+    }
+    this.pendingToolCalls.clear();
+    if (this.pendingPlanApproval) {
+      this.pendingPlanApproval.resolve(false);
+      this.pendingPlanApproval = null;
+    }
+    send(peer, { type: "done" });
+  }
+  acceptChanges(peer, authSecret) {
+    this.safePushUserMessage(`[System: Changes accepted and applied.]`);
+    send(peer, { type: "done", message: "Changes accepted" });
+    this.sendConversationState(peer, authSecret);
+  }
+  rejectChanges(peer, authSecret) {
+    this.safePushUserMessage(
+      `[System: Changes rejected. All pending changes have been reverted. The page is back to its previous state.]`
+    );
+    send(peer, { type: "done", message: "Changes rejected" });
+    this.sendConversationState(peer, authSecret);
+  }
+  approvePlan() {
+    if (this.pendingPlanApproval) {
+      this.pendingPlanApproval.resolve(true);
+      this.pendingPlanApproval = null;
+    }
+  }
+  rejectPlan() {
+    if (this.pendingPlanApproval) {
+      this.pendingPlanApproval.resolve(false);
+      this.pendingPlanApproval = null;
+    }
+  }
+  getTranscript(peer) {
+    send(peer, {
+      type: "transcript",
+      content: this.buildTranscript()
+    });
+  }
+  newConversation(peer, authSecret) {
+    this.abortController?.abort();
+    this.messages = [];
+    this.activatedLazyTools.clear();
+    this.loadedSkills.clear();
+    this.plan = null;
+    if (this.pendingPlanApproval) {
+      this.pendingPlanApproval.resolve(false);
+      this.pendingPlanApproval = null;
+    }
+    send(peer, { type: "done" });
+    this.sendConversationState(peer, authSecret);
+  }
+  cleanup() {
+    this.abortController?.abort();
+    for (const pending of this.pendingToolCalls.values()) {
+      pending.reject(new Error("Session closed"));
+    }
+    this.pendingToolCalls.clear();
+    if (this.pendingPlanApproval) {
+      this.pendingPlanApproval.resolve(false);
+      this.pendingPlanApproval = null;
+    }
+    this.plan = null;
+    this.messages = [];
+    this.tools = [];
+    this.lazyTools = [];
+    this.activatedLazyTools.clear();
+    this.loadedSkills.clear();
+    this.pageContext = void 0;
+  }
+  /**
+   * Build a persistence snapshot and send it to the client.
+   */
+  sendConversationState(peer, authSecret) {
+    const state = this.getConversationStateForPersistence(authSecret);
+    send(peer, { type: "conversation_state", state });
+  }
+  /**
+   * Create an aggressively pruned snapshot of the current conversation
+   * for client-side persistence.
+   */
+  getConversationStateForPersistence(authSecret) {
+    const prunedMessages = pruneForPersistence(this.messages);
+    const activatedLazyTools = Array.from(this.activatedLazyTools);
+    const hash = computeStateHash(
+      prunedMessages,
+      activatedLazyTools,
+      authSecret
+    );
+    return {
+      messages: prunedMessages,
+      activatedLazyTools,
+      hash
+    };
+  }
+  /**
+   * Restore conversation state from a client-provided snapshot.
+   * Verifies HMAC integrity before loading.
+   */
+  restoreConversation(state, authSecret) {
+    if (!verifyStateHash(state, authSecret)) {
+      return { success: false, reason: "Invalid state hash" };
+    }
+    const issues = validateMessages(state.messages);
+    if (issues.length > 0) {
+      return { success: false, reason: "Invalid message structure" };
+    }
+    this.messages = state.messages;
+    const validLazyToolNames = new Set(this.lazyTools.map((t) => t.name));
+    this.activatedLazyTools = new Set(
+      state.activatedLazyTools.filter((name) => validLazyToolNames.has(name))
+    );
+    this.plan = null;
+    return { success: true };
+  }
+  // --------------------------------------------------------------------------
+  // Private methods
+  // --------------------------------------------------------------------------
+  /**
+   * Convert the server plan to a client-facing plan (strips descriptions).
+   */
+  toClientPlan() {
+    if (!this.plan) return null;
+    return {
+      title: this.plan.title,
+      steps: this.plan.steps.map(
+        (s) => ({
+          label: s.label,
+          status: s.status
+        })
+      )
+    };
+  }
+  /**
+   * Build an ActivePlanContext from the current plan state,
+   * or undefined if no plan step is in progress.
+   */
+  getActivePlanContext() {
+    const currentStep = this.plan?.steps.find((s) => s.status === "in_progress");
+    if (!this.plan || !currentStep) return void 0;
+    return {
+      title: this.plan.title,
+      totalSteps: this.plan.steps.length,
+      completedSteps: this.plan.steps.filter((s) => s.status === "completed").length,
+      currentStep: {
+        label: currentStep.label,
+        description: currentStep.description
+      },
+      remainingSteps: this.plan.steps.filter((s) => s.status === "pending").map((s) => s.label)
+    };
+  }
+  /**
+   * Wait for the user to approve or reject the plan.
+   */
+  waitForPlanApproval() {
+    return new Promise((resolve) => {
+      this.pendingPlanApproval = { resolve };
+    });
+  }
+  async runAgentLoop(peer, prompt, apiKey, authSecret, selectedUuids) {
+    if (this.tools.length === 0) {
+      send(peer, {
+        type: "error",
+        errorType: "bad_request",
+        message: "No tools available. Client must send init message with tools first."
+      });
+      return;
+    }
+    if (!this.pageContext) {
+      send(peer, {
+        type: "error",
+        errorType: "bad_request",
+        message: "No page context available. Client must send init message with pageContext first."
+      });
+      return;
+    }
+    const resolvedSkills = resolveSkills(this.pageContext);
+    const lazyToolSummaries = this.lazyTools.map((t) => ({
+      name: t.name,
+      description: t.description
+    }));
+    const userParts = [];
+    if (selectedUuids?.length) {
+      userParts.push(
+        `[User has selected the following paragraphs: ${selectedUuids.join(", ")}]`
+      );
+    }
+    userParts.push(prompt);
+    const userMessage = {
+      role: "user",
+      content: userParts.join("\n\n")
+    };
+    this.messages.push(userMessage);
+    this.abortController = new AbortController();
+    this.isProcessing = true;
+    let toolCallCounter = 0;
+    let planRetryCount = 0;
+    let streamRetryCount = 0;
+    const MAX_STREAM_RETRIES = 1;
+    try {
+      while (true) {
+        if (this.abortController.signal.aborted) {
+          break;
+        }
+        send(peer, { type: "thinking" });
+        const assistantContent = [];
+        const toolResults = [];
+        const extraBlocks = [];
+        let messagesCommittedByPlanTool = false;
+        let currentToolUse = null;
+        let currentTextContent = "";
+        let inTextBlock = false;
+        const activatedTools = this.lazyTools.filter(
+          (t) => this.activatedLazyTools.has(t.name)
+        );
+        const unloadedLazyTools = this.lazyTools.filter(
+          (t) => !this.activatedLazyTools.has(t.name)
+        );
+        const defCtx = {
+          resolvedSkills,
+          plan: this.plan,
+          unloadedLazyTools
+        };
+        const serverToolDefs = serverTools.map((t) => buildDefinition(t, defCtx)).filter((d) => d !== null);
+        const allTools = [...serverToolDefs, ...this.tools, ...activatedTools];
+        const systemPrompt = buildSystemPrompt(
+          this.pageContext,
+          resolvedSkills,
+          lazyToolSummaries,
+          this.getActivePlanContext(),
+          this.loadedSkills
+        );
+        const stream = provider.createStream(
+          {
+            apiKey,
+            model: (models.find((m) => m.isDefault) || models[0]).name
+          },
+          {
+            systemPrompt,
+            messages: this.messages,
+            tools: allTools,
+            maxTokens: 4096,
+            signal: this.abortController.signal
+          }
+        );
+        try {
+          for await (const event of stream) {
+            if (this.abortController?.signal.aborted) {
+              break;
+            }
+            switch (event.type) {
+              case "debug_request":
+                this.lastProviderTools = event.tools;
+                break;
+              case "text_start":
+                inTextBlock = true;
+                currentTextContent = "";
+                break;
+              case "text_delta":
+                if (inTextBlock) {
+                  const transformed = transformText(event.text);
+                  currentTextContent += transformed;
+                  send(peer, { type: "text_delta", content: transformed });
+                }
+                break;
+              case "text_end":
+                if (inTextBlock && currentTextContent) {
+                  assistantContent.push({
+                    type: "text",
+                    text: currentTextContent
+                  });
+                }
+                currentTextContent = "";
+                inTextBlock = false;
+                break;
+              case "tool_use_start":
+                currentToolUse = {
+                  id: event.id,
+                  name: event.name,
+                  inputJson: ""
+                };
+                break;
+              case "tool_use_delta":
+                if (currentToolUse) {
+                  currentToolUse.inputJson += event.partial_json;
+                }
+                break;
+              case "tool_use_end":
+                if (currentToolUse) {
+                  let input;
+                  try {
+                    input = JSON.parse(currentToolUse.inputJson || "{}");
+                  } catch {
+                    input = {};
+                    assistantContent.push({
+                      type: "tool_use",
+                      id: currentToolUse.id,
+                      name: currentToolUse.name,
+                      input
+                    });
+                    toolResults.push({
+                      type: "tool_result",
+                      tool_use_id: currentToolUse.id,
+                      content: JSON.stringify({
+                        error: "Your tool call produced malformed JSON input. Please try again."
+                      }),
+                      is_error: true
+                    });
+                    currentToolUse = null;
+                    break;
+                  }
+                  const callId = `tc_${toolCallCounter++}`;
+                  assistantContent.push({
+                    type: "tool_use",
+                    id: currentToolUse.id,
+                    name: currentToolUse.name,
+                    input
+                  });
+                  const matchedServerTool = serverTools.find(
+                    (t) => t.name === currentToolUse.name
+                  );
+                  if (matchedServerTool) {
+                    if (matchedServerTool.isAvailable && !matchedServerTool.isAvailable(defCtx)) {
+                      toolResults.push({
+                        type: "tool_result",
+                        tool_use_id: currentToolUse.id,
+                        content: JSON.stringify({
+                          error: "This tool is not available right now."
+                        }),
+                        is_error: true
+                      });
+                      currentToolUse = null;
+                      break;
+                    }
+                    const handlerCtx = {
+                      toolUseId: currentToolUse.id,
+                      send: (msg) => send(peer, msg),
+                      resolvedSkills,
+                      lazyTools: this.lazyTools,
+                      activatedLazyTools: this.activatedLazyTools,
+                      loadedSkills: this.loadedSkills,
+                      plan: this.plan,
+                      setPlan: (p) => {
+                        this.plan = p;
+                      },
+                      toClientPlan: () => this.toClientPlan(),
+                      waitForPlanApproval: () => this.waitForPlanApproval(),
+                      assistantContent,
+                      commitMessagesEarly: (toolResult) => {
+                        if (assistantContent.length) {
+                          this.messages.push({
+                            role: "assistant",
+                            content: [...assistantContent]
+                          });
+                          assistantContent.length = 0;
+                        }
+                        this.messages.push({
+                          role: "user",
+                          content: [toolResult]
+                        });
+                      },
+                      updateLastToolResult: (toolUseId, content) => {
+                        const lastMsg = this.messages[this.messages.length - 1];
+                        if (lastMsg.role === "user" && Array.isArray(lastMsg.content)) {
+                          const resultBlock = lastMsg.content.find(
+                            (b) => b.type === "tool_result" && b.tool_use_id === toolUseId
+                          );
+                          if (resultBlock && resultBlock.type === "tool_result") {
+                            resultBlock.content = content;
+                          }
+                        }
+                      }
+                    };
+                    try {
+                      const parsed = matchedServerTool.inputSchema(defCtx).parse(input);
+                      const result = await matchedServerTool.handle(
+                        handlerCtx,
+                        parsed
+                      );
+                      toolResults.push(...result.toolResults);
+                      if (result.extraBlocks) {
+                        extraBlocks.push(...result.extraBlocks);
+                      }
+                      if (result.messagesCommitted) {
+                        messagesCommittedByPlanTool = true;
+                      }
+                    } catch (e) {
+                      toolResults.push({
+                        type: "tool_result",
+                        tool_use_id: currentToolUse.id,
+                        content: JSON.stringify({
+                          error: `Invalid input: ${e.message}`
+                        }),
+                        is_error: true
+                      });
+                    }
+                    currentToolUse = null;
+                    break;
+                  }
+                  send(peer, {
+                    type: "tool_call",
+                    callId,
+                    tool: currentToolUse.name,
+                    params: input
+                  });
+                  try {
+                    const clientResult = await this.waitForToolResult(callId);
+                    if (clientResult.error) {
+                      toolResults.push({
+                        type: "tool_result",
+                        tool_use_id: currentToolUse.id,
+                        content: JSON.stringify({
+                          error: clientResult.error
+                        }),
+                        is_error: true
+                      });
+                    } else {
+                      let resultForLLM = clientResult.result;
+                      if (typeof resultForLLM === "object" && resultForLLM !== null && "agentMessage" in resultForLLM) {
+                        const { agentMessage, ...rest } = resultForLLM;
+                        resultForLLM = { ...rest, label: agentMessage };
+                      }
+                      toolResults.push({
+                        type: "tool_result",
+                        tool_use_id: currentToolUse.id,
+                        content: JSON.stringify(resultForLLM)
+                      });
+                    }
+                  } catch (error) {
+                    toolResults.push({
+                      type: "tool_result",
+                      tool_use_id: currentToolUse.id,
+                      content: JSON.stringify({
+                        error: error.message
+                      }),
+                      is_error: true
+                    });
+                  }
+                  currentToolUse = null;
+                }
+                break;
+              case "message_end":
+                if (event.inputTokens !== void 0 && event.outputTokens !== void 0) {
+                  const defaultModel = models.find((m) => m.isDefault) || models[0];
+                  send(peer, {
+                    type: "usage",
+                    usage: {
+                      inputTokens: event.inputTokens,
+                      outputTokens: event.outputTokens,
+                      cacheCreationInputTokens: event.cacheCreationInputTokens ?? 0,
+                      cacheReadInputTokens: event.cacheReadInputTokens ?? 0,
+                      pricing: defaultModel?.pricing ?? null
+                    }
+                  });
+                }
+                break;
+              case "error":
+                throw event.error;
+            }
+          }
+        } catch (streamError) {
+          const isApiError = typeof streamError.status === "number";
+          const isRetryable = !isApiError && toolResults.length === 0;
+          if (isRetryable && streamRetryCount < MAX_STREAM_RETRIES) {
+            streamRetryCount++;
+            console.warn(
+              `[blokkli agent] Transient stream error, retrying (attempt ${streamRetryCount}/${MAX_STREAM_RETRIES}):`,
+              streamError
+            );
+            continue;
+          }
+          if (isRetryable && this.messages[this.messages.length - 1] === userMessage) {
+            this.messages.pop();
+            console.error(
+              "[blokkli agent] Transient stream error, retries exhausted:",
+              streamError
+            );
+            const classified = classifyError(streamError);
+            send(peer, { type: "error", ...classified, retryable: true });
+            break;
+          }
+          throw streamError;
+        }
+        if (this.abortController?.signal.aborted) {
+          break;
+        }
+        if (messagesCommittedByPlanTool) {
+          continue;
+        }
+        if (assistantContent.length) {
+          this.messages.push({
+            role: "assistant",
+            content: assistantContent
+          });
+        }
+        if (toolResults.length) {
+          this.messages.push({
+            role: "user",
+            content: [...toolResults, ...extraBlocks]
+          });
+          planRetryCount = 0;
+          streamRetryCount = 0;
+        }
+        if (toolResults.length === 0) {
+          const hasActivePlan = this.plan && this.plan.steps.some((s) => s.status === "in_progress");
+          if (hasActivePlan && planRetryCount < 2) {
+            planRetryCount++;
+            this.safePushUserMessage("[Continue with the plan.]");
+            continue;
+          }
+          const finalMessage = assistantContent.filter(
+            (c) => c.type === "text"
+          ).map((t) => t.text).join("\n").trim();
+          send(peer, { type: "done", message: finalMessage || void 0 });
+          break;
+        }
+      }
+    } catch (error) {
+      if (error.name !== "AbortError" && !this.abortController?.signal.aborted) {
+        console.error("Agent loop error:", error);
+        const classified = classifyError(error);
+        send(peer, { type: "error", ...classified });
+      }
+    } finally {
+      this.isProcessing = false;
+      this.abortController = null;
+      pruneMessages(
+        this.messages,
+        KEEP_RECENT_TURNS,
+        this.buildToolMetadataMap()
+      );
+      if (this.loadedSkills.size > 0) {
+        this.reconcileLoadedSkills();
+      }
+      this.sendConversationState(peer, authSecret);
+    }
+  }
+  waitForToolResult(callId) {
+    return new Promise((resolve, reject) => {
+      this.pendingToolCalls.set(callId, {
+        resolve: (result) => {
+          this.pendingToolCalls.delete(callId);
+          resolve(result);
+        },
+        reject: (error) => {
+          this.pendingToolCalls.delete(callId);
+          reject(error);
+        }
+      });
+    });
+  }
+  /**
+   * Bug 3 fix: safely push a user message, merging with the last message
+   * if it's also a user message to avoid consecutive same-role messages.
+   */
+  safePushUserMessage(text) {
+    const lastMessage = this.messages[this.messages.length - 1];
+    if (lastMessage && lastMessage.role === "user") {
+      if (typeof lastMessage.content === "string") {
+        lastMessage.content = lastMessage.content + "\n" + text;
+      } else {
+        lastMessage.content.push({ type: "text", text });
+      }
+    } else {
+      this.messages.push({ role: "user", content: text });
+    }
+  }
+  /**
+   * Remove skills from loadedSkills whose text content has been pruned
+   * from the conversation messages.
+   */
+  reconcileLoadedSkills() {
+    const presentSkills = /* @__PURE__ */ new Set();
+    for (const msg of this.messages) {
+      if (msg.role !== "user" || !Array.isArray(msg.content)) continue;
+      for (const block of msg.content) {
+        if (block.type === "skill") {
+          presentSkills.add(block.name);
+        }
+      }
+    }
+    for (const name of this.loadedSkills) {
+      if (!presentSkills.has(name)) {
+        this.loadedSkills.delete(name);
+      }
+    }
+  }
+  /**
+   * Build a map of tool name to pruning metadata from all known tools.
+   */
+  buildToolMetadataMap() {
+    const map = /* @__PURE__ */ new Map();
+    for (const tool of [...this.tools, ...this.lazyTools]) {
+      if (tool.volatile) {
+        map.set(tool.name, { volatile: true });
+      }
+    }
+    return map;
+  }
+  buildTranscript() {
+    const lines = [];
+    let systemPrompt = "(No page context available)";
+    if (this.pageContext) {
+      const resolvedSkills = resolveSkills(this.pageContext);
+      const lazyToolSummaries = this.lazyTools.map((t) => ({
+        name: t.name,
+        description: t.description
+      }));
+      systemPrompt = buildSystemPromptText(
+        this.pageContext,
+        resolvedSkills,
+        lazyToolSummaries,
+        this.getActivePlanContext()
+      );
+    }
+    lines.push("=".repeat(80));
+    lines.push("SYSTEM PROMPT");
+    lines.push("=".repeat(80));
+    lines.push(systemPrompt);
+    lines.push("");
+    if (this.lastProviderTools) {
+      lines.push("=".repeat(80));
+      lines.push("TOOLS (exact payload sent to provider)");
+      lines.push("=".repeat(80));
+      lines.push(JSON.stringify(this.lastProviderTools, null, 2));
+      lines.push("");
+    }
+    for (const message of this.messages) {
+      lines.push("=".repeat(80));
+      lines.push(`${message.role.toUpperCase()}`);
+      lines.push("=".repeat(80));
+      if (typeof message.content === "string") {
+        lines.push(message.content);
+      } else if (Array.isArray(message.content)) {
+        for (const block of message.content) {
+          if (block.type === "text") {
+            lines.push(block.text);
+          } else if (block.type === "skill") {
+            lines.push(`[Skill: ${block.name}]`);
+            lines.push(block.text);
+          } else if (block.type === "tool_use") {
+            lines.push(`[Tool Call: ${block.name}]`);
+            lines.push(JSON.stringify(block.input, null, 2));
+          } else if (block.type === "tool_result") {
+            lines.push(`[Tool Result: ${block.tool_use_id}]`);
+            try {
+              lines.push(JSON.stringify(JSON.parse(block.content), null, 2));
+            } catch {
+              lines.push(block.content);
+            }
+          }
+        }
+      }
+      lines.push("");
+    }
+    return lines.join("\n");
+  }
+}