Forceflush fix (#474)

* fix

* fix

* changeset

* fix

* fix

Author: shagun-singh-inkeep

.changeset/quick-moments-juggle.md

@@ -0,0 +1,5 @@
+---
+"@inkeep/agents-run-api": patch
+---
+
+force flush fix

agents-run-api/src/agents/Agent.ts

@@ -60,7 +60,6 @@ import { toolSessionManager } from './ToolSessionManager';
 import type { SystemPromptV1 } from './types';
 import { Phase1Config } from './versions/v1/Phase1Config';
 import { Phase2Config } from './versions/v1/Phase2Config';
-import { defaultBatchProcessor } from '../instrumentation';
 
 /**
  * Creates a stopWhen condition that stops when any tool call name starts with the given prefix
@@ -750,7 +749,6 @@ export class Agent {
         requestContext: requestContext || {},
         tenantId: this.config.tenantId,
       });
-      await defaultBatchProcessor.forceFlush();
 
       // Add built-in variables to resolved context
       const contextWithBuiltins = {
@@ -1473,7 +1471,6 @@ export class Agent {
               throw err;
             } finally {
               childSpan.end();
-              await defaultBatchProcessor.forceFlush();
             }
           }
         );
@@ -2044,7 +2041,6 @@ ${output}${structureHintsFormatted}`;
         // Mark span as successful
         span.setStatus({ code: SpanStatusCode.OK });
         span.end();
-        await defaultBatchProcessor.forceFlush();
 
         // Format response - handle object vs text responses differently
         // Only format if we don't already have formattedContent from streaming

agents-run-api/src/app.ts

@@ -12,7 +12,7 @@ import { cors } from 'hono/cors';
 import { HTTPException } from 'hono/http-exception';
 import { requestId } from 'hono/request-id';
 import type { StatusCode } from 'hono/utils/http-status';
-import { defaultBatchProcessor } from './instrumentation';
+import { flushBatchProcessor } from './instrumentation';
 import { getLogger } from './logger';
 import { apiKeyAuth } from './middleware/api-key-auth';
 import { setupOpenAPIRoutes } from './openapi';
@@ -263,19 +263,19 @@ function createExecutionHono(
 
   app.use('/tenants/*', async (_c, next) => {
     await next();
-    await defaultBatchProcessor.forceFlush();
+    await flushBatchProcessor();
   });
   app.use('/agents/*', async (_c, next) => {
     await next();
-    await defaultBatchProcessor.forceFlush();
+    await flushBatchProcessor();
   });
   app.use('/v1/*', async (_c, next) => {
     await next();
-    await defaultBatchProcessor.forceFlush();
+    await flushBatchProcessor();
   });
   app.use('/api/*', async (_c, next) => {
     await next();
-    await defaultBatchProcessor.forceFlush();
+    await flushBatchProcessor();
   });
 
   const baseApp = new Hono();

agents-run-api/src/env.ts

@@ -19,6 +19,8 @@ const envSchema = z.object({
   OPENAI_API_KEY: z.string().optional(),
   ANTHROPIC_API_KEY: z.string(),
   INKEEP_AGENTS_RUN_API_BYPASS_SECRET: z.string().optional(),
+  OTEL_BSP_SCHEDULE_DELAY: z.coerce.number().optional().default(500),
+  OTEL_BSP_MAX_EXPORT_BATCH_SIZE: z.coerce.number().optional().default(64),
 });
 
 const parseEnv = () => {

agents-run-api/src/handlers/executionHandler.ts

@@ -20,7 +20,6 @@ import { agentInitializingOp, completionOp, errorOp } from '../utils/agent-opera
 import type { StreamHelper } from '../utils/stream-helpers.js';
 import { MCPStreamHelper } from '../utils/stream-helpers.js';
 import { registerStreamHelper, unregisterStreamHelper } from '../utils/stream-registry.js';
-import { defaultBatchProcessor } from '../instrumentation';
 
 const logger = getLogger('ExecutionHandler');
 
@@ -454,7 +453,6 @@ export class ExecutionHandler {
               throw error;
             } finally {
               span.end();
-              await defaultBatchProcessor.forceFlush();
             }
           });
         }

agents-run-api/src/instrumentation.ts

@@ -19,6 +19,7 @@ import {
   NoopSpanProcessor,
 } from '@opentelemetry/sdk-trace-base';
 import { ATTR_SERVICE_NAME } from '@opentelemetry/semantic-conventions';
+import { env } from './env';
 import { getLogger } from './logger';
 
 const otlpExporter = new OTLPTraceExporter();
@@ -29,7 +30,8 @@ const logger = getLogger('instrumentation');
 function createSafeBatchProcessor(): SpanProcessor {
   try {
     return new BatchSpanProcessor(otlpExporter, {
-      scheduledDelayMillis: 1000,
+      scheduledDelayMillis: env.OTEL_BSP_SCHEDULE_DELAY,
+      maxExportBatchSize: env.OTEL_BSP_MAX_EXPORT_BATCH_SIZE,
     });
   } catch (error) {
     logger.warn({ error }, 'Failed to create batch processor');
@@ -84,3 +86,11 @@ export const defaultSDK = new NodeSDK({
   spanProcessors: defaultSpanProcessors,
   instrumentations: defaultInstrumentations,
 });
+
+export async function flushBatchProcessor(): Promise<void> {
+  try {
+    await defaultBatchProcessor.forceFlush();
+  } catch (error) {
+    logger.warn({ error }, 'Failed to flush batch processor');
+  }
+}

agents-run-api/src/routes/chatDataStream.ts

@@ -22,7 +22,6 @@ import { ExecutionHandler } from '../handlers/executionHandler';
 import { getLogger } from '../logger';
 import { errorOp } from '../utils/agent-operations';
 import { createVercelStreamHelper } from '../utils/stream-helpers';
-import { defaultBatchProcessor } from '../instrumentation';
 
 type AppVariables = {
   credentialStores: CredentialStoreRegistry;
@@ -170,7 +169,6 @@ app.openapi(chatDataStreamRoute, async (c) => {
       dbClient,
       credentialStores,
     });
-    await defaultBatchProcessor.forceFlush();
 
     // Store last user message
     const lastUserMessage = body.messages.filter((m: any) => m.role === 'user').slice(-1)[0];

agents-run-api/src/services/GraphSession.ts

@@ -19,7 +19,6 @@ import { getStreamHelper } from '../utils/stream-registry';
 import { setSpanWithError, tracer } from '../utils/tracer';
 import { ArtifactParser } from './ArtifactParser';
 import { ArtifactService } from './ArtifactService';
-import { defaultBatchProcessor } from '../instrumentation';
 
 const logger = getLogger('GraphSession');
 
@@ -959,7 +958,6 @@ ${this.statusUpdateState?.config.prompt?.trim() || ''}`;
           return { summaries: [] };
         } finally {
           span.end();
-          await defaultBatchProcessor.forceFlush();
         }
       }
     );
@@ -1474,7 +1472,6 @@ Make it specific and relevant.`;
         } finally {
           // Always end the main span
           span.end();
-          await defaultBatchProcessor.forceFlush();
         }
       }
     );

agents-run-api/src/services/ResponseFormatter.ts

@@ -3,7 +3,6 @@ import { getLogger } from '../logger';
 import { ArtifactParser, type StreamPart } from '../services/ArtifactParser';
 import { tracer, setSpanWithError } from '../utils/tracer';
 import { graphSessionManager } from '../services/GraphSession';
-import { defaultBatchProcessor } from '../instrumentation';
 
 const logger = getLogger('ResponseFormatter');
 
@@ -84,7 +83,6 @@ export class ResponseFormatter {
         };
       } finally {
         span.end();
-        await defaultBatchProcessor.forceFlush();
       }
     });
   }
@@ -151,7 +149,6 @@ export class ResponseFormatter {
         return { text: responseText };
       } finally {
         span.end();
-        await defaultBatchProcessor.forceFlush();
       }
     });
   }