@yeaft/webchat-agent 1.0.188 → 1.0.190

package/local-runtime/server/ws-agent.js

@@ -0,0 +1,319 @@
+import { randomUUID } from 'crypto';
+import { WebSocket } from 'ws';
+import { CONFIG } from './config.js';
+import { verifyAgent } from './auth.js';
+import { encodeKey } from './encryption.js';
+import { agents, pendingAgentConnections } from './context.js';
+import {
+  parseMessage, broadcastAgentList, clearAgentDirCache
+} from './ws-utils.js';
+import { handleAgentConversation } from './handlers/agent-conversation.js';
+import { handleAgentOutput } from './handlers/agent-output.js';
+import { handleAgentWorkCenter } from './handlers/agent-work-center.js';
+import { handleAgentFileTerminal } from './handlers/agent-file-terminal.js';
+import { handleAgentSync } from './handlers/agent-sync.js';
+import { recordPerfTraceEvent } from './perf-trace.js';
+import { markAgentHeartbeatSeen } from './heartbeat-policy.js';
+
+/**
+ * Build the internal Map key for an agent.
+ * Uses `${ownerId}:${agentName}` to prevent different users' same-named
+ * agents from colliding. Global (AGENT_SECRET) connections use `global:`.
+ */
+function buildAgentMapKey(ownerId, agentName) {
+  const prefix = ownerId || 'global';
+  return `${prefix}:${agentName}`;
+}
+
+export function handleAgentConnection(ws, url) {
+  const clientAgentId = url.searchParams.get('id') || randomUUID();
+  const agentName = url.searchParams.get('name') || `Agent-${clientAgentId.slice(0, 8)}`;
+  const instanceId = url.searchParams.get('instanceId') || clientAgentId;
+  const workDir = url.searchParams.get('workDir') || '';
+
+  // In development mode (SKIP_AUTH), register immediately
+  if (CONFIG.skipAuth) {
+    // Dev mode: no owner isolation, use clientAgentId as-is for backward compat
+    const capabilities = (url.searchParams.get('capabilities') || '').split(',').filter(Boolean);
+    completeAgentRegistration(ws, clientAgentId, agentName, workDir, null, capabilities, null, null, null, instanceId);
+
+    ws.on('message', async (data) => {
+      const agent = agents.get(clientAgentId);
+      if (!agent) {
+        console.error(`[Agent] No agent found for id: ${clientAgentId}`);
+        return;
+      }
+      markAgentHeartbeatSeen(agent);
+      const msg = await parseMessage(data, agent.sessionKey);
+      if (msg) {
+        console.log(`[Agent] Received message from ${clientAgentId}: ${msg.type}`);
+        if (msg.perfTraceId) {
+          recordPerfTraceEvent({
+            traceId: msg.perfTraceId,
+            source: 'server',
+            phase: 'websocket.agent_received',
+            at: Date.now(),
+            agentId: clientAgentId,
+            sessionId: msg.sessionId || null,
+            vpId: msg.vpId || null,
+            turnId: msg.turnId || null,
+            threadId: msg.threadId || null,
+            messageType: msg.type,
+            bytes: data.length || 0,
+          });
+        }
+        handleAgentMessage(clientAgentId, msg);
+      } else {
+        console.error(`[Agent] Failed to parse message from ${clientAgentId}`);
+      }
+    });
+
+    ws.on('close', () => {
+      handleAgentDisconnect(clientAgentId, agentName);
+    });
+
+    ws.on('error', (err) => {
+      console.error(`Agent error (${agentName}):`, err.message);
+    });
+    return;
+  }
+
+  // In production mode, wait for auth message with secret
+  const tempId = randomUUID();
+  // Mutable: will be updated to the owner-scoped key after auth succeeds
+  let resolvedAgentId = null;
+
+  const authTimeout = setTimeout(() => {
+    console.log(`Agent auth timeout: ${agentName}`);
+    pendingAgentConnections.delete(tempId);
+    ws.close(1008, 'Authentication timeout');
+  }, 30000);
+
+  pendingAgentConnections.set(tempId, {
+    ws,
+    agentId: clientAgentId,
+    agentName,
+    instanceId,
+    workDir,
+    timeout: authTimeout
+  });
+
+  // Request authentication
+  ws.send(JSON.stringify({
+    type: 'auth_required',
+    tempId
+  }));
+
+  ws.on('message', async (data) => {
+    const pending = pendingAgentConnections.get(tempId);
+    if (pending) {
+      // Still pending authentication
+      try {
+        const msg = JSON.parse(data.toString());
+        if (msg.type === 'auth' && msg.tempId === tempId) {
+          clearTimeout(pending.timeout);
+          pendingAgentConnections.delete(tempId);
+
+          const authResult = verifyAgent(msg.secret);
+          if (!authResult.valid) {
+            console.log(`Agent auth failed: ${agentName}`);
+            ws.close(1008, 'Invalid agent secret');
+            return;
+          }
+
+          const capabilities = msg.capabilities || [];
+          const agentVersion = msg.version || null;
+          // Build owner-scoped key to prevent cross-user collision. New agents
+          // identify local service instances separately from display names;
+          // old agents keep using their name as the stable id.
+          resolvedAgentId = buildAgentMapKey(authResult.userId, pending.instanceId || pending.agentId || pending.agentName);
+          completeAgentRegistration(ws, resolvedAgentId, pending.agentName, pending.workDir, authResult.sessionKey, capabilities, authResult.userId, authResult.username, agentVersion, pending.instanceId || pending.agentId || pending.agentName);
+        }
+      } catch (e) {
+        console.error('Failed to parse agent auth message:', e.message);
+      }
+    } else {
+      // Already authenticated, handle normally
+      if (!resolvedAgentId) return;
+      const agent = agents.get(resolvedAgentId);
+      if (!agent) {
+        console.error(`[Agent] No agent found for id: ${resolvedAgentId}`);
+        return;
+      }
+      markAgentHeartbeatSeen(agent);
+      const msg = await parseMessage(data, agent.sessionKey);
+      if (msg) {
+        console.log(`[Agent] Received message from ${resolvedAgentId}: ${msg.type}`);
+        if (msg.perfTraceId) {
+          recordPerfTraceEvent({
+            traceId: msg.perfTraceId,
+            source: 'server',
+            phase: 'websocket.agent_received',
+            at: Date.now(),
+            agentId: resolvedAgentId,
+            sessionId: msg.sessionId || null,
+            vpId: msg.vpId || null,
+            turnId: msg.turnId || null,
+            threadId: msg.threadId || null,
+            messageType: msg.type,
+            bytes: data.length || 0,
+          });
+        }
+        handleAgentMessage(resolvedAgentId, msg);
+      } else {
+        console.error(`[Agent] Failed to parse message from ${resolvedAgentId}`);
+      }
+    }
+  });
+
+  ws.on('close', () => {
+    const pending = pendingAgentConnections.get(tempId);
+    if (pending) {
+      clearTimeout(pending.timeout);
+      pendingAgentConnections.delete(tempId);
+    }
+    // Use resolvedAgentId if auth completed, otherwise nothing to clean
+    if (resolvedAgentId) {
+      handleAgentDisconnect(resolvedAgentId, agentName);
+    }
+  });
+
+  ws.on('error', (err) => {
+    console.error(`Agent error (${agentName}):`, err.message);
+  });
+}
+
+/**
+ * Shared disconnect handler: clean up and remove agent from the agents Map.
+ * Conversations are persisted in DB and will be restored on reconnect via
+ * get_agents (client-side recovery) and conversation_list (agent-side sync).
+ */
+function handleAgentDisconnect(agentId, agentName) {
+  const agent = agents.get(agentId);
+  // Phase 4: 清理目录缓存
+  clearAgentDirCache(agentId);
+  // Phase 1: 清理同步超时
+  if (agent?._syncTimeout) {
+    clearTimeout(agent._syncTimeout);
+  }
+  // Remove agent entirely — eliminates zombie agents from broadcastAgentList
+  agents.delete(agentId);
+  console.log(`Agent disconnected: ${agentName}`);
+  broadcastAgentList();
+}
+
+function completeAgentRegistration(ws, agentId, agentName, workDir, sessionKey, capabilities = [], ownerId = null, ownerUsername = null, agentVersion = null, instanceId = null) {
+  // 如果是重连，保留 conversations；否则（server 重启）创建空 Map
+  const existingAgent = agents.get(agentId);
+  const conversations = existingAgent?.conversations || new Map();
+  const proxyPorts = (existingAgent?.proxyPorts || []).map(p => ({ ...p, enabled: false }));
+  const slashCommands = existingAgent?.slashCommands || [];
+  const slashCommandDescriptions = existingAgent?.slashCommandDescriptions || {};
+
+  // 兼容旧版 agent：未上报 capabilities 时默认全部开启
+  const effectiveCapabilities = capabilities.length > 0
+    ? capabilities
+    : ['terminal', 'file_editor', 'background_tasks'];
+
+  // feat-ws-plaintext-negotiation: new agents advertise `plaintext-ok` in
+  // their capability list (either via the URL query string in skipAuth
+  // path or in the `auth` frame in prod path). When absent, treat as old
+  // agent and keep encrypting outbound (back-compat).
+  const encryptOutbound = !effectiveCapabilities.includes('plaintext-ok');
+
+  agents.set(agentId, {
+    ws,
+    name: agentName,
+    instanceId,
+    workDir,
+    conversations,
+    sessionKey,
+    isAlive: true,
+    lastSeenAt: Date.now(),
+    capabilities: effectiveCapabilities,
+    proxyPorts,
+    slashCommands,
+    slashCommandDescriptions,
+    status: 'syncing',
+    ownerId,
+    ownerUsername,
+    version: agentVersion,
+    encryptOutbound
+  });
+
+  // 同步超时保护：30 秒后强制 ready
+  const syncTimeout = setTimeout(() => {
+    const ag = agents.get(agentId);
+    if (ag && ag.status === 'syncing') {
+      console.warn(`[Sync] Agent ${agentName} sync timeout, forcing ready`);
+      ag.status = 'ready';
+      broadcastAgentList();
+    }
+  }, 30000);
+  agents.get(agentId)._syncTimeout = syncTimeout;
+
+  // 心跳响应处理 + latency 测量
+  ws.on('pong', () => {
+    const agent = agents.get(agentId);
+    if (agent) {
+      markAgentHeartbeatSeen(agent);
+      if (agent.pingSentAt) {
+        agent.latency = Date.now() - agent.pingSentAt;
+        agent.pingSentAt = null;
+      }
+    }
+  });
+
+  // Send registration (with session key only in production mode)
+  const latestAgentVersion = process.env.AGENT_LATEST_VERSION || null;
+  const upgradeAvailable = (latestAgentVersion && agentVersion && latestAgentVersion !== agentVersion) ? latestAgentVersion : null;
+
+  ws.send(JSON.stringify({
+    type: 'registered',
+    agentId,
+    sessionKey: sessionKey ? encodeKey(sessionKey) : null,
+    // feat-ws-plaintext-negotiation: tell new agents that this server
+    // will accept plaintext outbound from them. Old agents ignore the
+    // unknown field. New agents flip `serverEncryptionRequired = false`
+    // and stop calling encrypt() on the send path.
+    acceptPlaintext: true,
+    ...(upgradeAvailable && { upgradeAvailable })
+  }));
+
+  console.log(`Agent connected: ${agentName} (${agentId})${encryptOutbound ? '' : ' [plaintext mode]'}`);
+  broadcastAgentList();
+}
+
+async function handleAgentMessage(agentId, msg) {
+  const agent = agents.get(agentId);
+  if (!agent) return;
+
+  // Security: 需要 conversationId 的消息类型，验证该 conversation 属于此 agent.
+  // The conversation-id check only authorizes Chat flows where the
+  // conversationId IS the ownership key. Workbench responses and Yeaft
+  // events use `_requestUserId` for ownership (enforced in
+  // `forwardToClients`), so they bypass this gate. See PR #772.
+  const CONV_EXEMPT_TYPES = new Set([
+    'conversation_list', 'conversation_created', 'conversation_resumed',
+    'agent_sync_complete', 'sync_sessions', 'proxy_response', 'proxy_response_chunk',
+    'proxy_response_end', 'proxy_ports_update', 'proxy_ws_opened', 'proxy_ws_message',
+    'proxy_ws_closed', 'proxy_ws_error', 'restart_agent_ack', 'upgrade_agent_ack',
+    'directory_listing', 'folders_list', 'models_list', 'yeaft_output', 'yeaft_session_output', 'session_output', 'yeaft_asset_put',
+    'yeaft_history_chunk', 'yeaft_history_search_result', 'yeaft_history_window', 'slash_commands_update', 'agent_metrics',
+    'file_content', 'file_saved', 'file_op_result', 'file_search_result',
+    'git_status_result', 'git_diff_result', 'git_op_result'
+  ]);
+  if (msg.conversationId && !CONV_EXEMPT_TYPES.has(msg.type)) {
+    if (!agent.conversations.has(msg.conversationId)) {
+      console.warn(`[Security] Agent ${agentId} sent ${msg.type} for unknown conversation ${msg.conversationId}, ignoring`);
+      return;
+    }
+  }
+
+  // Dispatch to handler sub-modules
+  if (await handleAgentConversation(agentId, agent, msg)) return;
+  if (await handleAgentWorkCenter(agentId, msg)) return;
+  if (await handleAgentOutput(agentId, agent, msg)) return;
+  if (await handleAgentFileTerminal(agentId, agent, msg)) return;
+  if (await handleAgentSync(agentId, agent, msg)) return;
+}

package/local-runtime/server/ws-client.js

@@ -0,0 +1,214 @@
+import { randomUUID } from 'crypto';
+import { WebSocket } from 'ws';
+import { CONFIG } from './config.js';
+import { verifyToken, generateSkipAuthSession } from './auth.js';
+import { encodeKey } from './encryption.js';
+import { userDb } from './database.js';
+import { agents, webClients, isHeartbeatMessageType, trackRequest } from './context.js';
+import {
+  parseMessage, sendToWebClient, sendToAgent,
+  broadcastAgentList, resolveAgentAccessError
+} from './ws-utils.js';
+import { handleClientConversation } from './handlers/client-conversation.js';
+import { handleClientWorkbench } from './handlers/client-workbench.js';
+import { handleClientMisc } from './handlers/client-misc.js';
+import { clearWorkCenterRequestsForClient, handleClientWorkCenter } from './handlers/client-work-center.js';
+import { recordPerfTraceEvent } from './perf-trace.js';
+
+export function handleWebConnection(ws, url) {
+  const clientId = randomUUID();
+  const token = url.searchParams.get('token');
+
+  let authenticated = false;
+  let sessionKey = null;
+  let username = null;
+  let userId = null;
+  let role = null;
+
+  // Check for skip auth mode
+  if (CONFIG.skipAuth) {
+    authenticated = true;
+    const session = generateSkipAuthSession();
+    sessionKey = session.sessionKey;
+    username = 'dev-user';
+    role = 'admin';
+  } else if (token) {
+    const result = verifyToken(token);
+    if (result.valid) {
+      authenticated = true;
+      sessionKey = result.sessionKey;
+      username = result.username;
+      role = result.role === 'admin' ? 'admin' : 'pro';
+    }
+  }
+
+  // 获取或创建用户
+  if (authenticated && username) {
+    try {
+      const user = userDb.getOrCreate(username);
+      userId = user.id;
+      userDb.updateLogin(userId);
+    } catch (e) {
+      console.error('Failed to get/create user:', e.message);
+    }
+  }
+
+  webClients.set(clientId, {
+    ws,
+    authenticated,
+    username,
+    userId,
+    role,
+    currentAgent: null,
+    currentConversation: null,
+    sessionKey,
+    isAlive: true,
+    // feat-ws-plaintext-negotiation: per-client flag. Defaults `true`
+    // (= old client, encrypt outbound for back-compat). Flipped to
+    // `false` when the client sends `client_hello { plaintextOk: true }`
+    // — see early dispatch in handleWebMessage.
+    encryptOutbound: true
+  });
+
+  // 心跳响应处理
+  ws.on('pong', () => {
+    const client = webClients.get(clientId);
+    if (client) client.isAlive = true;
+  });
+
+  console.log(`Web client connected: ${clientId} (authenticated: ${authenticated})`);
+
+  const client = webClients.get(clientId);
+
+  if (authenticated) {
+    // Send auth result unencrypted (client doesn't have key yet).
+    // `acceptPlaintext: true` advertises that this server will accept
+    // plaintext from a new client. Old clients ignore the unknown field.
+    // The corresponding flip on the server side (stop encrypting outbound
+    // to this client) happens when the new client confirms it can speak
+    // plaintext via the `client_hello` frame — see handleWebMessage.
+    ws.send(JSON.stringify({
+      type: 'auth_result',
+      success: true,
+      sessionKey: sessionKey ? encodeKey(sessionKey) : null,
+      role,
+      acceptPlaintext: true
+    }));
+    setTimeout(() => broadcastAgentList(), 100);
+  } else {
+    ws.send(JSON.stringify({ type: 'auth_result', success: false, error: 'Authentication required' }));
+    ws.close(1008, 'Authentication required');
+    return;
+  }
+
+  ws.on('message', async (data) => {
+    const client = webClients.get(clientId);
+    const msg = await parseMessage(data, client?.sessionKey);
+    if (!msg) return;
+    // Stats tracking: exclude heartbeat/control frames. User-turn bytes are
+    // accounted at the send handlers, not here, so pings and dashboard polling
+    // don't inflate traffic.
+    if (!isHeartbeatMessageType(msg.type)) {
+      trackRequest(client?.userId, data.length || 0, msg.type);
+    }
+    if (msg.perfTraceId) {
+      recordPerfTraceEvent({
+        traceId: msg.perfTraceId,
+        source: 'server',
+        phase: 'websocket.web_received',
+        at: Date.now(),
+        userId: client?.userId || null,
+        agentId: msg.agentId || client?.currentAgent || null,
+        sessionId: msg.sessionId || null,
+        messageType: msg.type,
+        bytes: data.length || 0,
+      });
+    }
+    handleWebMessage(clientId, msg);
+  });
+
+  ws.on('close', () => {
+    const client = webClients.get(clientId);
+    // Web 客户端断开时，检查是否需要禁用其关联 Agent 的端口
+    if (client?.currentAgent) {
+      const agentId = client.currentAgent;
+      const agent = agents.get(agentId);
+      if (agent?.proxyPorts?.length > 0) {
+        // 检查是否还有其他 Web 客户端连接到同一个 agent
+        let otherClientsOnAgent = false;
+        for (const [otherId, otherClient] of webClients.entries()) {
+          if (otherId !== clientId && otherClient.currentAgent === agentId && otherClient.ws.readyState === WebSocket.OPEN) {
+            otherClientsOnAgent = true;
+            break;
+          }
+        }
+        // 只有当没有其他 Web 客户端连接到同一 agent 时才禁用
+        if (!otherClientsOnAgent) {
+          agent.proxyPorts = agent.proxyPorts.map(p => ({ ...p, enabled: false }));
+          if (agent.ws.readyState === WebSocket.OPEN) {
+            sendToAgent(agent, { type: 'proxy_update_ports', ports: agent.proxyPorts });
+          }
+          broadcastAgentList();
+        }
+      }
+    }
+    clearWorkCenterRequestsForClient(client);
+    webClients.delete(clientId);
+    console.log(`Web client disconnected: ${clientId}`);
+  });
+
+  ws.on('error', (err) => {
+    console.error(`Web client error (${clientId}):`, err.message);
+  });
+}
+
+// Workbench 功能（terminal、file、git、proxy）仅 admin/pro 可用
+const WORKBENCH_TYPES = new Set([
+  'terminal_create', 'terminal_input', 'terminal_resize', 'terminal_close',
+  'read_file', 'write_file', 'create_file', 'delete_files', 'move_files', 'copy_files', 'upload_to_dir', 'file_search',
+  'git_status', 'git_diff', 'git_add', 'git_reset', 'git_restore', 'git_commit', 'git_push',
+  'proxy_update_ports', 'update_file_tabs', 'restore_file_tabs'
+]);
+
+async function handleWebMessage(clientId, msg) {
+  const client = webClients.get(clientId);
+  if (!client || !client.authenticated) return;
+
+  // feat-ws-plaintext-negotiation: early capability frame from new web
+  // clients. Tells the server "you may stop encrypting outbound to me".
+  // Old clients never send this; their per-client `encryptOutbound` flag
+  // stays `true` and we keep the ciphertext path.
+  if (msg.type === 'client_hello') {
+    if (msg.plaintextOk === true) {
+      client.encryptOutbound = false;
+      console.log(`[WS] Client ${clientId} negotiated plaintext mode`);
+    }
+    return;
+  }
+
+  // Workbench 权限检查：仅 admin 和 pro 可用（当前所有用户都是 pro 或 admin）
+  if (!CONFIG.skipAuth && WORKBENCH_TYPES.has(msg.type) && client.role !== 'admin' && client.role !== 'pro') {
+    console.warn(`[Security] User ${client.userId} (role=${client.role}) denied workbench action: ${msg.type}`);
+    await sendToWebClient(client, { type: 'error', message: 'Permission denied: workbench access requires pro or admin role' });
+    return;
+  }
+
+  // Helper: check agent access (ownership + availability)
+  const checkAgentAccess = async (agentId) => {
+    const error = resolveAgentAccessError(agentId, client.userId, client.role);
+    if (!error) return true;
+    if (error === 'Agent access denied') {
+      console.warn(`[Security] User ${client.userId} denied access to agent ${agentId}`);
+    } else {
+      console.warn(`[WS] Agent unavailable for user ${client.userId}: ${agentId || '(none)'}`);
+    }
+    await sendToWebClient(client, { type: 'error', message: error });
+    return false;
+  };
+
+  // Dispatch to handler sub-modules
+  if (await handleClientConversation(clientId, client, msg, checkAgentAccess)) return;
+  if (await handleClientWorkbench(clientId, client, msg, checkAgentAccess)) return;
+  if (await handleClientWorkCenter(client, msg, checkAgentAccess)) return;
+  if (await handleClientMisc(clientId, client, msg, checkAgentAccess)) return;
+}