agentshield-sdk 7.2.0 → 7.2.1

package/src/index.js

@@ -53,6 +53,15 @@ const DEFAULT_CONFIG = {
   /** Custom callback when a threat is detected. */
   onThreat: null,
 
+  /** Maximum input size in bytes before truncation warning. */
+  maxInputSize: 1_000_000,
+
+  /** Maximum number of scan history entries to retain. */
+  maxScanHistory: 100,
+
+  /** Maximum recursion depth when flattening tool arguments. */
+  maxArgDepth: 10,
+
   /** Dangerous tool names that should be scrutinized more carefully. */
   dangerousTools: [
     'bash', 'shell', 'terminal', 'exec', 'execute',
@@ -113,8 +122,8 @@ class AgentShield {
     if (typeof text !== 'string') {
       throw new TypeError(`[Agent Shield] scan() expects a string, got ${typeof text}`);
     }
-    if (text.length > 1_000_000) {
-      console.warn('[Agent Shield] Input exceeds 1MB — consider scanning in chunks');
+    if (text.length > this.config.maxInputSize) {
+      console.warn('[Agent Shield] Input exceeds configured maxInputSize - consider scanning in chunks');
     }
     const result = scanText(text, {
       source: options.source || 'unknown',
@@ -133,7 +142,7 @@ class AgentShield {
       threatCount: result.threats.length,
       source: options.source || 'unknown'
     });
-    if (this.stats.scanHistory.length > 100) {
+    if (this.stats.scanHistory.length > this.config.maxScanHistory) {
       this.stats.scanHistory.shift();
     }
 
@@ -197,6 +206,7 @@ class AgentShield {
    * @param {object} [options] - Options.
    * @param {string} [options.source='user_input'] - Where the input came from.
    * @returns {object} Scan result with additional `blocked` field.
+   * @throws {TypeError} If text is not a string.
    */
   scanInput(text, options = {}) {
     if (typeof text !== 'string') {
@@ -214,6 +224,7 @@ class AgentShield {
    * @param {object} [options] - Options.
    * @param {string} [options.source='agent_output'] - Source label.
    * @returns {object} Scan result with additional `blocked` field.
+   * @throws {TypeError} If text is not a string.
    */
   scanOutput(text, options = {}) {
     if (typeof text !== 'string') {
@@ -232,8 +243,11 @@ class AgentShield {
    * @returns {object} Scan result with `blocked` and `warnings` fields.
    */
   scanToolCall(toolName, args = {}, options = {}) {
-    if (!toolName || typeof toolName !== 'string') {
-      return { status: 'safe', toolName: toolName || '', threats: [], warnings: ['Invalid tool name'], blocked: false, isDangerousTool: false, timestamp: Date.now() };
+    if (typeof toolName !== 'string') {
+      throw new TypeError(`[Agent Shield] scanToolCall() expects toolName to be a string, got ${typeof toolName}`);
+    }
+    if (!toolName) {
+      return { status: 'safe', toolName: '', threats: [], warnings: ['Empty tool name'], blocked: false, isDangerousTool: false, timestamp: Date.now() };
     }
     const warnings = [];
     const allThreats = [];
@@ -371,7 +385,8 @@ class AgentShield {
    * @param {object} args
    * @returns {string}
    */
-  _flattenArgs(args, maxDepth = 10) {
+  _flattenArgs(args, maxDepth) {
+    if (maxDepth == null) maxDepth = this.config.maxArgDepth;
     const parts = [];
     const flatten = (obj, depth) => {
       if (depth > maxDepth) return;
@@ -393,7 +408,8 @@ class AgentShield {
    * @param {object} args
    * @returns {Array<string>}
    */
-  _extractFilePaths(args, maxDepth = 10) {
+  _extractFilePaths(args, maxDepth) {
+    if (maxDepth == null) maxDepth = this.config.maxArgDepth;
     const paths = [];
     const fileKeys = [
       'file', 'path', 'file_path', 'filepath', 'filename', 'target',

package/src/main.js

@@ -27,7 +27,7 @@ function safeRequire(path, label) {
 // Core (these are critical — if they fail, we still export what we can)
 const { AgentShield } = safeRequire('./index', 'core');
 const { scanText, getPatterns, SEVERITY_ORDER } = safeRequire('./detector-core', 'detector-core');
-const { expressMiddleware, wrapAgent, shieldTools, extractTextFromBody } = safeRequire('./middleware', 'middleware');
+const { expressMiddleware, wrapAgent, shieldTools, extractTextFromBody, rateLimitMiddleware, shieldMiddleware } = safeRequire('./middleware', 'middleware');
 
 // Protection
 const { CircuitBreaker, shadowMode, RateLimiter, STATE } = safeRequire('./circuit-breaker', 'circuit-breaker');
@@ -51,7 +51,7 @@ const { SteganographyDetector, EncodingBruteforceDetector, StructuredDataScanner
 const { OutputWatermark, DifferentialPrivacy } = safeRequire('./watermark', 'watermark');
 
 // Utilities
-const { getGrade, getGradeLabel, makeBar, truncate, formatHeader, generateId } = safeRequire('./utils', 'utils');
+const { getGrade, getGradeLabel, makeBar, truncate, formatHeader, generateId, createGracefulShutdown, loadEnvFile } = safeRequire('./utils', 'utils');
 
 // Error codes & deprecation
 const { ERROR_CODES, createShieldError, deprecationWarning } = safeRequire('./errors', 'errors');
@@ -288,6 +288,8 @@ const _exports = {
   wrapAgent,
   shieldTools,
   extractTextFromBody,
+  rateLimitMiddleware,
+  shieldMiddleware,
 
   // Protection
   CircuitBreaker,
@@ -343,6 +345,8 @@ const _exports = {
   truncate,
   formatHeader,
   generateId,
+  createGracefulShutdown,
+  loadEnvFile,
 
   // Integrations
   ShieldCallbackHandler,

package/src/mcp-security-runtime.js

@@ -812,16 +812,41 @@ class MCPSecurityRuntime {
 
   /**
    * Shuts down the runtime and cleans up resources.
+   * @param {object} [options]
+   * @param {number} [options.timeoutMs=10000] - Max time to wait for drain before forced cleanup.
+   * @returns {Promise<void>}
    */
-  shutdown() {
+  async shutdown(options = {}) {
+    const timeoutMs = options.timeoutMs || 10000;
+
     if (this._cleanupInterval) {
       clearInterval(this._cleanupInterval);
       this._cleanupInterval = null;
     }
-    const sessionIds = [...this._sessions.keys()];
-    for (const sessionId of sessionIds) {
-      this.terminateSession(sessionId);
-    }
+
+    // Drain: wait for in-flight operations with timeout
+    const drainPromise = new Promise((resolve) => {
+      const sessionIds = [...this._sessions.keys()];
+      for (const sessionId of sessionIds) {
+        try {
+          this.terminateSession(sessionId);
+        } catch (err) {
+          console.error(`${LOG_PREFIX} Error terminating session ${sessionId}: ${err.message}`);
+        }
+      }
+      resolve();
+    });
+
+    const timeoutPromise = new Promise((resolve) => {
+      const timer = setTimeout(() => {
+        console.error(`${LOG_PREFIX} Shutdown drain timeout (${timeoutMs}ms), forcing cleanup`);
+        resolve();
+      }, timeoutMs);
+      if (timer.unref) timer.unref();
+    });
+
+    await Promise.race([drainPromise, timeoutPromise]);
+
     this._audit('runtime_shutdown', { totalProcessed: this.stats.toolCallsProcessed });
   }
 

package/src/mcp-server.js

@@ -707,6 +707,11 @@ class MCPServer {
 // =========================================================================
 
 if (require.main === module) {
+  const { createGracefulShutdown, loadEnvFile } = require('./utils');
+
+  // Load .env file if present
+  loadEnvFile();
+
   let config = {};
 
   // Parse --config flag
@@ -726,15 +731,14 @@ if (require.main === module) {
   const server = new MCPServer(config);
   server.start();
 
-  // Graceful shutdown
-  process.on('SIGINT', () => {
-    server.stop();
-    process.exit(0);
-  });
-  process.on('SIGTERM', () => {
-    server.stop();
-    process.exit(0);
+  // Graceful shutdown with timeout enforcement
+  const { shutdown } = createGracefulShutdown({
+    timeoutMs: parseInt(process.env.SHIELD_SHUTDOWN_TIMEOUT_MS, 10) || 10000,
+    cleanupFns: [() => server.stop()]
   });
+
+  process.on('SIGINT', () => shutdown('SIGINT').then(() => process.exit(0)));
+  process.on('SIGTERM', () => shutdown('SIGTERM').then(() => process.exit(0)));
 }
 
 module.exports = { MCPServer, MCPToolHandler };