neoagent 3.2.1-beta.2 → 3.2.1-beta.4

package/server/services/integrations/manager.js

@@ -10,6 +10,11 @@ const {
   normalizeAccessMode,
   withConnectionAccessMode,
 } = require('./access');
+const {
+  abortError,
+  waitForAbortableResult,
+} = require('./http');
+const { throwIfAborted } = require('../../utils/abort');
 
 const OAUTH_STATE_PATTERN = /^[a-f0-9]{32,128}$/i;
 
@@ -56,6 +61,7 @@ class IntegrationManager {
     this.registry = createIntegrationRegistry({
       io: this.app?.locals?.io || null,
     });
+    this.connectionExecutionQueues = new Map();
   }
 
   parseCredentials(credentialsJson) {
@@ -152,6 +158,23 @@ class IntegrationManager {
     ).run();
   }
 
+  claimOAuthState(state) {
+    return db.transaction(() => {
+      const stateRow = db
+        .prepare(
+          `SELECT * FROM integration_oauth_states
+           WHERE state = ? AND datetime(expires_at) > datetime('now')`,
+        )
+        .get(state);
+      if (!stateRow) return null;
+      const deleted = db.prepare(
+        `DELETE FROM integration_oauth_states
+         WHERE state = ? AND datetime(expires_at) > datetime('now')`,
+      ).run(state);
+      return deleted.changes === 1 ? stateRow : null;
+    })();
+  }
+
   listConnections(userId, providerKey = null, agentId = null) {
     const scopedAgentId = resolveAgentId(userId, agentId);
     const query = providerKey
@@ -225,6 +248,7 @@ class IntegrationManager {
         userId,
         agentId,
         appKey,
+        signal: options.signal || null,
       });
       const url = String(result?.url || '').trim();
       const absoluteUrl = url.startsWith('http')
@@ -247,6 +271,7 @@ class IntegrationManager {
       userId,
       agentId,
       appKey,
+      signal: options.signal || null,
     });
 
     db.prepare(
@@ -290,7 +315,8 @@ class IntegrationManager {
     );
   }
 
-  async finishOAuth(state, code) {
+  async finishOAuth(state, code, options = {}) {
+    throwIfAborted(options.signal, 'OAuth callback request aborted.');
     this.cleanupExpiredOauthStates();
     const normalizedState = String(state || '').trim();
     if (!OAUTH_STATE_PATTERN.test(normalizedState)) {
@@ -300,12 +326,7 @@ class IntegrationManager {
     if (!normalizedCode) {
       throw new Error('OAuth authorization code is required.');
     }
-    const stateRow = db
-      .prepare(
-        `SELECT * FROM integration_oauth_states
-         WHERE state = ? AND datetime(expires_at) > datetime('now')`,
-      )
-      .get(normalizedState);
+    const stateRow = this.claimOAuthState(normalizedState);
     if (!stateRow) {
       throw new Error('OAuth state is missing or expired.');
     }
@@ -322,6 +343,7 @@ class IntegrationManager {
       code: normalizedCode,
       codeVerifier: decryptValue(stateRow.code_verifier),
       appKey: stateRow.app_key,
+      signal: options.signal || null,
     });
 
     const mergedCredentials = this.mergeWithReusableCredentials(
@@ -386,10 +408,6 @@ class IntegrationManager {
         result.accountEmail,
       );
 
-    db.prepare('DELETE FROM integration_oauth_states WHERE state = ?').run(
-      stateRow.state,
-    );
-
     return {
       provider: provider.key,
       appId: stateRow.app_key,
@@ -420,7 +438,11 @@ class IntegrationManager {
       };
     }
 
-    await provider.disconnect(connection).catch(() => {});
+    if (typeof provider.disconnect === 'function') {
+      await provider.disconnect(connection, {
+        signal: options.signal || null,
+      }).catch(() => {});
+    }
     db.prepare('DELETE FROM integration_connections WHERE user_id = ? AND agent_id = ? AND id = ?').run(
       userId,
       agentId,
@@ -636,7 +658,46 @@ class IntegrationManager {
     };
   }
 
-  async executeTool(userId, toolName, args, agentId = null) {
+  connectionExecutionKey(connection) {
+    return [
+      connection.user_id,
+      connection.agent_id,
+      connection.provider_key,
+      String(connection.account_email || '').trim().toLowerCase(),
+    ].join(':');
+  }
+
+  async acquireConnectionExecution(connection, signal) {
+    const key = this.connectionExecutionKey(connection);
+    const previous = this.connectionExecutionQueues.get(key) || Promise.resolve();
+    let releaseGate;
+    const gate = new Promise((resolve) => {
+      releaseGate = resolve;
+    });
+    const tail = previous.catch(() => {}).then(() => gate);
+    this.connectionExecutionQueues.set(key, tail);
+    void tail.finally(() => {
+      if (this.connectionExecutionQueues.get(key) === tail) {
+        this.connectionExecutionQueues.delete(key);
+      }
+    });
+
+    try {
+      await waitForAbortableResult(previous.catch(() => {}), signal);
+    } catch (error) {
+      releaseGate();
+      throw error;
+    }
+
+    let released = false;
+    return () => {
+      if (released) return;
+      released = true;
+      releaseGate();
+    };
+  }
+
+  async executeTool(userId, toolName, args, agentId = null, options = {}) {
     let foundSupportingProvider = false;
     for (const provider of this.registry.list()) {
       if (!provider.supportsTool(toolName)) continue;
@@ -661,54 +722,86 @@ class IntegrationManager {
         };
       }
 
-      let execution;
+      const storedCredentials = this.parseCredentials(
+        selection.connection.credentials_json,
+      );
+      const needsCredentialLock =
+        provider.requiresRefreshToken === true ||
+        Boolean(storedCredentials.refresh_token);
+      const releaseExecution = needsCredentialLock
+        ? await this.acquireConnectionExecution(selection.connection, options.signal)
+        : () => {};
       try {
-        execution = await provider.executeTool(
-          toolName,
-          args,
-          selection.connection,
+        const connection = this.getConnectionById(
+          userId,
+          selection.connection.id,
+          agentId,
         );
-      } catch (err) {
-        if (isLikelyExpiredConnectionError(err)) {
-          db.prepare(
-            `UPDATE integration_connections
-             SET status = 'expired', updated_at = datetime('now')
-             WHERE id = ? AND user_id = ? AND agent_id = ?`,
-          ).run(
-            selection.connection.id,
+        if (!connection || connection.status !== 'connected') {
+          return {
+            error: `The selected ${provider.label} connection is no longer available.`,
+          };
+        }
+
+        let execution;
+        try {
+          execution = await provider.executeTool(
+            toolName,
+            args,
+            connection,
+            { signal: options.signal || null },
+          );
+        } catch (err) {
+          if (
+            options.signal?.aborted ||
+            err?.name === 'AbortError' ||
+            err?.code === 'ABORT_ERR'
+          ) {
+            throw options.signal?.aborted ? abortError(options.signal) : err;
+          }
+          if (isLikelyExpiredConnectionError(err)) {
+            db.prepare(
+              `UPDATE integration_connections
+               SET status = 'expired', updated_at = datetime('now')
+               WHERE id = ? AND user_id = ? AND agent_id = ?`,
+            ).run(
+              connection.id,
+              userId,
+              resolveAgentId(userId, agentId),
+            );
+            return {
+              error: `Your ${provider.label} account connection has expired and needs to be reconnected. Open Official Integrations to reconnect your account.`,
+              expired: true,
+              connectionId: connection.id,
+            };
+          }
+          return { error: err?.message || 'execution_error' };
+        }
+        if (!execution) {
+          return { error: 'execution_failed' };
+        }
+
+        if (execution.credentials) {
+          const existingCredentials = this.parseCredentials(
+            connection.credentials_json,
+          );
+          const mergedCredentials = this.mergeUpdatedCredentials(
+            existingCredentials,
+            execution.credentials,
+          );
+          this.persistSharedCredentials(
             userId,
             resolveAgentId(userId, agentId),
+            provider.key,
+            connection.account_email,
+            mergedCredentials,
           );
-          return {
-            error: `Your ${provider.label} account connection has expired and needs to be reconnected. Open Official Integrations to reconnect your account.`,
-            expired: true,
-            connectionId: selection.connection.id,
-          };
         }
-        return { error: err?.message || 'execution_error' };
-      }
-      if (!execution) {
-        return { error: 'execution_failed' };
-      }
 
-      if (execution.credentials) {
-        const existingCredentials = this.parseCredentials(
-          selection.connection.credentials_json,
-        );
-        const mergedCredentials = this.mergeUpdatedCredentials(
-          existingCredentials,
-          execution.credentials,
-        );
-        this.persistSharedCredentials(
-          userId,
-          resolveAgentId(userId, agentId),
-          provider.key,
-          selection.connection.account_email,
-          mergedCredentials,
-        );
+        return execution.result;
+      } finally {
+        releaseExecution();
       }
-
-      return execution.result;
     }
 
     return foundSupportingProvider
@@ -716,6 +809,18 @@ class IntegrationManager {
       : { error: 'no_provider_support' };
   }
 
+  async shutdown() {
+    const providers = this.registry.list();
+    await Promise.allSettled(
+      providers.map((provider) => {
+        if (typeof provider.shutdown !== 'function') return null;
+        return provider.shutdown();
+      }),
+    );
+    this.connectionExecutionQueues.clear();
+    return { state: 'stopped', providerCount: providers.length };
+  }
+
   summarizeConnectedProviders(userId, agentId = null) {
     const scopedAgentId = resolveAgentId(userId, agentId);
     const ingestionService = this.app?.locals?.memoryIngestionService || null;

package/server/services/integrations/microsoft/provider.js

@@ -202,19 +202,83 @@ function graphUrl(path, query) {
   return url.toString();
 }
 
-async function graphRequest(credentials, { method = 'GET', path, query, body }) {
-  return fetchJson(
+function expiresAtFromSeconds(expiresIn) {
+  return new Date(
+    Date.now() + Math.max(1, Number(expiresIn) || 3600) * 1000,
+  ).toISOString();
+}
+
+function tokenExpiresSoon(credentials) {
+  const expiresAt = Date.parse(String(credentials?.expires_at || ''));
+  return Number.isFinite(expiresAt) && expiresAt <= Date.now() + 60 * 1000;
+}
+
+async function refreshMicrosoftCredentials(credentials, signal) {
+  const refreshToken = String(credentials?.refresh_token || '').trim();
+  if (!refreshToken) {
+    throw new Error(
+      'Microsoft 365 refresh token is missing. Reconnect this integration account.',
+    );
+  }
+  const { config, tokenUrl } = getMicrosoftEndpoints();
+  const token = await fetchJson(
+    tokenUrl,
+    {
+      method: 'POST',
+      form: {
+        client_id: config.clientId,
+        client_secret: config.clientSecret,
+        grant_type: 'refresh_token',
+        refresh_token: refreshToken,
+        scope: credentials.scope || undefined,
+      },
+      signal,
+    },
+    { serviceName: 'Microsoft 365 token refresh' },
+  );
+  if (!String(token?.access_token || '').trim()) {
+    throw new Error('Microsoft 365 token refresh did not return an access token.');
+  }
+  return {
+    ...credentials,
+    access_token: token.access_token,
+    refresh_token: token.refresh_token || refreshToken,
+    expires_in: token.expires_in,
+    expires_at: expiresAtFromSeconds(token.expires_in),
+    scope: token.scope || credentials.scope,
+    token_type: token.token_type || credentials.token_type,
+  };
+}
+
+async function graphRequest(context, { method = 'GET', path, query, body }) {
+  const { signal } = context;
+  let credentials = context.credentials;
+  if (tokenExpiresSoon(credentials)) {
+    credentials = await refreshMicrosoftCredentials(credentials, signal);
+    context.updateCredentials(credentials);
+  }
+  const performRequest = (activeCredentials) => fetchJson(
     graphUrl(path, query),
     {
       method: String(method || 'GET').toUpperCase(),
-      headers: { Authorization: `Bearer ${credentials.access_token}` },
+      headers: { Authorization: `Bearer ${activeCredentials.access_token}` },
       ...(body === undefined ? {} : { json: body }),
+      signal,
     },
     { serviceName: 'Microsoft Graph' },
   );
+
+  try {
+    return await performRequest(credentials);
+  } catch (error) {
+    if (error?.status !== 401 || !credentials.refresh_token) throw error;
+    credentials = await refreshMicrosoftCredentials(credentials, signal);
+    context.updateCredentials(credentials);
+    return performRequest(credentials);
+  }
 }
 
-async function executeMicrosoftTool(toolName, args, { credentials }) {
+async function executeMicrosoftTool(toolName, args, context) {
   switch (toolName) {
     case 'microsoft_365_outlook_list_messages': {
       const folder = String(args.folder_id || 'inbox').trim();
@@ -222,7 +286,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
         ? '/v1.0/me/mailFolders/inbox/messages'
         : `/v1.0/me/mailFolders/${encodeURIComponent(folder)}/messages`;
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           path,
           query: {
             '$top': Math.max(1, Math.min(Number(args.top) || 10, 50)),
@@ -233,7 +297,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
       };
     }
     case 'microsoft_365_outlook_send_mail':
-      await graphRequest(credentials, {
+      await graphRequest(context, {
         method: 'POST',
         path: '/v1.0/me/sendMail',
         body: {
@@ -255,7 +319,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
     case 'microsoft_365_calendar_list_events': {
       const hasWindow = args.start && args.end;
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           path: hasWindow ? '/v1.0/me/calendarView' : '/v1.0/me/events',
           query: {
             '$top': Math.max(1, Math.min(Number(args.top) || 10, 50)),
@@ -268,7 +332,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
     }
     case 'microsoft_365_calendar_create_event':
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           method: 'POST',
           path: '/v1.0/me/events',
           body: {
@@ -294,7 +358,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
       };
     case 'microsoft_365_onedrive_list_children':
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           path: args.item_id
             ? `/v1.0/me/drive/items/${encodeURIComponent(String(args.item_id))}/children`
             : '/v1.0/me/drive/root/children',
@@ -303,14 +367,14 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
       };
     case 'microsoft_365_teams_list_chats':
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           path: '/v1.0/me/chats',
           query: { '$top': Math.max(1, Math.min(Number(args.top) || 20, 50)) },
         }),
       };
     case 'microsoft_365_teams_send_chat_message':
       return {
-        result: await graphRequest(credentials, {
+        result: await graphRequest(context, {
           method: 'POST',
           path: `/v1.0/chats/${encodeURIComponent(requireText(args.chat_id, 'chat_id'))}/messages`,
           body: {
@@ -327,7 +391,7 @@ async function executeMicrosoftTool(toolName, args, { credentials }) {
           throw new Error('microsoft_365_*_graph_request tools are disabled by default. Set NEOAGENT_ENABLE_MICROSOFT_DYNAMIC_GRAPH_REQUEST=true to enable them.');
         }
         return {
-          result: await graphRequest(credentials, {
+          result: await graphRequest(context, {
             method: args.method,
             path: requireText(args.path, 'path'),
             query: args.query,
@@ -371,7 +435,7 @@ function createMicrosoftProvider() {
         appId: app.id,
       };
     },
-    async finishOAuth({ code, app }) {
+    async finishOAuth({ code, app, signal }) {
       const { config, tokenUrl } = getMicrosoftEndpoints();
       const token = await fetchJson(
         tokenUrl,
@@ -385,6 +449,7 @@ function createMicrosoftProvider() {
             redirect_uri: config.redirectUri,
             scope: escapeScope(app.scopes),
           },
+          signal,
         },
         { serviceName: 'Microsoft 365' },
       );
@@ -395,6 +460,7 @@ function createMicrosoftProvider() {
           headers: {
             Authorization: `Bearer ${token.access_token}`,
           },
+          signal,
         },
         { serviceName: 'Microsoft 365' },
       );
@@ -413,6 +479,7 @@ function createMicrosoftProvider() {
           access_token: token.access_token,
           refresh_token: token.refresh_token,
           expires_in: token.expires_in,
+          expires_at: expiresAtFromSeconds(token.expires_in),
           scope: token.scope,
           token_type: token.token_type,
         },