testeiya 0.2.0 → 0.3.9

package/src/connection.ts

@@ -1,658 +0,0 @@
-import { randomUUID } from "node:crypto";
-import { mkdirSync, writeFileSync, readFileSync } from "node:fs";
-import { join } from "node:path";
-import { log } from "evlog";
-import { PROJECT_DIR } from "./project-dir.js";
-import { createTesteiyaSession } from "./session-factory.js";
-import { transformEvent } from "./bridge.js";
-import { attachAiDebug } from "./ai-debug.js";
-import { getSession } from "./session-store.js";
-import { resolveSessionShellEnv } from "./api/testomatio-target.js";
-import { dedupeSkillsByName, loadBundledSkills, loadCustomSkills } from "./skills.js";
-import { browserStateNotice } from "./api/playwright-cli.js";
-import {
-  contextStamp,
-  contextUpdateNotice,
-  listContext,
-  listContextFolders,
-} from "./context-store.js";
-
-// Bundled skills are static for the process; load once. Custom skills are
-// re-read per prompt (cheap dir scan) so a freshly dropped-in skill is
-// mentionable without a restart. Custom last — it overrides by name, matching
-// session-factory. Each entry has name + description + filePath (see skills.ts).
-let bundledSkillsCache: ReturnType<typeof loadBundledSkills> | null = null;
-function getSkills(cwd: string) {
-  if (!bundledSkillsCache) bundledSkillsCache = loadBundledSkills();
-  return dedupeSkillsByName([...bundledSkillsCache, ...loadCustomSkills(cwd)]);
-}
-
-/** The server emits a lightweight `ping` every HEARTBEAT_MS for the whole
- *  life of the connection. It lets the client's stall watchdog tell "alive but
- *  quiet" (long tool calls, an ask_question waiting for the user's answer, SDK
- *  continuations that outlive session.prompt()) apart from a genuinely dropped
- *  socket, so legitimate silence never surfaces a false "connection lost". */
-const HEARTBEAT_MS = 15000;
-
-/** Max rate for forwarding a running tool's streamed output tail (per tool). */
-const PARTIAL_THROTTLE_MS = 250;
-
-/** Transport-agnostic sender — `ws.send(JSON.stringify(...))` for either the
- *  `ws` library or Bun's native ServerWebSocket. */
-export type Send = (data: Record<string, any>) => void;
-
-export interface Connection {
-  /** Feed one raw text frame from the socket. */
-  handleRaw(raw: string): Promise<void>;
-  /** Tear down the agent subscription when the socket closes. */
-  close(): void;
-}
-
-/**
- * Per-connection agent driver. Holds the session/subscription state in a
- * closure so it can back any WebSocket transport — used by the unified Bun
- * app-server's WS upgrade.
- */
-export function createConnection(
-  storedSessionId: string | null,
-  send: Send
-): Connection {
-  let session: any = null;
-  let askChannel: any = null;
-  let widgetChannel: any = null;
-  let unsubscribe: (() => void) | null = null;
-  let aiDebugUnsub: (() => void) | null = null;
-  let messageId = "";
-  let lastTurnError: string | null = null;
-  let promptStartedAt = 0;
-  let promptCalls = 0;
-  let promptOutTok = 0;
-  let promptPeakCtx = 0;
-  let promptRetries = 0;
-  let sessionModel = "";
-  let conversationId: string | null = null;
-  let sessionCwd = process.cwd();
-  // The workspace-context stamp the agent last saw (system prompt at session
-  // creation, or the last update notice). A change between prompts — the Add
-  // to Workspace UI ran, or a pull seeded manual-tests — appends a refresh
-  // notice to the next prompt.
-  let contextSeenStamp = "";
-  const heartbeat = setInterval(() => send({ type: "ping" }), HEARTBEAT_MS);
-  let resumeConversationId: string | null = null;
-  // `tool-output-partial` coalescing: the SDK streams a tool's full output tail
-  // on every chunk (up to 50KB each), so forwarding raw would flood the socket.
-  // Leading-edge send + one trailing flush per PARTIAL_THROTTLE_MS, per tool.
-  const partialTimers = new Map<string, ReturnType<typeof setTimeout>>();
-  const pendingPartials = new Map<string, Record<string, any>>();
-
-  function queuePartial(msg: Record<string, any>): void {
-    const id = msg.toolCallId as string;
-    if (partialTimers.has(id)) {
-      pendingPartials.set(id, msg);
-      return;
-    }
-    send(msg);
-    partialTimers.set(
-      id,
-      setTimeout(() => {
-        partialTimers.delete(id);
-        const pending = pendingPartials.get(id);
-        if (!pending) return;
-        pendingPartials.delete(id);
-        queuePartial(pending);
-      }, PARTIAL_THROTTLE_MS)
-    );
-  }
-
-  function dropPartial(id: string): void {
-    const timer = partialTimers.get(id);
-    if (timer) clearTimeout(timer);
-    partialTimers.delete(id);
-    pendingPartials.delete(id);
-  }
-
-  async function handleMessage(msg: any): Promise<void> {
-    switch (msg.type) {
-      case "prompt": {
-        try {
-          // Create session on first prompt
-          if (!session) {
-            const { type: _, message: __, ...sessionParams } = msg;
-
-            // Merge stored session config if a session ID was provided
-            if (storedSessionId) {
-              const stored = getSession(storedSessionId);
-              if (stored) {
-                sessionParams.cwd = stored.cwd;
-                sessionParams.promptContext = stored.promptContext;
-                sessionParams.backendUrl = stored.backendUrl;
-                sessionParams.tokens = stored.tokens;
-                sessionParams.projects = stored.projects;
-                sessionParams.trusted = stored.trusted;
-              }
-            }
-
-            sessionCwd = sessionParams.cwd || process.cwd();
-            if (sessionParams.backendUrl) {
-              process.env.TESTOMATIO_URL = sessionParams.backendUrl;
-            }
-            const tokenEntries = sessionParams.tokens
-              ? Object.entries(sessionParams.tokens)
-              : [];
-            // NOTE: process.env is process-GLOBAL and the unified server is one
-            // process serving many WS connections, so concurrent sessions share
-            // the last-set token (accepted: desktop and single-user web are the
-            // real deployments; per-session isolation needs an SDK env hook).
-            // The shell-env.ts patch overlays these live values onto the SDK's
-            // otherwise frozen-after-first-use shell env, so whatever is set
-            // here is what every agent bash command actually sees.
-            let tokenSource = "none";
-            if (tokenEntries.length === 1) {
-              process.env.TESTOMATIO = tokenEntries[0][1] as string;
-              tokenSource = "managed";
-            } else if (tokenEntries.length > 1) {
-              // A single global TESTOMATIO would be wrong for all but one
-              // project — unset it; the agent uses per-project MCP / app Sync.
-              delete process.env.TESTOMATIO;
-              tokenSource = "managed-multi";
-            } else {
-              const resolved = await resolveSessionShellEnv(
-                {
-                  tokens: sessionParams.tokens,
-                  backendUrl: sessionParams.backendUrl,
-                  projects: sessionParams.projects ?? [],
-                },
-                sessionCwd
-              );
-              if (resolved) {
-                Object.assign(process.env, resolved.env);
-                tokenSource = resolved.source;
-                sessionParams.connection = {
-                  tokenAvailable: true,
-                  projectId: resolved.projectId,
-                  title: resolved.title,
-                };
-              } else {
-                delete process.env.TESTOMATIO;
-                delete process.env.TESTOMATIO_PROJECT_ID;
-                sessionParams.connection = { tokenAvailable: false };
-              }
-            }
-
-            console.log(
-              "[session] cwd:", sessionParams.cwd || `default (${process.cwd()})`,
-              "| promptContext:", sessionParams.promptContext ? "yes" : "no",
-              "| TESTOMATIO:", process.env.TESTOMATIO ? "set" : "unset",
-              `(${tokenSource})`,
-              "| TESTOMATIO_URL:", process.env.TESTOMATIO_URL || "default"
-            );
-            const result = await createTesteiyaSession({
-              ...sessionParams,
-              mode: "web",
-              resumeConversationId: resumeConversationId ?? undefined,
-            });
-            session = result.session;
-            sessionModel = `${result.model.provider}/${result.model.id}`;
-            conversationId = (result as any).conversationId ?? null;
-            contextSeenStamp = contextStamp(sessionCwd);
-            aiDebugUnsub = attachAiDebug(session, conversationId);
-            askChannel = (result as any).askChannel ?? null;
-            widgetChannel = (result as any).widgetChannel ?? null;
-
-            const mcpManager = (result as any).mcpManager;
-            let mcpTools: string[] = [];
-            const mcpServerStatus: Record<string, string> = {};
-            try {
-              if (mcpManager) {
-                const tools = mcpManager.getTools?.() ?? [];
-                mcpTools = tools.map((t: any) => t.name ?? String(t));
-                const connections = mcpManager.getConnections?.() ?? [];
-                for (const name of Object.keys(connections)) {
-                  mcpServerStatus[name] = mcpManager.getConnectionStatus?.(name) ?? "unknown";
-                }
-              }
-            } catch (e) {
-              console.warn("[session] failed to enumerate MCP tools:", e);
-            }
-
-            // Enumerate ALL tools available to the agent (including custom tools like render_list).
-            try {
-              const sessionTools = (session as any)?.tools;
-              if (sessionTools instanceof Map) {
-                const allNames = Array.from(sessionTools.keys() as Iterable<string>);
-                console.log(`[session] all tools (${allNames.length}):`, allNames.filter((n) => !n.startsWith("mcp_")).join(", "));
-              } else if (sessionTools && typeof sessionTools === "object") {
-                console.log(`[session] tools keys:`, Object.keys(sessionTools).slice(0, 30).join(", "));
-              }
-            } catch (e) {
-              console.warn("[session] tool enumeration error:", e);
-            }
-
-            console.log(
-              `[session] created | mcp servers: ${Object.keys(mcpServerStatus).length} | mcp tools: ${mcpTools.length}`
-            );
-            if (Object.keys(mcpServerStatus).length > 0) {
-              console.log("  mcp server status:", mcpServerStatus);
-            }
-            if (mcpTools.length) {
-              console.log("  mcp tools:", mcpTools.join(", "));
-            } else if (mcpManager) {
-              console.log("  (mcpManager present but no tools returned — server may have failed to start; check stderr)");
-            } else {
-              console.log("  (mcpManager is undefined — enableMCP may be false or discovery found no configs)");
-            }
-
-            send({
-              type: "session_created",
-              model: sessionModel,
-              cwd: sessionParams.cwd || process.cwd(),
-              conversationId,
-              mcpTools,
-              mcpServers: mcpServerStatus,
-            });
-          }
-
-          // Clean up previous subscription
-          unsubscribe?.();
-
-          // New message ID per prompt
-          messageId = `msg_${randomUUID().slice(0, 12)}`;
-          lastTurnError = null;
-          promptStartedAt = Date.now();
-          promptCalls = 0;
-          promptOutTok = 0;
-          promptPeakCtx = 0;
-          promptRetries = 0;
-
-          // Subscribe to agent events and bridge them
-          unsubscribe = session.subscribe((event: any) => {
-            // The agent can re-enter its loop within one prompt (auto-retry on
-            // overload, TTSR/todo continuations) — each re-entry emits a fresh
-            // `agent_start`. Give every run its own id so the UI renders a
-            // distinct message instead of two with the same React key.
-            if (event.type === "agent_start") {
-              messageId = `msg_${randomUUID().slice(0, 12)}`;
-            }
-            // A failed model turn is otherwise buried behind a plain `finish`
-            // (see ai-debug.ts). turn_end(error) fires on every failed retry
-            // attempt, so remember only the latest and surface it when the SDK
-            // has finally given up — retries exhausted, or the run ends still
-            // in error. A recovered retry clears it.
-            if (event.type === "turn_end") {
-              lastTurnError =
-                event.message?.stopReason === "error"
-                  ? event.message?.errorMessage || "The model request failed."
-                  : null;
-              if (event.message?.role === "assistant") {
-                promptCalls++;
-                promptOutTok += event.message.usage?.output ?? 0;
-                const ctx = event.message.usage?.totalTokens ?? 0;
-                if (ctx > promptPeakCtx) promptPeakCtx = ctx;
-              }
-            }
-            if (event.type === "auto_retry_start") {
-              promptRetries++;
-            }
-            if (event.type === "auto_retry_end" && event.success) {
-              lastTurnError = null;
-            }
-            if (event.type === "auto_retry_end" && !event.success) {
-              const err = event.finalError || lastTurnError || "The model request failed.";
-              send({ type: "error", error: err });
-              lastTurnError = null;
-            }
-            const transformed = transformEvent(event, messageId);
-            if (transformed?.type === "tool-output-partial") {
-              queuePartial(transformed);
-            } else if (transformed) {
-              // The final output supersedes any queued partial — drop it so a
-              // trailing flush can't arrive after (and visually stomp) the result.
-              if (transformed.type === "tool-output-available") {
-                dropPartial(transformed.toolCallId);
-              }
-              send(transformed);
-            }
-            if (event.type === "agent_end") {
-              const promptError = lastTurnError;
-              if (lastTurnError) {
-                send({ type: "error", error: lastTurnError });
-                lastTurnError = null;
-              }
-              if (promptCalls > 0) {
-                const durationMs = Date.now() - promptStartedAt;
-                console.log(
-                  `[ai] prompt done — ${promptCalls} calls · ${(durationMs / 1000).toFixed(1)}s · out=${promptOutTok} tok · peak ctx=${promptPeakCtx} tok`
-                );
-                // One queryable wide event per prompt in the NDJSON app log.
-                const wide: Record<string, unknown> = {
-                  channel: "prompt",
-                  model: sessionModel,
-                  session: conversationId,
-                  calls: promptCalls,
-                  retries: promptRetries,
-                  durationMs,
-                  outputTokens: promptOutTok,
-                  peakContext: promptPeakCtx,
-                };
-                if (promptError) {
-                  wide.error = promptError;
-                  log.error(wide);
-                } else {
-                  log.info(wide);
-                }
-                promptCalls = 0;
-                promptOutTok = 0;
-                promptPeakCtx = 0;
-                promptRetries = 0;
-              }
-              send({ type: "done" });
-            }
-          });
-
-          // Append a live browser-state notice (if a browser is open) so the
-          // agent knows a session is running and what it shows. Best-effort.
-          const browserNotice = await browserStateNotice(sessionCwd);
-
-          // User-added context that changed since the system prompt was built.
-          let contextNotice: string | null = null;
-          const currentContextStamp = contextStamp(sessionCwd);
-          if (currentContextStamp !== contextSeenStamp) {
-            contextSeenStamp = currentContextStamp;
-            contextNotice = contextUpdateNotice(
-              listContext(sessionCwd),
-              listContextFolders(sessionCwd)
-            );
-          }
-
-          // The client serializes the widget the user is currently viewing (id +
-          // its declared actions) into `<active_widget>`; append it per-message
-          // so the agent can drive it via the `ui_widget` tool.
-          const widgetNotice =
-            typeof msg.activeWidget === "string" && msg.activeWidget.trim()
-              ? msg.activeWidget
-              : null;
-
-          // A leading `/<skill-name>` expands the skill's content into the
-          // prompt (the model receives the workflow, the UI keeps showing the
-          // short command). Announce it so the UI can render a "using skill"
-          // banner on the agent's reply.
-          const skills = getSkills(sessionCwd);
-          const skill = matchSkillCommand(msg.message, skills);
-          if (skill) {
-            send({
-              type: "skill",
-              name: skill.entry.name,
-              description: skill.entry.description ?? "",
-            });
-            await session.prompt(
-              withActiveWidget(
-                widgetNotice,
-                withBrowserState(
-                  browserNotice,
-                  withContextUpdate(contextNotice, buildSkillPrompt(skill.entry, skill.args))
-                )
-              )
-            );
-            break;
-          }
-
-          // `/<skill-name>` mentioned anywhere else in the message keeps the
-          // user's text intact and appends each mentioned skill's instructions.
-          const mentioned = matchSkillMentions(msg.message, skills);
-          for (const entry of mentioned) {
-            send({
-              type: "skill",
-              name: entry.name,
-              description: entry.description ?? "",
-            });
-          }
-
-          // Images go to the model natively; other files are written into the
-          // workspace so the agent can read them with its file tools.
-          const { text, images } = prepareAttachments(
-            msg.message,
-            msg.files,
-            sessionCwd,
-            messageId
-          );
-          await session.prompt(
-            withActiveWidget(
-              widgetNotice,
-              withBrowserState(
-                browserNotice,
-                withContextUpdate(contextNotice, withSkillMentions(text, mentioned))
-              )
-            ),
-            images.length > 0 ? { images } : undefined
-          );
-        } catch (err: any) {
-          console.error("[session] prompt error:", err?.message || err);
-          send({ type: "error", error: err.message || String(err) });
-        }
-        break;
-      }
-
-      case "answer": {
-        if (askChannel && typeof msg.toolCallId === "string") {
-          askChannel.answer(msg.toolCallId, String(msg.value ?? ""));
-        }
-        break;
-      }
-
-      case "widget_result": {
-        if (widgetChannel && typeof msg.toolCallId === "string") {
-          const value =
-            typeof msg.result === "string" ? msg.result : JSON.stringify(msg.result ?? "");
-          widgetChannel.respond(msg.toolCallId, value);
-        }
-        break;
-      }
-
-      case "abort": {
-        if (session) {
-          try {
-            askChannel?.cancelAll();
-            widgetChannel?.cancelAll();
-            await session.abort();
-            send({ type: "abort" });
-          } catch (err: any) {
-            send({ type: "error", error: err.message || String(err) });
-          }
-        }
-        break;
-      }
-
-      case "new_session": {
-        unsubscribe?.();
-        aiDebugUnsub?.();
-        aiDebugUnsub = null;
-        askChannel?.cancelAll();
-        widgetChannel?.cancelAll();
-        session = null;
-        askChannel = null;
-        widgetChannel = null;
-        resumeConversationId = null;
-        send({ type: "session_cleared" });
-        break;
-      }
-
-      case "open_conversation": {
-        unsubscribe?.();
-        aiDebugUnsub?.();
-        aiDebugUnsub = null;
-        askChannel?.cancelAll();
-        widgetChannel?.cancelAll();
-        session = null;
-        askChannel = null;
-        widgetChannel = null;
-        resumeConversationId =
-          typeof msg.conversationId === "string" ? msg.conversationId : null;
-        send({ type: "conversation_opened", conversationId: resumeConversationId });
-        break;
-      }
-
-      default:
-        send({ type: "error", error: `Unknown message type: ${msg.type}` });
-    }
-  }
-
-  return {
-    async handleRaw(raw: string): Promise<void> {
-      let msg: any;
-      try {
-        msg = JSON.parse(raw);
-      } catch {
-        send({ type: "error", error: "Invalid JSON" });
-        return;
-      }
-      await handleMessage(msg);
-    },
-    close(): void {
-      clearInterval(heartbeat);
-      for (const timer of partialTimers.values()) clearTimeout(timer);
-      partialTimers.clear();
-      pendingPartials.clear();
-      askChannel?.cancelAll();
-      widgetChannel?.cancelAll();
-      unsubscribe?.();
-      aiDebugUnsub?.();
-    },
-  };
-}
-
-/**
- * Split prompt attachments into model-native images and written-to-disk files.
- * Each attachment is a base64 data URL. Images become `ImageContent` passed to
- * the model directly; everything else is written under `<cwd>/.testeiya/uploads/
- * <messageId>/` and referenced by relative path so the agent can `read` it.
- */
-function prepareAttachments(
-  message: string,
-  files: PromptFile[] | undefined,
-  cwd: string,
-  messageId: string
-): { text: string; images: ImageContent[] } {
-  const images: ImageContent[] = [];
-  if (!files || files.length === 0) return { text: message, images };
-
-  const savedPaths: string[] = [];
-  let uploadDir = "";
-  for (const file of files) {
-    const match = /^data:([^;]+);base64,(.*)$/s.exec(file.dataUrl);
-    if (!match) continue;
-    const mimeType = file.mediaType || match[1];
-    const base64 = match[2];
-    if (mimeType.startsWith("image/")) {
-      images.push({ type: "image", data: base64, mimeType });
-      continue;
-    }
-    if (!uploadDir) {
-      uploadDir = join(cwd, PROJECT_DIR, "uploads", messageId);
-      mkdirSync(uploadDir, { recursive: true });
-    }
-    const safeName = file.filename.replace(/[/\\]/g, "_");
-    writeFileSync(join(uploadDir, safeName), Buffer.from(base64, "base64"));
-    savedPaths.push(join(PROJECT_DIR, "uploads", messageId, safeName));
-  }
-
-  if (savedPaths.length === 0) {
-    // A bare image with no text leaves an empty text block, which some
-    // providers reject — give it a minimal caption.
-    if (images.length > 0 && !message.trim()) {
-      return { text: "(see attached image)", images };
-    }
-    return { text: message, images };
-  }
-
-  const note = savedPaths.map((p) => `- ${p}`).join("\n");
-  const base = message.trim() ? message : "(see attached files)";
-  const text = `${base}\n\nAttached files saved into the workspace (read them as needed):\n${note}`;
-  return { text, images };
-}
-
-type SkillEntry = ReturnType<typeof loadBundledSkills>[number];
-
-/** Match a leading `/<skill-name>` against the loaded skills; null if none. */
-function matchSkillCommand(
-  message: string,
-  skills: SkillEntry[]
-): { entry: SkillEntry; args: string } | null {
-  const text = (message ?? "").trimStart();
-  if (!text.startsWith("/")) return null;
-  const match = /^\/([^\s]+)\s*([\s\S]*)$/.exec(text);
-  if (!match) return null;
-  const name = match[1].toLowerCase();
-  const entry = skills.find((s) => s.name.toLowerCase() === name);
-  if (!entry) return null;
-  return { entry, args: match[2].trim() };
-}
-
-/** Collect every `/<skill-name>` at a word boundary anywhere in the message. */
-function matchSkillMentions(message: string, skills: SkillEntry[]): SkillEntry[] {
-  const found = new Map<string, SkillEntry>();
-  for (const match of (message ?? "").matchAll(/(?:^|\s)\/([a-z0-9][a-z0-9-]*)/gi)) {
-    const name = match[1].toLowerCase();
-    if (found.has(name)) continue;
-    const entry = skills.find((s) => s.name.toLowerCase() === name);
-    if (entry) found.set(name, entry);
-  }
-  return [...found.values()];
-}
-
-/** Append each mentioned skill's instructions after the user's message. */
-function withSkillMentions(text: string, entries: SkillEntry[]): string {
-  if (entries.length === 0) return text;
-  const blocks = entries.map(
-    (entry) => `<skill name="${entry.name}">\n${readSkillBody(entry)}\n</skill>`
-  );
-  return `${text}\n\nThe /<name> tokens in the message above reference skills. Follow each referenced skill's instructions for this request:\n\n${blocks.join("\n\n")}`;
-}
-
-/** Append the `<browser_state>` notice (when present) to a prompt. */
-function withBrowserState(notice: string | null, text: string): string {
-  if (!notice) return text;
-  return `${text}\n\n${notice}`;
-}
-
-/** Append the `<active_widget>` notice (when present) to a prompt. */
-function withActiveWidget(notice: string | null, text: string): string {
-  if (!notice) return text;
-  return `${text}\n\n${notice}`;
-}
-
-/** Append the `<workspace-context-update>` notice (when present) to a prompt. */
-function withContextUpdate(notice: string | null, text: string): string {
-  if (!notice) return text;
-  return `${text}\n\n${notice}`;
-}
-
-/** Expand a skill into a prompt: its workflow content plus any user args. */
-function buildSkillPrompt(entry: SkillEntry, args: string): string {
-  const parts = [readSkillBody(entry)];
-  if (args) parts.push(`---\n\nUser request: ${args}`);
-  return parts.join("\n\n");
-}
-
-/** The skill's SKILL.md content with the YAML frontmatter stripped. */
-function readSkillBody(entry: SkillEntry): string {
-  let content = "";
-  try {
-    content = readFileSync(entry.filePath, "utf8");
-  } catch {
-    content = "";
-  }
-  const body = content.replace(/^---\n[\s\S]*?\n---\n?/, "").trim();
-  return body || `Apply the "${entry.name}" skill.`;
-}
-
-interface PromptFile {
-  filename: string;
-  mediaType: string;
-  dataUrl: string;
-}
-
-interface ImageContent {
-  type: "image";
-  data: string;
-  mimeType: string;
-}