@iqai/adk 0.1.20 → 0.1.22

package/dist/index.mjs

@@ -53,7 +53,7 @@ var init_logger = __esm({
       }
       info(message, ...args) {
         const time = (/* @__PURE__ */ new Date()).toLocaleTimeString();
-        console.info(
+        console.debug(
           this.colorize(`[${time}] \u2139\uFE0F [${this.name}] ${message}`),
           ...args
         );
@@ -229,7 +229,7 @@ var init_base_tool = __esm({
        * @param context The context of the tool
        * @returns The result of running the tool
        */
-      async runAsync(args, context) {
+      async runAsync(args, context4) {
         throw new Error(`${this.constructor.name} runAsync is not implemented`);
       }
       /**
@@ -253,6 +253,12 @@ var init_base_tool = __esm({
           if (!toolWithFunctionDeclarations.functionDeclarations) {
             toolWithFunctionDeclarations.functionDeclarations = [];
           }
+          const alreadyExists = toolWithFunctionDeclarations.functionDeclarations.some(
+            (fd) => fd?.name === functionDeclaration.name
+          );
+          if (alreadyExists) {
+            return;
+          }
           toolWithFunctionDeclarations.functionDeclarations.push(
             functionDeclaration
           );
@@ -281,7 +287,7 @@ var init_base_tool = __esm({
        * @param context Tool execution context
        * @returns Result of the tool execution or error information
        */
-      async safeExecute(args, context) {
+      async safeExecute(args, context4) {
         if (!this.validateArguments(args)) {
           return {
             error: "Invalid arguments",
@@ -302,7 +308,7 @@ var init_base_tool = __esm({
               );
               await new Promise((resolve) => setTimeout(resolve, delay));
             }
-            const result = await this.runAsync(args, context);
+            const result = await this.runAsync(args, context4);
             return { result };
           } catch (error) {
             lastError = error instanceof Error ? error : new Error(String(error));
@@ -500,7 +506,7 @@ var init_function_tool = __esm({
       /**
        * Executes the wrapped function with the provided arguments.
        */
-      async runAsync(args, context) {
+      async runAsync(args, context4) {
         try {
           const missingArgs = this.getMissingMandatoryArgs(args);
           if (missingArgs.length > 0) {
@@ -513,13 +519,13 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
           }
           const argsToCall = { ...args };
           if (this.functionAcceptsToolContext()) {
-            argsToCall.toolContext = context;
+            argsToCall.toolContext = context4;
           }
           const funcParams = this.getFunctionParameters();
           const argValues = [];
           for (const paramName of funcParams) {
             if (paramName === "toolContext" && this.functionAcceptsToolContext()) {
-              argValues.push(context);
+              argValues.push(context4);
             } else if (paramName in argsToCall) {
               const convertedValue = this.convertArgumentType(
                 argsToCall[paramName],
@@ -954,6 +960,7 @@ init_logger();
 import {
   DiagConsoleLogger,
   DiagLogLevel,
+  context,
   diag,
   trace
 } from "@opentelemetry/api";
@@ -994,13 +1001,24 @@ var TelemetryService = class {
     this.sdk = new NodeSDK({
       resource,
       traceExporter,
-      instrumentations: [getNodeAutoInstrumentations()]
+      instrumentations: [
+        getNodeAutoInstrumentations({
+          // Follow Python ADK approach: let all HTTP instrumentation through.
+          // This provides transparency and aligns with standard OpenTelemetry behavior.
+          // High-level LLM tracing is provided through dedicated ADK spans.
+          "@opentelemetry/instrumentation-http": {
+            ignoreIncomingRequestHook: (req) => {
+              return true;
+            }
+          }
+        })
+      ]
     });
     try {
       this.sdk.start();
       this.isInitialized = true;
       this.tracer = trace.getTracer("iqai-adk", config.appVersion || "0.1.0");
-      diag.info("OpenTelemetry SDK started successfully.");
+      diag.debug("OpenTelemetry SDK started successfully.");
     } catch (error) {
       diag.error("Error starting OpenTelemetry SDK:", error);
       throw error;
@@ -1043,7 +1061,7 @@ var TelemetryService = class {
       });
       await Promise.race([this.sdk.shutdown(), timeoutPromise]);
       this.isInitialized = false;
-      diag.info("Telemetry terminated successfully.");
+      diag.debug("Telemetry terminated successfully.");
     } catch (error) {
       if (error instanceof Error && error.message.includes("timeout")) {
         diag.warn("Telemetry shutdown timed out, some traces may be lost");
@@ -1071,7 +1089,7 @@ var TelemetryService = class {
       }
     }
     span.setAttributes({
-      "gen_ai.system.name": "iqai-adk",
+      "gen_ai.system": "iqai-adk",
       "gen_ai.operation.name": "execute_tool",
       "gen_ai.tool.name": tool.name,
       "gen_ai.tool.description": tool.description,
@@ -1085,7 +1103,7 @@ var TelemetryService = class {
       ...process.env.NODE_ENV && {
         "deployment.environment.name": process.env.NODE_ENV
       },
-      // Tool-specific data
+      // ADK-specific attributes (matching Python namespace pattern)
       "adk.tool_call_args": this._safeJsonStringify(args),
       "adk.event_id": functionResponseEvent.invocationId,
       "adk.tool_response": this._safeJsonStringify(toolResponse),
@@ -1101,9 +1119,8 @@ var TelemetryService = class {
     if (!span) return;
     const requestData = this._buildLlmRequestForTrace(llmRequest);
     span.setAttributes({
-      // Standard OpenTelemetry attributes
-      "gen_ai.system.name": "iqai-adk",
-      "gen_ai.operation.name": "generate",
+      // Standard OpenTelemetry attributes (following Python pattern)
+      "gen_ai.system": "iqai-adk",
       "gen_ai.request.model": llmRequest.model,
       // Session and user tracking (maps to Langfuse sessionId, userId)
       "session.id": invocationContext.session.id,
@@ -1116,15 +1133,21 @@ var TelemetryService = class {
       "gen_ai.request.max_tokens": llmRequest.config.maxOutputTokens || 0,
       "gen_ai.request.temperature": llmRequest.config.temperature || 0,
       "gen_ai.request.top_p": llmRequest.config.topP || 0,
-      // Legacy ADK attributes (keep for backward compatibility)
       "adk.system_name": "iqai-adk",
       "adk.request_model": llmRequest.model,
-      "adk.invocation_id": invocationContext.session.id,
+      // ADK-specific attributes (matching Python namespace pattern)
+      "adk.invocation_id": invocationContext.invocationId,
       "adk.session_id": invocationContext.session.id,
       "adk.event_id": eventId,
       "adk.llm_request": this._safeJsonStringify(requestData),
       "adk.llm_response": this._safeJsonStringify(llmResponse)
     });
+    if (llmResponse.usageMetadata) {
+      span.setAttributes({
+        "gen_ai.usage.input_tokens": llmResponse.usageMetadata.promptTokenCount || 0,
+        "gen_ai.usage.output_tokens": llmResponse.usageMetadata.candidatesTokenCount || 0
+      });
+    }
     span.addEvent("gen_ai.content.prompt", {
       "gen_ai.prompt": this._safeJsonStringify(requestData.messages)
     });
@@ -1137,9 +1160,14 @@ var TelemetryService = class {
    */
   async *traceAsyncGenerator(spanName, generator) {
     const span = this.tracer.startSpan(spanName);
+    const spanContext = trace.setSpan(context.active(), span);
     try {
-      for await (const item of generator) {
-        yield item;
+      while (true) {
+        const result = await context.with(spanContext, () => generator.next());
+        if (result.done) {
+          break;
+        }
+        yield result.value;
       }
     } catch (error) {
       span.recordException(error);
@@ -1226,7 +1254,7 @@ var traceLlmCall = (invocationContext, eventId, llmRequest, llmResponse) => tele
 // src/models/base-llm.ts
 var BaseLlm = class {
   /**
-   * The name of the LLM, e.g. gemini-1.5-flash or gemini-1.5-flash-001.
+   * The name of the LLM, e.g. gemini-2.5-flash or gemini-2.5-flash-001.
    */
   model;
   logger = new Logger({ name: "BaseLlm" });
@@ -1353,7 +1381,7 @@ var BaseLlm = class {
    * @param llmRequest LlmRequest, the request to send to the LLM.
    * @returns BaseLLMConnection, the connection to the LLM.
    */
-  connect(llmRequest) {
+  connect(_llmRequest) {
     throw new Error(`Live connection is not supported for ${this.model}.`);
   }
 };
@@ -1915,7 +1943,7 @@ var GoogleLlm = class extends BaseLlm {
   /**
    * Constructor for Gemini
    */
-  constructor(model = "gemini-1.5-flash") {
+  constructor(model = "gemini-2.5-flash") {
     super(model);
   }
   /**
@@ -3954,10 +3982,10 @@ var CreatedTool = class extends BaseTool {
   /**
    * Executes the tool function with validation
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     try {
       const validatedArgs = this.schema.parse(args);
-      const result = await Promise.resolve(this.func(validatedArgs, context));
+      const result = await Promise.resolve(this.func(validatedArgs, context4));
       return result ?? {};
     } catch (error) {
       if (error instanceof z.ZodError) {
@@ -4215,7 +4243,7 @@ var AgentTool = class extends BaseTool {
   /**
    * Execute the tool by running the agent with the provided input
    */
-  async runAsync(params, context) {
+  async runAsync(params, context4) {
     try {
       const input = params.input || Object.values(params)[0];
       if (!isLlmAgent(this.agent)) {
@@ -4223,7 +4251,7 @@ var AgentTool = class extends BaseTool {
           `Agent ${this.name} does not support running as a tool`
         );
       }
-      const parentInvocation = context._invocationContext;
+      const parentInvocation = context4._invocationContext;
       const childInvocationContext = new InvocationContext({
         invocationId: uuidv42(),
         agent: this.agent,
@@ -4260,8 +4288,8 @@ var AgentTool = class extends BaseTool {
       } catch {
         toolResult = mergedText;
       }
-      if (this.outputKey && context?.state) {
-        context.state[this.outputKey] = toolResult;
+      if (this.outputKey && context4?.state) {
+        context4.state[this.outputKey] = toolResult;
       }
       return toolResult;
     } catch (error) {
@@ -4809,9 +4837,9 @@ var UserInteractionTool = class extends BaseTool {
   /**
    * Execute the user interaction
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     try {
-      const actions = context.actions;
+      const actions = context4.actions;
       if (!actions || !actions.promptUser) {
         return {
           success: false,
@@ -4859,9 +4887,9 @@ var ExitLoopTool = class extends BaseTool {
   /**
    * Execute the exit loop action
    */
-  async runAsync(_args, context) {
+  async runAsync(_args, context4) {
     this.logger.debug("Executing exit loop tool");
-    context.actions.escalate = true;
+    context4.actions.escalate = true;
   }
 };
 
@@ -4912,14 +4940,14 @@ var GetUserChoiceTool = class extends BaseTool {
    * This is a long running operation that will return null initially
    * and the actual choice will be provided asynchronously
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     this.logger.debug(
       `Executing get_user_choice with options: ${args.options.join(", ")}`
     );
     if (args.question) {
       this.logger.debug(`Question: ${args.question}`);
     }
-    context.actions.skipSummarization = true;
+    context4.actions.skipSummarization = true;
     return null;
   }
 };
@@ -4961,9 +4989,9 @@ var TransferToAgentTool = class extends BaseTool {
   /**
    * Execute the transfer to agent action
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     this.logger.debug(`Executing transfer to agent: ${args.agent_name}`);
-    context.actions.transferToAgent = args.agent_name;
+    context4.actions.transferToAgent = args.agent_name;
   }
 };
 
@@ -5004,10 +5032,10 @@ var LoadMemoryTool = class extends BaseTool {
   /**
    * Execute the memory loading action
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     this.logger.debug(`Executing load_memory with query: ${args.query}`);
     try {
-      const searchResult = await context.searchMemory(args.query);
+      const searchResult = await context4.searchMemory(args.query);
       return {
         memories: searchResult.memories || [],
         count: searchResult.memories?.length || 0
@@ -5057,7 +5085,7 @@ var LoadArtifactsTool = class extends BaseTool {
   /**
    * Execute the load artifacts operation
    */
-  async runAsync(args, context) {
+  async runAsync(args, context4) {
     const artifactNames = args.artifact_names || [];
     return { artifact_names: artifactNames };
   }
@@ -6088,12 +6116,12 @@ var McpToolset = class {
    * Checks if a tool should be included based on the tool filter.
    * Similar to Python's _is_selected method.
    */
-  isSelected(tool, context) {
+  isSelected(tool, context4) {
     if (!this.toolFilter) {
       return true;
     }
     if (typeof this.toolFilter === "function") {
-      return this.toolFilter(tool, context);
+      return this.toolFilter(tool, context4);
     }
     if (Array.isArray(this.toolFilter)) {
       return this.toolFilter.includes(tool.name);
@@ -6146,7 +6174,7 @@ var McpToolset = class {
    * Retrieves tools from the MCP server and converts them to BaseTool instances.
    * Similar to Python's get_tools method.
    */
-  async getTools(context) {
+  async getTools(context4) {
     try {
       if (this.isClosing) {
         throw new McpError(
@@ -6168,7 +6196,7 @@ var McpToolset = class {
       }
       const tools = [];
       for (const mcpTool of toolsResponse.tools) {
-        if (this.isSelected(mcpTool, context)) {
+        if (this.isSelected(mcpTool, context4)) {
           try {
             const tool = await createTool2(mcpTool, client);
             tools.push(tool);
@@ -6205,9 +6233,9 @@ var McpToolset = class {
   /**
    * Refreshes the tool cache by clearing it and fetching tools again
    */
-  async refreshTools(context) {
+  async refreshTools(context4) {
     this.tools = [];
-    return this.getTools(context);
+    return this.getTools(context4);
   }
   /**
    * Closes the connection to the MCP server.
@@ -6251,6 +6279,7 @@ async function getMcpTools(config, toolFilter) {
 }
 
 // src/flows/llm-flows/functions.ts
+import { context as context2, trace as trace2 } from "@opentelemetry/api";
 var AF_FUNCTION_CALL_ID_PREFIX = "adk-";
 var REQUEST_EUC_FUNCTION_CALL_NAME = "adk_request_credential";
 function generateClientFunctionCallId() {
@@ -6340,23 +6369,40 @@ async function handleFunctionCallsAsync(invocationContext, functionCallEvent, to
       toolsDict
     );
     const functionArgs = functionCall.args || {};
-    const functionResponse = await callToolAsync(
-      tool,
-      functionArgs,
-      toolContext
-    );
-    if (tool.isLongRunning) {
+    const tracer2 = telemetryService.getTracer();
+    const span = tracer2.startSpan(`execute_tool ${tool.name}`);
+    const spanContext = trace2.setSpan(context2.active(), span);
+    try {
+      const functionResponse = await context2.with(spanContext, async () => {
+        const result = await callToolAsync(tool, functionArgs, toolContext);
+        if (tool.isLongRunning && !result) {
+          return null;
+        }
+        const functionResponseEvent = buildResponseEvent(
+          tool,
+          result,
+          toolContext,
+          invocationContext
+        );
+        telemetryService.traceToolCall(
+          tool,
+          functionArgs,
+          functionResponseEvent
+        );
+        return { result, event: functionResponseEvent };
+      });
       if (!functionResponse) {
         continue;
       }
+      functionResponseEvents.push(functionResponse.event);
+      span.setStatus({ code: 1 });
+    } catch (error) {
+      span.recordException(error);
+      span.setStatus({ code: 2, message: error.message });
+      throw error;
+    } finally {
+      span.end();
     }
-    const functionResponseEvent = buildResponseEvent(
-      tool,
-      functionResponse,
-      toolContext,
-      invocationContext
-    );
-    functionResponseEvents.push(functionResponseEvent);
   }
   if (!functionResponseEvents.length) {
     return null;
@@ -6456,7 +6502,7 @@ var BaseLlmFlow = class {
   responseProcessors = [];
   logger = new Logger({ name: "BaseLlmFlow" });
   async *runAsync(invocationContext) {
-    this.logger.info(`Agent '${invocationContext.agent.name}' started.`);
+    this.logger.debug(`Agent '${invocationContext.agent.name}' started.`);
     let stepCount = 0;
     while (true) {
       stepCount++;
@@ -6466,7 +6512,7 @@ var BaseLlmFlow = class {
         yield event;
       }
       if (!lastEvent || lastEvent.isFinalResponse()) {
-        this.logger.info(
+        this.logger.debug(
           `Agent '${invocationContext.agent.name}' finished after ${stepCount} steps.`
         );
         break;
@@ -6496,7 +6542,7 @@ var BaseLlmFlow = class {
       yield event;
     }
     if (invocationContext.endInvocation) {
-      this.logger.info("Invocation ended during preprocessing.");
+      this.logger.debug("Invocation ended during preprocessing.");
       return;
     }
     const modelResponseEvent = new Event({
@@ -6611,7 +6657,7 @@ var BaseLlmFlow = class {
         yield functionResponseEvent;
         const transferToAgent = functionResponseEvent.actions?.transferToAgent;
         if (transferToAgent) {
-          this.logger.info(`\u{1F504} Live transfer to agent '${transferToAgent}'`);
+          this.logger.debug(`\u{1F504} Live transfer to agent '${transferToAgent}'`);
           const agentToRun = this._getAgentToRun(
             invocationContext,
             transferToAgent
@@ -6650,7 +6696,7 @@ var BaseLlmFlow = class {
       yield functionResponseEvent;
       const transferToAgent = functionResponseEvent.actions?.transferToAgent;
       if (transferToAgent) {
-        this.logger.info(`\u{1F504} Transferring to agent '${transferToAgent}'`);
+        this.logger.debug(`\u{1F504} Transferring to agent '${transferToAgent}'`);
         const agentToRun = this._getAgentToRun(
           invocationContext,
           transferToAgent
@@ -7074,8 +7120,6 @@ var BasicLlmRequestProcessor = class extends BaseLlmRequestProcessor {
     llmRequest.liveConnectConfig.realtimeInputConfig = runConfig.realtimeInputConfig;
     llmRequest.liveConnectConfig.enableAffectiveDialog = runConfig.enableAffectiveDialog;
     llmRequest.liveConnectConfig.proactivity = runConfig.proactivity;
-    const tools = await agent.canonicalTools();
-    llmRequest.appendTools(tools);
     for await (const _ of []) {
       yield _;
     }
@@ -8088,6 +8132,9 @@ var IdentityLlmRequestProcessor = class extends BaseLlmRequestProcessor {
 };
 var requestProcessor5 = new IdentityLlmRequestProcessor();
 
+// src/flows/llm-flows/instructions.ts
+import { zodToJsonSchema as zodToJsonSchema2 } from "zod-to-json-schema";
+
 // src/utils/instructions-utils.ts
 async function injectSessionState(template, readonlyContext) {
   const invocationContext = readonlyContext._invocationContext;
@@ -8202,6 +8249,22 @@ var InstructionsLlmRequestProcessor = class extends BaseLlmRequestProcessor {
       }
       llmRequest.appendInstructions([instruction]);
     }
+    if (agent.outputSchema) {
+      try {
+        const raw = zodToJsonSchema2(agent.outputSchema, {
+          target: "jsonSchema7",
+          $refStrategy: "none"
+        });
+        const { $schema, ...json } = raw || {};
+        llmRequest.appendInstructions([
+          "You must respond with application/json that validates against this JSON Schema:",
+          "```json",
+          JSON.stringify(json, null, 2),
+          "```"
+        ]);
+      } catch {
+      }
+    }
     for await (const _ of []) {
       yield _;
     }
@@ -8493,12 +8556,87 @@ function removeThoughtFromRequest(llmRequest) {
 var requestProcessor7 = new NlPlanningRequestProcessor();
 var responseProcessor2 = new NlPlanningResponseProcessor();
 
+// src/flows/llm-flows/output-schema.ts
+init_logger();
+var OutputSchemaResponseProcessor = class extends BaseLlmResponseProcessor {
+  logger = new Logger({ name: "OutputSchemaResponseProcessor" });
+  async *runAsync(invocationContext, llmResponse) {
+    if (!llmResponse || !llmResponse.content || !llmResponse.content.parts || llmResponse.content.parts.length === 0) {
+      return;
+    }
+    const agent = invocationContext.agent;
+    if (!("outputSchema" in agent) || !agent.outputSchema) {
+      return;
+    }
+    let textContent = llmResponse.content.parts.map((part) => {
+      if (part && typeof part === "object" && "text" in part) {
+        return part.text || "";
+      }
+      return "";
+    }).join("");
+    if (!textContent.trim()) {
+      return;
+    }
+    try {
+      const parsed = JSON.parse(textContent);
+      const validated = agent.outputSchema.parse(parsed);
+      textContent = JSON.stringify(validated, null, 2);
+      llmResponse.content.parts = llmResponse.content.parts.map((part) => {
+        if (part && typeof part === "object" && "text" in part) {
+          return {
+            ...part,
+            text: textContent
+          };
+        }
+        return part;
+      });
+      this.logger.debug("Output schema validation successful", {
+        agent: agent.name,
+        originalLength: textContent.length,
+        validatedKeys: Object.keys(validated)
+      });
+    } catch (error) {
+      const errorMessage = error instanceof Error ? error.message : String(error);
+      const detailedError = `Output schema validation failed for agent '${agent.name}': ${errorMessage}`;
+      this.logger.error(detailedError, {
+        agent: agent.name,
+        responseContent: textContent.substring(0, 200) + (textContent.length > 200 ? "..." : ""),
+        error: errorMessage
+      });
+      llmResponse.errorCode = "OUTPUT_SCHEMA_VALIDATION_FAILED";
+      llmResponse.errorMessage = detailedError;
+      llmResponse.error = new Error(detailedError);
+      const errorEvent = new Event({
+        id: Event.newId(),
+        invocationId: invocationContext.invocationId,
+        author: agent.name,
+        branch: invocationContext.branch,
+        content: {
+          role: "assistant",
+          parts: [
+            {
+              text: `Error: ${detailedError}`
+            }
+          ]
+        }
+      });
+      errorEvent.errorCode = "OUTPUT_SCHEMA_VALIDATION_FAILED";
+      errorEvent.errorMessage = detailedError;
+      errorEvent.error = new Error(detailedError);
+      yield errorEvent;
+    }
+  }
+};
+var responseProcessor3 = new OutputSchemaResponseProcessor();
+
 // src/flows/llm-flows/shared-memory.ts
 var SharedMemoryRequestProcessor = class extends BaseLlmRequestProcessor {
   async *runAsync(invocationContext, llmRequest) {
     const memoryService = invocationContext.memoryService;
     if (!memoryService) return;
-    const lastUserEvent = invocationContext.session.events.findLast((e) => e.author === "user" && e.content?.parts?.length);
+    const lastUserEvent = invocationContext.session.events.findLast(
+      (e) => e.author === "user" && e.content?.parts?.length
+    );
     if (!lastUserEvent) return;
     const query = (lastUserEvent.content.parts ?? []).map((p) => p.text || "").join(" ");
     const results = await memoryService.searchMemory({
@@ -8557,8 +8695,10 @@ var SingleFlow = class extends BaseLlmFlow {
     this.responseProcessors.push(
       responseProcessor2,
       // Phase 5: NL Planning
+      responseProcessor3,
+      // Phase 6: Output Schema validation and parsing - validates response against agent's output schema
       responseProcessor
-      // Phase 5: Code Execution (placeholder)
+      // Phase 7: Code Execution (placeholder)
     );
     this.logger.debug("SingleFlow initialized with processors");
   }
@@ -8731,12 +8871,26 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * The input schema when agent is used as a tool
    */
   inputSchema;
-  // Schema type - depends on specific implementation
   /**
    * The output schema when agent replies
    */
   outputSchema;
-  // Schema type - depends on specific implementation
+  /**
+   * Callback or list of callbacks to be called before calling the LLM
+   */
+  beforeModelCallback;
+  /**
+   * Callback or list of callbacks to be called after calling the LLM
+   */
+  afterModelCallback;
+  /**
+   * Callback or list of callbacks to be called before calling a tool
+   */
+  beforeToolCallback;
+  /**
+   * Callback or list of callbacks to be called after calling a tool
+   */
+  afterToolCallback;
   logger = new Logger({ name: "LlmAgent" });
   /**
    * Constructor for LlmAgent
@@ -8767,6 +8921,11 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
     this.generateContentConfig = config.generateContentConfig;
     this.inputSchema = config.inputSchema;
     this.outputSchema = config.outputSchema;
+    this.beforeModelCallback = config.beforeModelCallback;
+    this.afterModelCallback = config.afterModelCallback;
+    this.beforeToolCallback = config.beforeToolCallback;
+    this.afterToolCallback = config.afterToolCallback;
+    this.validateOutputSchemaConfig();
   }
   /**
    * The resolved model field as BaseLLM
@@ -8783,13 +8942,15 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
       return new AiSdkLlm(this.model);
     }
     let ancestorAgent = this.parentAgent;
-    while (ancestorAgent !== null) {
+    while (ancestorAgent !== null && ancestorAgent !== void 0) {
       if (ancestorAgent instanceof _LlmAgent) {
         return ancestorAgent.canonicalModel;
       }
       ancestorAgent = ancestorAgent.parentAgent;
     }
-    throw new Error(`No model found for ${this.name}.`);
+    throw new Error(
+      `No model found for agent "${this.name}". Please specify a model directly on this agent using the 'model' property`
+    );
   }
   /**
    * The resolved instruction field to construct instruction for this agent
@@ -8829,6 +8990,80 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
     }
     return resolvedTools;
   }
+  /**
+   * Gets the canonical before model callbacks as an array
+   */
+  get canonicalBeforeModelCallbacks() {
+    if (!this.beforeModelCallback) {
+      return [];
+    }
+    if (Array.isArray(this.beforeModelCallback)) {
+      return this.beforeModelCallback;
+    }
+    return [this.beforeModelCallback];
+  }
+  /**
+   * Gets the canonical after model callbacks as an array
+   */
+  get canonicalAfterModelCallbacks() {
+    if (!this.afterModelCallback) {
+      return [];
+    }
+    if (Array.isArray(this.afterModelCallback)) {
+      return this.afterModelCallback;
+    }
+    return [this.afterModelCallback];
+  }
+  /**
+   * Gets the canonical before tool callbacks as an array
+   */
+  get canonicalBeforeToolCallbacks() {
+    if (!this.beforeToolCallback) {
+      return [];
+    }
+    if (Array.isArray(this.beforeToolCallback)) {
+      return this.beforeToolCallback;
+    }
+    return [this.beforeToolCallback];
+  }
+  /**
+   * Gets the canonical after tool callbacks as an array
+   */
+  get canonicalAfterToolCallbacks() {
+    if (!this.afterToolCallback) {
+      return [];
+    }
+    if (Array.isArray(this.afterToolCallback)) {
+      return this.afterToolCallback;
+    }
+    return [this.afterToolCallback];
+  }
+  /**
+   * Validates output schema configuration
+   * This matches the Python implementation's __check_output_schema
+   */
+  validateOutputSchemaConfig() {
+    if (!this.outputSchema) {
+      return;
+    }
+    if (!this.disallowTransferToParent || !this.disallowTransferToPeers) {
+      this.logger.warn(
+        `Invalid config for agent ${this.name}: output_schema cannot co-exist with agent transfer configurations. Setting disallow_transfer_to_parent=true, disallow_transfer_to_peers=true`
+      );
+      this.disallowTransferToParent = true;
+      this.disallowTransferToPeers = true;
+    }
+    if (this.subAgents && this.subAgents.length > 0) {
+      throw new Error(
+        `Invalid config for agent ${this.name}: if output_schema is set, sub_agents must be empty to disable agent transfer.`
+      );
+    }
+    if (this.tools && this.tools.length > 0) {
+      throw new Error(
+        `Invalid config for agent ${this.name}: if output_schema is set, tools must be empty`
+      );
+    }
+  }
   /**
    * Gets the appropriate LLM flow for this agent
    * This matches the Python implementation's _llm_flow property
@@ -8844,8 +9079,28 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * This matches the Python implementation's __maybe_save_output_to_state
    */
   maybeSaveOutputToState(event) {
+    if (event.author !== this.name) {
+      this.logger.debug(
+        `Skipping output save for agent ${this.name}: event authored by ${event.author}`
+      );
+      return;
+    }
     if (this.outputKey && event.isFinalResponse() && event.content?.parts) {
-      const result = event.content.parts.map((part) => part.text || "").join("");
+      let result = event.content.parts.map((part) => part.text || "").join("");
+      if (this.outputSchema) {
+        if (!result.trim()) {
+          return;
+        }
+        try {
+          const parsed = JSON.parse(result);
+          result = this.outputSchema.parse(parsed);
+        } catch (error) {
+          this.logger.error("Failed to validate output with schema:", error);
+          throw new Error(
+            `Output validation failed: ${error instanceof Error ? error.message : String(error)}`
+          );
+        }
+      }
       if (result) {
         if (!event.actions.stateDelta) {
           event.actions.stateDelta = {};
@@ -8858,19 +9113,19 @@ var LlmAgent = class _LlmAgent extends BaseAgent {
    * Core logic to run this agent via text-based conversation
    * This matches the Python implementation's _run_async_impl
    */
-  async *runAsyncImpl(context) {
+  async *runAsyncImpl(context4) {
     this.logger.debug(`Starting LlmAgent execution for "${this.name}"`);
     try {
-      for await (const event of this.llmFlow.runAsync(context)) {
+      for await (const event of this.llmFlow.runAsync(context4)) {
         this.maybeSaveOutputToState(event);
         yield event;
       }
     } catch (error) {
       this.logger.error("Error in LlmAgent execution:", error);
       const errorEvent = new Event({
-        invocationId: context.invocationId,
+        invocationId: context4.invocationId,
         author: this.name,
-        branch: context.branch,
+        branch: context4.branch,
         content: {
           parts: [
             {
@@ -9138,7 +9393,7 @@ var LangGraphAgent = class extends BaseAgent {
   /**
    * Gets the next nodes to execute based on the current node and its result
    */
-  async getNextNodes(currentNode, lastEvent, context) {
+  async getNextNodes(currentNode, lastEvent, context4) {
     if (!currentNode.targets || currentNode.targets.length === 0) {
       return [];
     }
@@ -9150,7 +9405,7 @@ var LangGraphAgent = class extends BaseAgent {
         continue;
       }
       if (targetNode.condition) {
-        const shouldExecute = await targetNode.condition(lastEvent, context);
+        const shouldExecute = await targetNode.condition(lastEvent, context4);
         if (!shouldExecute) {
           this.logger.debug(`Skipping node "${targetName}" due to condition`);
           continue;
@@ -9163,7 +9418,7 @@ var LangGraphAgent = class extends BaseAgent {
   /**
    * Core logic to run this agent via text-based conversation.
    */
-  async *runAsyncImpl(context) {
+  async *runAsyncImpl(context4) {
     this.logger.debug(
       `Starting graph execution from root node "${this.rootNode}"`
     );
@@ -9185,7 +9440,7 @@ var LangGraphAgent = class extends BaseAgent {
       return;
     }
     let stepCount = 0;
-    const nodesToExecute = [{ node: rootNode, context }];
+    const nodesToExecute = [{ node: rootNode, context: context4 }];
     const executedNodes = [];
     let lastEvent = null;
     while (nodesToExecute.length > 0 && stepCount < this.maxSteps) {
@@ -9193,7 +9448,7 @@ var LangGraphAgent = class extends BaseAgent {
       const { node } = nodesToExecute.shift();
       this.logger.debug(`Step ${stepCount}: Executing node "${node.name}"`);
       executedNodes.push(node.name);
-      const childContext = context.createChildContext(node.agent);
+      const childContext = context4.createChildContext(node.agent);
       try {
         const nodeEvents = [];
         for await (const event of node.agent.runAsync(childContext)) {
@@ -9206,7 +9461,7 @@ var LangGraphAgent = class extends BaseAgent {
           events: nodeEvents
         });
         if (lastEvent) {
-          const nextNodes = await this.getNextNodes(node, lastEvent, context);
+          const nextNodes = await this.getNextNodes(node, lastEvent, context4);
           for (const nextNode of nextNodes) {
             nodesToExecute.push({
               node: nextNode,
@@ -9249,8 +9504,8 @@ var LangGraphAgent = class extends BaseAgent {
    * Core logic to run this agent via video/audio-based conversation.
    * For LangGraph, this follows the same execution pattern as text-based.
    */
-  async *runLiveImpl(context) {
-    yield* this.runAsyncImpl(context);
+  async *runLiveImpl(context4) {
+    yield* this.runAsyncImpl(context4);
   }
   /**
    * Gets the execution results from the last run
@@ -9303,7 +9558,7 @@ var LangGraphAgent = class extends BaseAgent {
 import { generateId } from "ai";
 
 // src/runners.ts
-import { SpanStatusCode } from "@opentelemetry/api";
+import { SpanStatusCode, context as context3, trace as trace3 } from "@opentelemetry/api";
 
 // src/agents/run-config.ts
 var StreamingMode = /* @__PURE__ */ ((StreamingMode2) => {
@@ -9963,11 +10218,11 @@ var Runner = class {
     runConfig = new RunConfig()
   }) {
     const span = tracer.startSpan("invocation");
+    const spanContext = trace3.setSpan(context3.active(), span);
     try {
-      const session = await this.sessionService.getSession(
-        this.appName,
-        userId,
-        sessionId
+      const session = await context3.with(
+        spanContext,
+        () => this.sessionService.getSession(this.appName, userId, sessionId)
       );
       if (!session) {
         throw new Error(`Session not found: ${sessionId}`);
@@ -9977,22 +10232,34 @@ var Runner = class {
         runConfig
       });
       if (newMessage) {
-        await this._appendNewMessageToSession(
-          session,
-          newMessage,
-          invocationContext,
-          runConfig.saveInputBlobsAsArtifacts || false
+        await context3.with(
+          spanContext,
+          () => this._appendNewMessageToSession(
+            session,
+            newMessage,
+            invocationContext,
+            runConfig.saveInputBlobsAsArtifacts || false
+          )
         );
       }
       invocationContext.agent = this._findAgentToRun(session, this.agent);
-      for await (const event of invocationContext.agent.runAsync(
-        invocationContext
-      )) {
+      const agentGenerator = invocationContext.agent.runAsync(invocationContext);
+      while (true) {
+        const result = await context3.with(
+          spanContext,
+          () => agentGenerator.next()
+        );
+        if (result.done) {
+          break;
+        }
+        const event = result.value;
         if (!event.partial) {
-          await this.sessionService.appendEvent(session, event);
-          if (this.memoryService) {
-            await this.memoryService.addSessionToMemory(session);
-          }
+          await context3.with(spanContext, async () => {
+            await this.sessionService.appendEvent(session, event);
+            if (this.memoryService) {
+              await this.memoryService.addSessionToMemory(session);
+            }
+          });
         }
         yield event;
       }
@@ -10188,6 +10455,14 @@ var AgentBuilder = class _AgentBuilder {
     this.config.instruction = instruction;
     return this;
   }
+  withInputSchema(schema) {
+    this.config.inputSchema = schema;
+    return this;
+  }
+  withOutputSchema(schema) {
+    this.config.outputSchema = schema;
+    return this;
+  }
   /**
    * Add tools to the agent
    * @param tools Tools to add to the agent
@@ -10394,6 +10669,14 @@ var AgentBuilder = class _AgentBuilder {
     }
     return { agent, runner, session };
   }
+  /**
+   * Type-safe build method for agents with output schemas
+   * Provides better type inference for the ask method return type
+   */
+  async buildWithSchema() {
+    const result = await this.build();
+    return result;
+  }
   /**
    * Quick execution helper - build and run a message
    * @param message Message to send to the agent (string or full message object)
@@ -10428,7 +10711,9 @@ var AgentBuilder = class _AgentBuilder {
           memoryService: this.memoryService,
           artifactService: this.artifactService,
           outputKey: this.config.outputKey,
-          sessionService: this.sessionService
+          sessionService: this.sessionService,
+          inputSchema: this.config.inputSchema,
+          outputSchema: this.config.outputSchema
         });
       }
       case "sequential":
@@ -10487,14 +10772,16 @@ var AgentBuilder = class _AgentBuilder {
     return `app-${this.config.name}`;
   }
   /**
-   * Create enhanced runner with simplified API
+   * Create enhanced runner with simplified API and proper typing
    * @param baseRunner The base runner instance
    * @param session The session instance
    * @returns Enhanced runner with simplified API
    */
   createEnhancedRunner(baseRunner, session) {
     const sessionOptions = this.sessionOptions;
+    const outputSchema = this.config.outputSchema;
     return {
+      __outputSchema: outputSchema,
       async ask(message) {
         const newMessage = typeof message === "string" ? { parts: [{ text: message }] } : typeof message === "object" && "contents" in message ? { parts: message.contents[message.contents.length - 1].parts } : message;
         let response = "";
@@ -10515,6 +10802,18 @@ var AgentBuilder = class _AgentBuilder {
             }
           }
         }
+        if (outputSchema) {
+          try {
+            const parsed = JSON.parse(response);
+            return outputSchema.parse(parsed);
+          } catch (parseError) {
+            try {
+              return outputSchema.parse(response);
+            } catch (validationError) {
+              return response.trim();
+            }
+          }
+        }
         return response.trim();
       },
       runAsync(params) {
@@ -10575,7 +10874,7 @@ var VertexAiSessionService = class extends BaseSessionService {
       path: `reasoningEngines/${reasoningEngineId}/sessions`,
       request_dict: sessionJsonDict
     });
-    console.info("Create Session response", apiResponse);
+    console.debug("Create Session response", apiResponse);
     const createdSessionId = apiResponse.name.split("/").slice(-3, -2)[0];
     const operationId = apiResponse.name.split("/").pop();
     let maxRetryAttempt = 5;