@juspay/neurolink 7.54.0 → 8.0.0

package/dist/neurolink.js

@@ -44,7 +44,7 @@ import { directToolsServer } from "./mcp/servers/agent/directToolsServer.js";
 // Import orchestration components
 import { ModelRouter } from "./utils/modelRouter.js";
 import { BinaryTaskClassifier } from "./utils/taskClassifier.js";
-import { initializeOpenTelemetry, shutdownOpenTelemetry, flushOpenTelemetry, getLangfuseHealthStatus, } from "./services/server/ai/observability/instrumentation.js";
+import { initializeOpenTelemetry, shutdownOpenTelemetry, flushOpenTelemetry, getLangfuseHealthStatus, setLangfuseContext, } from "./services/server/ai/observability/instrumentation.js";
 export class NeuroLink {
     mcpInitialized = false;
     emitter = new EventEmitter();
@@ -94,6 +94,40 @@ export class NeuroLink {
     // Mem0 memory instance and config for conversation context
     mem0Instance;
     mem0Config;
+    /**
+     * Extract and set Langfuse context from options with proper async scoping
+     */
+    async setLangfuseContextFromOptions(options, callback) {
+        if (options.context &&
+            typeof options.context === "object" &&
+            options.context !== null) {
+            try {
+                const ctx = options.context;
+                if (ctx.userId || ctx.sessionId) {
+                    return await new Promise((resolve, reject) => {
+                        setLangfuseContext({
+                            userId: typeof ctx.userId === "string" ? ctx.userId : null,
+                            sessionId: typeof ctx.sessionId === "string" ? ctx.sessionId : null,
+                        }, async () => {
+                            try {
+                                const result = await callback();
+                                resolve(result);
+                            }
+                            catch (error) {
+                                reject(error);
+                            }
+                        });
+                    });
+                }
+            }
+            catch (error) {
+                logger.warn("Failed to set Langfuse context from options", {
+                    error: error instanceof Error ? error.message : String(error),
+                });
+            }
+        }
+        return await callback();
+    }
     /**
      * Simple sync config setup for mem0
      */
@@ -518,7 +552,7 @@ export class NeuroLink {
                     langfuseInitStartTimeNs: langfuseInitStartTime.toString(),
                     message: "Starting Langfuse observability initialization",
                 });
-                // Initialize OpenTelemetry FIRST (required for Langfuse v4)
+                // Initialize OpenTelemetry (sets defaults from config)
                 initializeOpenTelemetry(langfuseConfig);
                 const healthStatus = getLangfuseHealthStatus();
                 const langfuseInitDurationNs = process.hrtime.bigint() - langfuseInitStartTime;
@@ -1146,198 +1180,201 @@ export class NeuroLink {
         if (!options.input?.text || typeof options.input.text !== "string") {
             throw new Error("Input text is required and must be a non-empty string");
         }
-        if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
-            options.context?.userId) {
-            try {
-                const mem0 = await this.ensureMem0Ready();
-                if (!mem0) {
-                    logger.debug("Mem0 not available, continuing without memory retrieval");
+        // Set session and user IDs from context for Langfuse spans and execute with proper async scoping
+        return await this.setLangfuseContextFromOptions(options, async () => {
+            if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
+                options.context?.userId) {
+                try {
+                    const mem0 = await this.ensureMem0Ready();
+                    if (!mem0) {
+                        logger.debug("Mem0 not available, continuing without memory retrieval");
+                    }
+                    else {
+                        const memories = await mem0.search(options.input.text, {
+                            userId: options.context.userId,
+                            limit: 5,
+                        });
+                        if (memories?.results?.length > 0) {
+                            // Enhance the input with memory context
+                            const memoryContext = memories.results
+                                .map((m) => m.memory)
+                                .join("\n");
+                            options.input.text = this.formatMemoryContext(memoryContext, options.input.text);
+                        }
+                    }
                 }
-                else {
-                    const memories = await mem0.search(options.input.text, {
-                        userId: options.context.userId,
-                        limit: 5,
+                catch (error) {
+                    logger.warn("Mem0 memory retrieval failed:", error);
+                }
+            }
+            const startTime = Date.now();
+            // Apply orchestration if enabled and no specific provider/model requested
+            if (this.enableOrchestration && !options.provider && !options.model) {
+                try {
+                    const orchestratedOptions = await this.applyOrchestration(options);
+                    logger.debug("Orchestration applied", {
+                        originalProvider: options.provider || "auto",
+                        orchestratedProvider: orchestratedOptions.provider,
+                        orchestratedModel: orchestratedOptions.model,
+                        prompt: options.input.text.substring(0, 100),
                     });
-                    if (memories?.results?.length > 0) {
-                        // Enhance the input with memory context
-                        const memoryContext = memories.results
-                            .map((m) => m.memory)
-                            .join("\n");
-                        options.input.text = this.formatMemoryContext(memoryContext, options.input.text);
-                    }
+                    // Use orchestrated options
+                    Object.assign(options, orchestratedOptions);
+                }
+                catch (error) {
+                    logger.warn("Orchestration failed, continuing with original options", {
+                        error: error instanceof Error ? error.message : String(error),
+                        originalProvider: options.provider || "auto",
+                    });
+                    // Continue with original options if orchestration fails
                 }
             }
-            catch (error) {
-                logger.warn("Mem0 memory retrieval failed:", error);
+            // Emit generation start event (NeuroLink format - keep existing)
+            this.emitter.emit("generation:start", {
+                provider: options.provider || "auto",
+                timestamp: startTime,
+            });
+            // ADD: Bedrock-compatible response:start event
+            this.emitter.emit("response:start");
+            // ADD: Bedrock-compatible message event
+            this.emitter.emit("message", `Starting ${options.provider || "auto"} text generation...`);
+            // Process factory configuration
+            const factoryResult = processFactoryOptions(options);
+            // Validate factory configuration if present
+            if (factoryResult.hasFactoryConfig && options.factoryConfig) {
+                const validation = validateFactoryConfig(options.factoryConfig);
+                if (!validation.isValid) {
+                    logger.warn("Invalid factory configuration detected", {
+                        errors: validation.errors,
+                    });
+                    // Continue with warning rather than throwing - graceful degradation
+                }
             }
-        }
-        const startTime = Date.now();
-        // Apply orchestration if enabled and no specific provider/model requested
-        if (this.enableOrchestration && !options.provider && !options.model) {
-            try {
-                const orchestratedOptions = await this.applyOrchestration(options);
-                logger.debug("Orchestration applied", {
-                    originalProvider: options.provider || "auto",
-                    orchestratedProvider: orchestratedOptions.provider,
-                    orchestratedModel: orchestratedOptions.model,
-                    prompt: options.input.text.substring(0, 100),
-                });
-                // Use orchestrated options
-                Object.assign(options, orchestratedOptions);
+            // 🔧 CRITICAL FIX: Convert to TextGenerationOptions while preserving the input object for multimodal support
+            const baseOptions = {
+                prompt: options.input.text,
+                provider: options.provider,
+                model: options.model,
+                temperature: options.temperature,
+                maxTokens: options.maxTokens,
+                systemPrompt: options.systemPrompt,
+                schema: options.schema,
+                output: options.output,
+                disableTools: options.disableTools,
+                enableAnalytics: options.enableAnalytics,
+                enableEvaluation: options.enableEvaluation,
+                context: options.context,
+                evaluationDomain: options.evaluationDomain,
+                toolUsageContext: options.toolUsageContext,
+                input: options.input, // This includes text, images, and content arrays
+                region: options.region,
+            };
+            // Apply factory enhancement using centralized utilities
+            const textOptions = enhanceTextGenerationOptions(baseOptions, factoryResult);
+            // Pass conversation memory config if available
+            if (this.conversationMemory) {
+                textOptions.conversationMemoryConfig = this.conversationMemory.config;
+                // Include original prompt for context summarization
+                textOptions.originalPrompt = originalPrompt;
             }
-            catch (error) {
-                logger.warn("Orchestration failed, continuing with original options", {
-                    error: error instanceof Error ? error.message : String(error),
-                    originalProvider: options.provider || "auto",
+            // Detect and execute domain-specific tools
+            const { toolResults, enhancedPrompt } = await this.detectAndExecuteTools(textOptions.prompt || options.input.text, factoryResult.domainType);
+            // Update prompt with tool results if available
+            if (enhancedPrompt !== textOptions.prompt) {
+                textOptions.prompt = enhancedPrompt;
+                logger.debug("Enhanced prompt with tool results", {
+                    originalLength: options.input.text.length,
+                    enhancedLength: enhancedPrompt.length,
+                    toolResults: toolResults.length,
                 });
-                // Continue with original options if orchestration fails
             }
-        }
-        // Emit generation start event (NeuroLink format - keep existing)
-        this.emitter.emit("generation:start", {
-            provider: options.provider || "auto",
-            timestamp: startTime,
-        });
-        // ADD: Bedrock-compatible response:start event
-        this.emitter.emit("response:start");
-        // ADD: Bedrock-compatible message event
-        this.emitter.emit("message", `Starting ${options.provider || "auto"} text generation...`);
-        // Process factory configuration
-        const factoryResult = processFactoryOptions(options);
-        // Validate factory configuration if present
-        if (factoryResult.hasFactoryConfig && options.factoryConfig) {
-            const validation = validateFactoryConfig(options.factoryConfig);
-            if (!validation.isValid) {
-                logger.warn("Invalid factory configuration detected", {
-                    errors: validation.errors,
+            // Use redesigned generation logic
+            const textResult = await this.generateTextInternal(textOptions);
+            // Emit generation completion event (NeuroLink format - enhanced with content)
+            this.emitter.emit("generation:end", {
+                provider: textResult.provider,
+                responseTime: Date.now() - startTime,
+                toolsUsed: textResult.toolsUsed,
+                timestamp: Date.now(),
+                result: textResult, // Enhanced: include full result
+            });
+            // ADD: Bedrock-compatible response:end event with content
+            this.emitter.emit("response:end", textResult.content || "");
+            // ADD: Bedrock-compatible message event
+            this.emitter.emit("message", `Generation completed in ${Date.now() - startTime}ms`);
+            // Convert back to GenerateResult
+            const generateResult = {
+                content: textResult.content,
+                provider: textResult.provider,
+                model: textResult.model,
+                usage: textResult.usage
+                    ? {
+                        input: textResult.usage.input || 0,
+                        output: textResult.usage.output || 0,
+                        total: textResult.usage.total || 0,
+                    }
+                    : undefined,
+                responseTime: textResult.responseTime,
+                toolsUsed: textResult.toolsUsed,
+                toolExecutions: transformToolExecutions(textResult.toolExecutions),
+                enhancedWithTools: textResult.enhancedWithTools,
+                availableTools: transformAvailableTools(textResult.availableTools),
+                analytics: textResult.analytics,
+                evaluation: textResult.evaluation
+                    ? {
+                        ...textResult.evaluation,
+                        isOffTopic: textResult.evaluation
+                            .isOffTopic ?? false,
+                        alertSeverity: textResult.evaluation
+                            .alertSeverity ??
+                            "none",
+                        reasoning: textResult.evaluation
+                            .reasoning ?? "No evaluation provided",
+                        evaluationModel: textResult.evaluation
+                            .evaluationModel ?? "unknown",
+                        evaluationTime: textResult.evaluation
+                            .evaluationTime ?? Date.now(),
+                        // Include evaluationDomain from original options
+                        evaluationDomain: textResult.evaluation
+                            .evaluationDomain ??
+                            textOptions.evaluationDomain ??
+                            factoryResult.domainType,
+                    }
+                    : undefined,
+            };
+            if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
+                options.context?.userId &&
+                generateResult.content) {
+                // Non-blocking memory storage - run in background
+                setImmediate(async () => {
+                    try {
+                        const mem0 = await this.ensureMem0Ready();
+                        if (mem0) {
+                            // Store complete conversation turn (user + AI messages)
+                            const conversationTurn = [
+                                { role: "user", content: options.input.text },
+                                { role: "system", content: generateResult.content },
+                            ];
+                            await mem0.add(JSON.stringify(conversationTurn), {
+                                userId: options.context?.userId,
+                                metadata: {
+                                    timestamp: new Date().toISOString(),
+                                    provider: generateResult.provider,
+                                    model: generateResult.model,
+                                    type: "conversation_turn",
+                                    async_mode: true,
+                                },
+                            });
+                        }
+                    }
+                    catch (error) {
+                        // Non-blocking: Log error but don't fail the generation
+                        logger.warn("Mem0 memory storage failed:", error);
+                    }
                 });
-                // Continue with warning rather than throwing - graceful degradation
             }
-        }
-        // 🔧 CRITICAL FIX: Convert to TextGenerationOptions while preserving the input object for multimodal support
-        const baseOptions = {
-            prompt: options.input.text,
-            provider: options.provider,
-            model: options.model,
-            temperature: options.temperature,
-            maxTokens: options.maxTokens,
-            systemPrompt: options.systemPrompt,
-            schema: options.schema,
-            output: options.output,
-            disableTools: options.disableTools,
-            enableAnalytics: options.enableAnalytics,
-            enableEvaluation: options.enableEvaluation,
-            context: options.context,
-            evaluationDomain: options.evaluationDomain,
-            toolUsageContext: options.toolUsageContext,
-            input: options.input, // This includes text, images, and content arrays
-            region: options.region,
-        };
-        // Apply factory enhancement using centralized utilities
-        const textOptions = enhanceTextGenerationOptions(baseOptions, factoryResult);
-        // Pass conversation memory config if available
-        if (this.conversationMemory) {
-            textOptions.conversationMemoryConfig = this.conversationMemory.config;
-            // Include original prompt for context summarization
-            textOptions.originalPrompt = originalPrompt;
-        }
-        // Detect and execute domain-specific tools
-        const { toolResults, enhancedPrompt } = await this.detectAndExecuteTools(textOptions.prompt || options.input.text, factoryResult.domainType);
-        // Update prompt with tool results if available
-        if (enhancedPrompt !== textOptions.prompt) {
-            textOptions.prompt = enhancedPrompt;
-            logger.debug("Enhanced prompt with tool results", {
-                originalLength: options.input.text.length,
-                enhancedLength: enhancedPrompt.length,
-                toolResults: toolResults.length,
-            });
-        }
-        // Use redesigned generation logic
-        const textResult = await this.generateTextInternal(textOptions);
-        // Emit generation completion event (NeuroLink format - enhanced with content)
-        this.emitter.emit("generation:end", {
-            provider: textResult.provider,
-            responseTime: Date.now() - startTime,
-            toolsUsed: textResult.toolsUsed,
-            timestamp: Date.now(),
-            result: textResult, // Enhanced: include full result
+            return generateResult;
         });
-        // ADD: Bedrock-compatible response:end event with content
-        this.emitter.emit("response:end", textResult.content || "");
-        // ADD: Bedrock-compatible message event
-        this.emitter.emit("message", `Generation completed in ${Date.now() - startTime}ms`);
-        // Convert back to GenerateResult
-        const generateResult = {
-            content: textResult.content,
-            provider: textResult.provider,
-            model: textResult.model,
-            usage: textResult.usage
-                ? {
-                    input: textResult.usage.input || 0,
-                    output: textResult.usage.output || 0,
-                    total: textResult.usage.total || 0,
-                }
-                : undefined,
-            responseTime: textResult.responseTime,
-            toolsUsed: textResult.toolsUsed,
-            toolExecutions: transformToolExecutions(textResult.toolExecutions),
-            enhancedWithTools: textResult.enhancedWithTools,
-            availableTools: transformAvailableTools(textResult.availableTools),
-            analytics: textResult.analytics,
-            evaluation: textResult.evaluation
-                ? {
-                    ...textResult.evaluation,
-                    isOffTopic: textResult.evaluation
-                        .isOffTopic ?? false,
-                    alertSeverity: textResult.evaluation
-                        .alertSeverity ??
-                        "none",
-                    reasoning: textResult.evaluation
-                        .reasoning ?? "No evaluation provided",
-                    evaluationModel: textResult.evaluation
-                        .evaluationModel ?? "unknown",
-                    evaluationTime: textResult.evaluation
-                        .evaluationTime ?? Date.now(),
-                    // Include evaluationDomain from original options
-                    evaluationDomain: textResult.evaluation
-                        .evaluationDomain ??
-                        textOptions.evaluationDomain ??
-                        factoryResult.domainType,
-                }
-                : undefined,
-        };
-        if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
-            options.context?.userId &&
-            generateResult.content) {
-            // Non-blocking memory storage - run in background
-            setImmediate(async () => {
-                try {
-                    const mem0 = await this.ensureMem0Ready();
-                    if (mem0) {
-                        // Store complete conversation turn (user + AI messages)
-                        const conversationTurn = [
-                            { role: "user", content: options.input.text },
-                            { role: "system", content: generateResult.content },
-                        ];
-                        await mem0.add(JSON.stringify(conversationTurn), {
-                            userId: options.context?.userId,
-                            metadata: {
-                                timestamp: new Date().toISOString(),
-                                provider: generateResult.provider,
-                                model: generateResult.model,
-                                type: "conversation_turn",
-                                async_mode: true,
-                            },
-                        });
-                    }
-                }
-                catch (error) {
-                    // Non-blocking: Log error but don't fail the generation
-                    logger.warn("Mem0 memory storage failed:", error);
-                }
-            });
-        }
-        return generateResult;
     }
     /**
      * BACKWARD COMPATIBILITY: Legacy generateText method
@@ -1863,153 +1900,155 @@ export class NeuroLink {
         const originalPrompt = options.input.text; // Store the original prompt for memory storage
         await this.validateStreamInput(options);
         this.emitStreamStartEvents(options, startTime);
-        let enhancedOptions;
-        let factoryResult;
-        try {
-            // Initialize conversation memory if needed (for lazy loading)
-            await this.initializeConversationMemoryForGeneration(streamId, startTime, hrTimeStart);
-            // Initialize MCP
-            await this.initializeMCP();
-            const _originalPrompt = options.input.text;
-            if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
-                options.context?.userId) {
-                try {
-                    const mem0 = await this.ensureMem0Ready();
-                    if (!mem0) {
-                        // Continue without memories if mem0 is not available
-                        logger.debug("Mem0 not available, continuing without memory retrieval");
-                    }
-                    else {
-                        const memories = await mem0.search(options.input.text, {
-                            userId: options.context.userId,
-                            limit: 5,
-                        });
-                        if (memories?.results?.length > 0) {
-                            // Enhance the input with memory context
-                            const memoryContext = memories.results
-                                .map((m) => m.memory)
-                                .join("\n");
-                            options.input.text = this.formatMemoryContext(memoryContext, options.input.text);
+        // Set session and user IDs from context for Langfuse spans and execute with proper async scoping
+        return await this.setLangfuseContextFromOptions(options, async () => {
+            let enhancedOptions;
+            let factoryResult;
+            try {
+                // Initialize conversation memory if needed (for lazy loading)
+                await this.initializeConversationMemoryForGeneration(streamId, startTime, hrTimeStart);
+                // Initialize MCP
+                await this.initializeMCP();
+                const _originalPrompt = options.input.text;
+                if (this.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
+                    options.context?.userId) {
+                    try {
+                        const mem0 = await this.ensureMem0Ready();
+                        if (!mem0) {
+                            // Continue without memories if mem0 is not available
+                            logger.debug("Mem0 not available, continuing without memory retrieval");
+                        }
+                        else {
+                            const memories = await mem0.search(options.input.text, {
+                                userId: options.context.userId,
+                                limit: 5,
+                            });
+                            if (memories?.results?.length > 0) {
+                                // Enhance the input with memory context
+                                const memoryContext = memories.results
+                                    .map((m) => m.memory)
+                                    .join("\n");
+                                options.input.text = this.formatMemoryContext(memoryContext, options.input.text);
+                            }
                         }
                     }
+                    catch (error) {
+                        // Non-blocking: Log error but continue with streaming
+                        logger.warn("Mem0 memory retrieval failed:", error);
+                    }
                 }
-                catch (error) {
-                    // Non-blocking: Log error but continue with streaming
-                    logger.warn("Mem0 memory retrieval failed:", error);
-                }
-            }
-            // Apply orchestration if enabled and no specific provider/model requested
-            if (this.enableOrchestration && !options.provider && !options.model) {
-                try {
-                    const orchestratedOptions = await this.applyStreamOrchestration(options);
-                    logger.debug("Stream orchestration applied", {
-                        originalProvider: options.provider || "auto",
-                        orchestratedProvider: orchestratedOptions.provider,
-                        orchestratedModel: orchestratedOptions.model,
-                        prompt: options.input.text?.substring(0, 100),
-                    });
-                    // Use orchestrated options
-                    Object.assign(options, orchestratedOptions);
-                }
-                catch (error) {
-                    logger.warn("Stream orchestration failed, continuing with original options", {
-                        error: error instanceof Error ? error.message : String(error),
-                        originalProvider: options.provider || "auto",
-                    });
-                    // Continue with original options if orchestration fails
-                }
-            }
-            factoryResult = processStreamingFactoryOptions(options);
-            enhancedOptions = createCleanStreamOptions(options);
-            if (options.input?.text) {
-                const { toolResults: _toolResults, enhancedPrompt } = await this.detectAndExecuteTools(options.input.text, undefined);
-                if (enhancedPrompt !== options.input.text) {
-                    enhancedOptions.input.text = enhancedPrompt;
+                // Apply orchestration if enabled and no specific provider/model requested
+                if (this.enableOrchestration && !options.provider && !options.model) {
+                    try {
+                        const orchestratedOptions = await this.applyStreamOrchestration(options);
+                        logger.debug("Stream orchestration applied", {
+                            originalProvider: options.provider || "auto",
+                            orchestratedProvider: orchestratedOptions.provider,
+                            orchestratedModel: orchestratedOptions.model,
+                            prompt: options.input.text?.substring(0, 100),
+                        });
+                        // Use orchestrated options
+                        Object.assign(options, orchestratedOptions);
+                    }
+                    catch (error) {
+                        logger.warn("Stream orchestration failed, continuing with original options", {
+                            error: error instanceof Error ? error.message : String(error),
+                            originalProvider: options.provider || "auto",
+                        });
+                        // Continue with original options if orchestration fails
+                    }
                 }
-            }
-            const { stream: mcpStream, provider: providerName } = await this.createMCPStream(enhancedOptions);
-            // Create a wrapper around the stream that accumulates content
-            let accumulatedContent = "";
-            const processedStream = (async function* (self) {
-                try {
-                    for await (const chunk of mcpStream) {
-                        if (chunk &&
-                            "content" in chunk &&
-                            typeof chunk.content === "string") {
-                            accumulatedContent += chunk.content;
-                            // Emit chunk event for compatibility
-                            self.emitter.emit("response:chunk", chunk.content);
-                        }
-                        yield chunk; // Preserve original streaming behavior
+                factoryResult = processStreamingFactoryOptions(options);
+                enhancedOptions = createCleanStreamOptions(options);
+                if (options.input?.text) {
+                    const { toolResults: _toolResults, enhancedPrompt } = await this.detectAndExecuteTools(options.input.text, undefined);
+                    if (enhancedPrompt !== options.input.text) {
+                        enhancedOptions.input.text = enhancedPrompt;
                     }
                 }
-                finally {
-                    // Store memory after stream consumption is complete
-                    if (self.conversationMemory && enhancedOptions.context?.sessionId) {
-                        const sessionId = enhancedOptions.context?.sessionId;
-                        const userId = enhancedOptions.context
-                            ?.userId;
-                        try {
-                            await self.conversationMemory.storeConversationTurn(sessionId, userId, originalPrompt ?? "", accumulatedContent, new Date(startTime));
-                            logger.debug("Stream conversation turn stored", {
-                                sessionId,
-                                userInputLength: originalPrompt?.length ?? 0,
-                                responseLength: accumulatedContent.length,
-                            });
-                        }
-                        catch (error) {
-                            logger.warn("Failed to store stream conversation turn", {
-                                error: error instanceof Error ? error.message : String(error),
-                            });
+                const { stream: mcpStream, provider: providerName } = await this.createMCPStream(enhancedOptions);
+                // Create a wrapper around the stream that accumulates content
+                let accumulatedContent = "";
+                const processedStream = (async function* (self) {
+                    try {
+                        for await (const chunk of mcpStream) {
+                            if (chunk &&
+                                "content" in chunk &&
+                                typeof chunk.content === "string") {
+                                accumulatedContent += chunk.content;
+                                // Emit chunk event for compatibility
+                                self.emitter.emit("response:chunk", chunk.content);
+                            }
+                            yield chunk; // Preserve original streaming behavior
                         }
                     }
-                    if (self.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
-                        enhancedOptions.context?.userId &&
-                        accumulatedContent.trim()) {
-                        // Non-blocking memory storage - run in background
-                        setImmediate(async () => {
+                    finally {
+                        // Store memory after stream consumption is complete
+                        if (self.conversationMemory && enhancedOptions.context?.sessionId) {
+                            const sessionId = enhancedOptions.context?.sessionId;
+                            const userId = enhancedOptions.context?.userId;
                             try {
-                                const mem0 = await self.ensureMem0Ready();
-                                if (mem0) {
-                                    // Store complete conversation turn (user + AI messages)
-                                    const conversationTurn = [
-                                        { role: "user", content: originalPrompt },
-                                        { role: "system", content: accumulatedContent.trim() },
-                                    ];
-                                    await mem0.add(JSON.stringify(conversationTurn), {
-                                        userId: enhancedOptions.context?.userId,
-                                        metadata: {
-                                            timestamp: new Date().toISOString(),
-                                            type: "conversation_turn_stream",
-                                            userMessage: originalPrompt,
-                                            async_mode: true,
-                                            aiResponse: accumulatedContent.trim(),
-                                        },
-                                    });
-                                }
+                                await self.conversationMemory.storeConversationTurn(sessionId, userId, originalPrompt ?? "", accumulatedContent, new Date(startTime));
+                                logger.debug("Stream conversation turn stored", {
+                                    sessionId,
+                                    userInputLength: originalPrompt?.length ?? 0,
+                                    responseLength: accumulatedContent.length,
+                                });
                             }
                             catch (error) {
-                                logger.warn("Mem0 memory storage failed:", error);
+                                logger.warn("Failed to store stream conversation turn", {
+                                    error: error instanceof Error ? error.message : String(error),
+                                });
                             }
-                        });
+                        }
+                        if (self.conversationMemoryConfig?.conversationMemory?.mem0Enabled &&
+                            enhancedOptions.context?.userId &&
+                            accumulatedContent.trim()) {
+                            // Non-blocking memory storage - run in background
+                            setImmediate(async () => {
+                                try {
+                                    const mem0 = await self.ensureMem0Ready();
+                                    if (mem0) {
+                                        // Store complete conversation turn (user + AI messages)
+                                        const conversationTurn = [
+                                            { role: "user", content: originalPrompt },
+                                            { role: "system", content: accumulatedContent.trim() },
+                                        ];
+                                        await mem0.add(JSON.stringify(conversationTurn), {
+                                            userId: enhancedOptions.context?.userId,
+                                            metadata: {
+                                                timestamp: new Date().toISOString(),
+                                                type: "conversation_turn_stream",
+                                                userMessage: originalPrompt,
+                                                async_mode: true,
+                                                aiResponse: accumulatedContent.trim(),
+                                            },
+                                        });
+                                    }
+                                }
+                                catch (error) {
+                                    logger.warn("Mem0 memory storage failed:", error);
+                                }
+                            });
+                        }
                     }
-                }
-            })(this);
-            const streamResult = await this.processStreamResult(mcpStream, enhancedOptions, factoryResult);
-            const responseTime = Date.now() - startTime;
-            this.emitStreamEndEvents(streamResult);
-            return this.createStreamResponse(streamResult, processedStream, {
-                providerName,
-                options,
-                startTime,
-                responseTime,
-                streamId,
-                fallback: false,
-            });
-        }
-        catch (error) {
-            return this.handleStreamError(error, options, startTime, streamId, undefined, undefined);
-        }
+                })(this);
+                const streamResult = await this.processStreamResult(mcpStream, enhancedOptions, factoryResult);
+                const responseTime = Date.now() - startTime;
+                this.emitStreamEndEvents(streamResult);
+                return this.createStreamResponse(streamResult, processedStream, {
+                    providerName,
+                    options,
+                    startTime,
+                    responseTime,
+                    streamId,
+                    fallback: false,
+                });
+            }
+            catch (error) {
+                return this.handleStreamError(error, options, startTime, streamId, undefined, undefined);
+            }
+        });
     }
     /**
      * Validate stream input with comprehensive error reporting