@wrongstack/webui 0.1.9 → 0.2.0

package/dist/server/index.js

@@ -2,7 +2,6 @@
 import * as fs from "fs/promises";
 import * as os from "os";
 import * as path from "path";
-import { WebSocketServer, WebSocket } from "ws";
 import {
   Agent,
   AutoCompactionMiddleware,
@@ -13,6 +12,7 @@ import {
   DefaultErrorHandler,
   DefaultLogger,
   DefaultMemoryStore,
+  DefaultModeStore,
   DefaultModelsRegistry,
   DefaultPathResolver,
   DefaultPermissionPolicy,
@@ -21,25 +21,22 @@ import {
   DefaultSecretVault,
   DefaultSessionStore,
   DefaultSkillLoader,
-  DefaultModeStore,
   DefaultSystemPromptBuilder,
   DefaultTokenCounter,
   EventBus,
   HybridCompactor,
   ProviderRegistry,
-  ToolRegistry,
   TOKENS,
-  createDefaultPipelines,
+  ToolRegistry,
   atomicWrite,
+  createDefaultPipelines,
   migratePlaintextSecrets,
   resolveWstackPaths
 } from "@wrongstack/core";
-import {
-  buildProviderFactoriesFromRegistry,
-  makeProviderFromConfig
-} from "@wrongstack/providers";
+import { buildProviderFactoriesFromRegistry, makeProviderFromConfig } from "@wrongstack/providers";
+import { forgetTool, rememberTool } from "@wrongstack/tools";
 import { builtinTools } from "@wrongstack/tools/builtin";
-import { rememberTool, forgetTool } from "@wrongstack/tools";
+import { WebSocket, WebSocketServer } from "ws";
 async function bootConfig() {
   const cwd = process.cwd();
   const pathResolver = new DefaultPathResolver(cwd);
@@ -98,11 +95,14 @@ async function startWebUI(opts = {}) {
   container.bind(TOKENS.RetryPolicy, () => new DefaultRetryPolicy());
   container.bind(TOKENS.ErrorHandler, () => new DefaultErrorHandler());
   container.bind(TOKENS.ModelsRegistry, () => modelsRegistry);
-  container.bind(TOKENS.PermissionPolicy, () => new DefaultPermissionPolicy({
-    trustFile: wpaths.projectTrust,
-    yolo: false,
-    promptDelegate: void 0
-  }));
+  container.bind(
+    TOKENS.PermissionPolicy,
+    () => new DefaultPermissionPolicy({
+      trustFile: wpaths.projectTrust,
+      yolo: false,
+      promptDelegate: void 0
+    })
+  );
   const providerRegistry = new ProviderRegistry();
   try {
     const factories = await buildProviderFactoriesFromRegistry({
@@ -253,10 +253,24 @@ async function startWebUI(opts = {}) {
       mode: modeId
     };
   }
-  const wssPrimary = new WebSocketServer({ port: wsPort, host: wsHost });
-  const wssSecondary = wsHost === "127.0.0.1" ? new WebSocketServer({ port: wsPort, host: "::1" }) : null;
+  const verifyClient = (info) => {
+    const origin = info.origin;
+    if (!origin) return true;
+    try {
+      const { hostname } = new URL(origin);
+      return hostname === "localhost" || hostname === "127.0.0.1" || hostname === "::1";
+    } catch {
+      return false;
+    }
+  };
+  const wssPrimary = new WebSocketServer({
+    port: wsPort,
+    host: wsHost,
+    verifyClient
+  });
+  const wssSecondary = wsHost === "127.0.0.1" ? new WebSocketServer({ port: wsPort, host: "::1", verifyClient }) : null;
   const clients = /* @__PURE__ */ new Map();
-  let abortController = null;
+  let runLock = null;
   console.log(
     `[WebUI] WebSocket server running on ws://${wsHost}:${wsPort}` + (wssSecondary ? ` (and ws://[::1]:${wsPort})` : "")
   );
@@ -270,6 +284,9 @@ async function startWebUI(opts = {}) {
     events.on("provider.text_delta", (e) => {
       broadcast({ type: "provider.text_delta", payload: { text: e.text, messageId: "current" } });
     });
+    events.on("provider.thinking_delta", (e) => {
+      broadcast({ type: "provider.thinking_delta", payload: { text: e.text } });
+    });
     events.on("tool.started", (e) => {
       broadcast({
         type: "tool.started",
@@ -391,10 +408,20 @@ async function startWebUI(opts = {}) {
     switch (msg.type) {
       case "user_message": {
         const content = msg.payload.content;
-        abortController?.abort();
-        abortController = new AbortController();
+        if (runLock) {
+          send(ws, {
+            type: "error",
+            payload: {
+              phase: "user_message",
+              message: "Agent is already processing a request. Wait for the current run to finish."
+            }
+          });
+          break;
+        }
+        runLock = new AbortController();
+        const thisRun = runLock;
         try {
-          const result = await agent.run(content, { signal: abortController.signal });
+          const result = await agent.run(content, { signal: thisRun.signal });
           send(ws, {
             type: "run.result",
             payload: {
@@ -417,12 +444,14 @@ async function startWebUI(opts = {}) {
             }
           });
         } finally {
-          abortController = null;
+          if (runLock === thisRun) {
+            runLock = null;
+          }
         }
         break;
       }
       case "abort":
-        abortController?.abort();
+        runLock?.abort();
         broadcast({ type: "error", payload: { phase: "abort", message: "User aborted" } });
         break;
       case "ping":
@@ -468,10 +497,7 @@ async function startWebUI(opts = {}) {
             return String(c);
           }
         };
-        const sysTokens = context.systemPrompt.reduce(
-          (acc, b) => acc + estimate(b.text ?? ""),
-          0
-        );
+        const sysTokens = context.systemPrompt.reduce((acc, b) => acc + estimate(b.text ?? ""), 0);
         const tools = toolRegistry.list();
         const toolBreakdown = tools.map((t) => {
           const schema = t.inputSchema ?? {};
@@ -511,7 +537,11 @@ async function startWebUI(opts = {}) {
             total,
             systemPrompt: sysTokens,
             tools: { total: toolTokens, count: tools.length, breakdown: toolBreakdown },
-            messages: { total: msgTokens, count: context.messages.length, breakdown: messageBreakdown }
+            messages: {
+              total: msgTokens,
+              count: context.messages.length,
+              breakdown: messageBreakdown
+            }
           }
         });
         break;
@@ -786,7 +816,11 @@ async function startWebUI(opts = {}) {
         const { text, scope } = msg.payload;
         try {
           const removed = await memoryStore.forget(text, scope ?? "project-memory");
-          sendResult(ws, removed > 0, removed > 0 ? `Removed ${removed} entr${removed === 1 ? "y" : "ies"}` : "No matching entries");
+          sendResult(
+            ws,
+            removed > 0,
+            removed > 0 ? `Removed ${removed} entr${removed === 1 ? "y" : "ies"}` : "No matching entries"
+          );
         } catch (err) {
           sendResult(ws, false, err instanceof Error ? err.message : String(err));
         }
@@ -819,7 +853,11 @@ async function startWebUI(opts = {}) {
         } catch (err) {
           send(ws, {
             type: "skills.list",
-            payload: { skills: [], enabled: true, error: err instanceof Error ? err.message : String(err) }
+            payload: {
+              skills: [],
+              enabled: true,
+              error: err instanceof Error ? err.message : String(err)
+            }
           });
         }
         break;
@@ -894,7 +932,8 @@ async function startWebUI(opts = {}) {
           for (const e of entries) {
             if (results.length >= 600) return;
             if (e.name.startsWith(".") && e.name !== ".wrongstack" && e.name !== ".env.example") {
-              if (e.name !== ".gitignore" && e.name !== ".eslintrc" && e.name !== ".prettierrc") continue;
+              if (e.name !== ".gitignore" && e.name !== ".eslintrc" && e.name !== ".prettierrc")
+                continue;
             }
             const childRel = rel ? `${rel}/${e.name}` : e.name;
             if (e.isDirectory()) {
@@ -948,7 +987,11 @@ async function startWebUI(opts = {}) {
         } catch (err) {
           send(ws, {
             type: "modes.list",
-            payload: { modes: [], activeId: "default", error: err instanceof Error ? err.message : String(err) }
+            payload: {
+              modes: [],
+              activeId: "default",
+              error: err instanceof Error ? err.message : String(err)
+            }
           });
         }
         break;
@@ -964,6 +1007,22 @@ async function startWebUI(opts = {}) {
             await modeStore.setActiveMode(id);
           }
           modeId = id;
+          const modePrompt2 = id === "default" ? "" : (await modeStore.getMode(id))?.prompt ?? "";
+          const freshBuilder = new DefaultSystemPromptBuilder({
+            memoryStore,
+            skillLoader,
+            modeStore,
+            modeId: id,
+            modePrompt: modePrompt2,
+            modelCapabilities
+          });
+          context.systemPrompt = await freshBuilder.build({
+            cwd: projectRoot,
+            projectRoot,
+            tools: toolRegistry.list(),
+            provider: config.provider,
+            model: config.model
+          });
           sendResult(ws, true, `Switched to mode "${id}"`);
           broadcast({
             type: "session.start",
@@ -1125,7 +1184,9 @@ async function startWebUI(opts = {}) {
         baseUrl: payload.baseUrl
       };
       if (payload.apiKey) {
-        newProv.apiKeys = [{ label: "default", apiKey: payload.apiKey, createdAt: (/* @__PURE__ */ new Date()).toISOString() }];
+        newProv.apiKeys = [
+          { label: "default", apiKey: payload.apiKey, createdAt: (/* @__PURE__ */ new Date()).toISOString() }
+        ];
         newProv.activeKey = "default";
       }
       providers[payload.id] = newProv;
@@ -1152,7 +1213,11 @@ async function startWebUI(opts = {}) {
   const shutdown = async () => {
     console.log("[WebUI] Shutting down...");
     try {
-      await session.append({ type: "session_end", ts: (/* @__PURE__ */ new Date()).toISOString(), usage: tokenCounter.total() });
+      await session.append({
+        type: "session_end",
+        ts: (/* @__PURE__ */ new Date()).toISOString(),
+        usage: tokenCounter.total()
+      });
       await session.close();
     } catch (e) {
       console.warn("[WebUI] Error closing session:", e);