add run manager for observability callbacks

Author: Pederzh

src/agents/mcp_agent.ts

@@ -1,16 +1,18 @@
 import type { BaseCallbackHandler } from '@langchain/core/callbacks/base'
+import type { CallbackManagerForChainRun } from '@langchain/core/callbacks/manager'
 import type { BaseLanguageModelInterface, LanguageModelLike } from '@langchain/core/language_models/base'
+import type { Serialized } from '@langchain/core/load/serializable'
 import type {
   BaseMessage,
 } from '@langchain/core/messages'
 import type { StructuredToolInterface, ToolInterface } from '@langchain/core/tools'
 import type { StreamEvent } from '@langchain/core/tracers/log_stream'
 import type { AgentFinish, AgentStep } from 'langchain/agents'
-
 import type { ZodSchema } from 'zod'
 import type { MCPClient } from '../client.js'
 import type { BaseConnector } from '../connectors/base.js'
 import type { MCPSession } from '../session.js'
+import { CallbackManager } from '@langchain/core/callbacks/manager'
 import {
   AIMessage,
   HumanMessage,
@@ -22,10 +24,7 @@ import {
   ChatPromptTemplate,
   MessagesPlaceholder,
 } from '@langchain/core/prompts'
-import {
-  AgentExecutor,
-  createToolCallingAgent,
-} from 'langchain/agents'
+import { AgentExecutor, createToolCallingAgent } from 'langchain/agents'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 import { LangChainAdapter } from '../adapters/langchain_adapter.js'
 import { logger } from '../logging.js'
@@ -503,6 +502,20 @@ export class MCPAgent {
       let nameToToolMap: Record<string, StructuredToolInterface> = Object.fromEntries(this._tools.map(t => [t.name, t]))
       logger.info(`🏁 Starting agent execution with max_steps=${steps}`)
 
+      // Create a run manager with our callbacks if we have any - ONCE for the entire execution
+      let runManager: CallbackManagerForChainRun | undefined
+      if (this.callbacks?.length > 0) {
+        // Create an async callback manager with our callbacks
+        const callbackManager = new CallbackManager(undefined, {
+          handlers: this.callbacks,
+          inheritableHandlers: this.callbacks,
+        })
+        // Create a run manager for this chain execution
+        runManager = await callbackManager.handleChainStart({
+          name: 'MCPAgent (mcp-use)',
+        } as Serialized, inputs)
+      }
+
       for (let stepNum = 0; stepNum < steps; stepNum++) {
         stepsTaken = stepNum + 1
         if (this.useServerManager && this.serverManager) {
@@ -531,15 +544,17 @@ export class MCPAgent {
 
         try {
           logger.debug('Starting agent step execution')
-          const nextStepOutput = await this._agentExecutor._takeNextStep(
+          const nextStepOutput: AgentStep[] | AgentFinish = await this._agentExecutor._takeNextStep(
             nameToToolMap as Record<string, ToolInterface>,
             inputs,
             intermediateSteps,
+            runManager,
           )
-          // Agent finish handling
-          if ((nextStepOutput as AgentFinish).returnValues) {
+          // Agent finish handling (AgentFinish contains returnValues property)
+          if ('returnValues' in nextStepOutput) {
             logger.info(`✅ Agent finished at step ${stepNum + 1}`)
-            result = (nextStepOutput as AgentFinish).returnValues?.output ?? 'No output generated'
+            result = nextStepOutput.returnValues?.output ?? 'No output generated'
+            runManager?.handleChainEnd({ output: result })
 
             // If structured output is requested, attempt to create it
             if (outputSchema && structuredLlm) {
@@ -618,10 +633,10 @@ export class MCPAgent {
           // Detect direct return
           if (stepArray.length) {
             const lastStep = stepArray[stepArray.length - 1]
-            const toolReturn = await this._agentExecutor._getToolReturn(lastStep)
+            const toolReturn: AgentFinish | null = await this._agentExecutor._getToolReturn(lastStep)
             if (toolReturn) {
               logger.info(`🏆 Tool returned directly at step ${stepNum + 1}`)
-              result = (toolReturn as unknown as AgentFinish).returnValues?.output ?? 'No output generated'
+              result = toolReturn.returnValues?.output ?? 'No output generated'
               break
             }
           }
@@ -630,11 +645,13 @@ export class MCPAgent {
           if (e instanceof OutputParserException) {
             logger.error(`❌ Output parsing error during step ${stepNum + 1}: ${e}`)
             result = `Agent stopped due to a parsing error: ${e}`
+            runManager?.handleChainError(result)
             break
           }
           logger.error(`❌ Error during agent execution step ${stepNum + 1}: ${e}`)
           console.error(e)
           result = `Agent stopped due to an error: ${e}`
+          runManager?.handleChainError(result)
           break
         }
       }
@@ -643,6 +660,7 @@ export class MCPAgent {
       if (!result) {
         logger.warn(`⚠️ Agent stopped after reaching max iterations (${steps})`)
         result = `Agent stopped after reaching the maximum number of steps (${steps}).`
+        runManager?.handleChainEnd({ output: result })
       }
 
       logger.info('🎉 Agent execution complete')
@@ -802,6 +820,8 @@ export class MCPAgent {
       // Prepare inputs
       const inputs = { input: query, chat_history: langchainHistory }
 
+      logger.info('callbacks', this.callbacks)
+
       // Stream events from the agent executor
       const eventStream = agentExecutor.streamEvents(
         inputs,