@aigne/core 1.19.0 → 1.20.1

package/lib/esm/agents/agent.d.ts

@@ -2,6 +2,7 @@ import { nodejs } from "@aigne/platform-helpers/nodejs/index.js";
 import { ZodObject, type ZodType } from "zod";
 import type { AgentEvent, Context, UserContext } from "../aigne/context.js";
 import type { MessagePayload } from "../aigne/message-queue.js";
+import type { ContextUsage } from "../aigne/usage.js";
 import type { Memory, MemoryAgent } from "../memory/memory.js";
 import type { MemoryRecorderInput } from "../memory/recorder.js";
 import type { MemoryRetrieverInput } from "../memory/retriever.js";
@@ -12,7 +13,11 @@ export * from "./types.js";
 /**
  * Basic message type that can contain any key-value pairs
  */
-export type Message = Record<string, unknown>;
+export interface Message extends Record<string, unknown> {
+    $meta?: {
+        usage: ContextUsage;
+    };
+}
 /**
  * Topics the agent subscribes to, can be a single topic string or an array of topic strings
  */
@@ -328,7 +333,7 @@ export declare abstract class Agent<I extends Message = Message, O extends Messa
      * Regular mode waits for the agent to complete processing and return the final result,
      * suitable for scenarios where a complete result is needed at once.
      *
-     * @param input Input message to the agent, can be a string or structured object
+     * @param input Input message to the agent
      * @param options Invocation options, must set streaming to false or leave unset
      * @returns Final JSON response
      *
@@ -336,7 +341,7 @@ export declare abstract class Agent<I extends Message = Message, O extends Messa
      * Here's an example of invoking an agent with regular mode:
      * {@includeCode ../../test/agents/agent.test.ts#example-invoke}
      */
-    invoke(input: I | string, options?: Partial<AgentInvokeOptions> & {
+    invoke(input: I, options?: Partial<AgentInvokeOptions> & {
         streaming?: false;
     }): Promise<O>;
     /**
@@ -346,7 +351,7 @@ export declare abstract class Agent<I extends Message = Message, O extends Messa
      * suitable for scenarios requiring real-time progress updates, such as
      * chat bot typing effects.
      *
-     * @param input Input message to the agent, can be a string or structured object
+     * @param input Input message to the agent
      * @param options Invocation options, must set streaming to true for this overload
      * @returns Streaming response object
      *
@@ -354,7 +359,7 @@ export declare abstract class Agent<I extends Message = Message, O extends Messa
      * Here's an example of invoking an agent with streaming response:
      * {@includeCode ../../test/agents/agent.test.ts#example-invoke-streaming}
      */
-    invoke(input: I | string, options: Partial<AgentInvokeOptions> & {
+    invoke(input: I, options: Partial<AgentInvokeOptions> & {
         streaming: true;
     }): Promise<AgentResponseStream<O>>;
     /**
@@ -366,7 +371,7 @@ export declare abstract class Agent<I extends Message = Message, O extends Messa
      * @param options Invocation options
      * @returns Agent response (streaming or regular)
      */
-    invoke(input: I | string, options?: Partial<AgentInvokeOptions>): Promise<AgentResponse<O>>;
+    invoke(input: I, options?: Partial<AgentInvokeOptions>): Promise<AgentResponse<O>>;
     protected invokeSkill<I extends Message, O extends Message>(skill: Agent<I, O>, input: I, options: AgentInvokeOptions): Promise<O>;
     /**
      * Process agent output

package/lib/esm/agents/agent.js

@@ -1,6 +1,5 @@
 import { nodejs } from "@aigne/platform-helpers/nodejs/index.js";
 import { ZodObject, z } from "zod";
-import { createMessage, getMessage } from "../prompt/prompt-builder.js";
 import { logger } from "../utils/logger.js";
 import { agentResponseStreamToObject, asyncGeneratorToReadableStream, isAsyncGenerator, objectToAgentResponseStream, onAgentResponseStreamEnd, } from "../utils/stream-utils.js";
 import { checkArguments, createAccessorArray, isEmpty, orArrayToArray, } from "../utils/type-utils.js";
@@ -273,7 +272,7 @@ export class Agent {
         for (const memory of this.memories) {
             const ms = (await memory.retrieve({
                 ...input,
-                search: typeof input === "string" ? input : input && getMessage(input),
+                search: typeof input.search === "string" ? input.search : JSON.stringify(input.search),
                 limit: input.limit ?? this.maxRetrieveMemoryCount,
             }, options.context)).memories;
             memories.push(...ms);
@@ -300,13 +299,12 @@ export class Agent {
             opts.context.memories.push(...options.memories);
             options.memories = undefined;
         }
-        const message = typeof input === "string" ? createMessage(input) : input;
         logger.debug("Invoke agent %s started with input: %O", this.name, input);
         if (!this.disableEvents)
-            opts.context.emit("agentStarted", { agent: this, input: message });
+            opts.context.emit("agentStarted", { agent: this, input });
         try {
-            await this.hooks.onStart?.({ context: opts.context, input: message });
-            const parsedInput = checkArguments(`Agent ${this.name} input`, this.inputSchema, message);
+            await this.hooks.onStart?.({ context: opts.context, input });
+            const parsedInput = checkArguments(`Agent ${this.name} input`, this.inputSchema, input);
             await this.preprocess(parsedInput, opts);
             this.checkContextStatus(opts);
             let response = await this.process(parsedInput, opts);
@@ -319,23 +317,23 @@ export class Agent {
                     : isAsyncGenerator(response)
                         ? asyncGeneratorToReadableStream(response)
                         : objectToAgentResponseStream(response);
-                return this.checkResponseByGuideRails(message, onAgentResponseStreamEnd(stream, {
+                return this.checkResponseByGuideRails(input, onAgentResponseStreamEnd(stream, {
                     onResult: async (result) => {
                         return await this.processAgentOutput(parsedInput, result, opts);
                     },
                     onError: async (error) => {
-                        return await this.processAgentError(message, error, opts);
+                        return await this.processAgentError(input, error, opts);
                     },
                 }), opts);
             }
-            return await this.checkResponseByGuideRails(message, this.processAgentOutput(parsedInput, response instanceof ReadableStream
+            return await this.checkResponseByGuideRails(input, this.processAgentOutput(parsedInput, response instanceof ReadableStream
                 ? await agentResponseStreamToObject(response)
                 : isAsyncGenerator(response)
                     ? await agentResponseStreamToObject(response)
                     : response, opts), opts);
         }
         catch (error) {
-            throw await this.processAgentError(message, error, opts);
+            throw await this.processAgentError(input, error, opts);
         }
     }
     async invokeSkill(skill, input, options) {

package/lib/esm/agents/ai-agent.d.ts

@@ -12,7 +12,7 @@ import type { GuideRailAgentOutput } from "./guide-rail-agent.js";
  * @template I The input message type the agent accepts
  * @template O The output message type the agent returns
  */
-export interface AIAgentOptions<I extends Message = Message, O extends Message = Message> extends AgentOptions<I, O> {
+export interface AIAgentOptions<InputKey extends string = string, I extends Message & InputMessage<InputKey> = Message & InputMessage<InputKey>, O extends Message = Message> extends AgentOptions<Omit<I, InputKey> & Partial<InputMessage<InputKey>>, O> {
     /**
      * The language model to use for this agent
      *
@@ -26,10 +26,14 @@ export interface AIAgentOptions<I extends Message = Message, O extends Message =
      * more complex prompt templates
      */
     instructions?: string | PromptBuilder;
+    /**
+     * Pick a message from input to use as the user's message
+     */
+    inputKey?: InputKey;
     /**
      * Custom key to use for text output in the response
      *
-     * Defaults to $message if not specified
+     * Defaults to `message` if not specified
      */
     outputKey?: string;
     /**
@@ -106,6 +110,9 @@ export declare const aiAgentToolChoiceSchema: z.ZodUnion<[z.ZodNativeEnum<typeof
 export declare const aiAgentOptionsSchema: ZodObject<{
     [key in keyof AIAgentOptions]: ZodType<AIAgentOptions[key]>;
 }>;
+type InputMessage<K> = K extends string ? {
+    [key in K]: string;
+} : Message;
 /**
  * AI-powered agent that leverages language models
  *
@@ -126,7 +133,7 @@ export declare const aiAgentOptionsSchema: ZodObject<{
  * Basic AIAgent creation:
  * {@includeCode ../../test/agents/ai-agent.test.ts#example-ai-agent-basic}
  */
-export declare class AIAgent<I extends Message = Message, O extends Message = Message> extends Agent<I, O> {
+export declare class AIAgent<InputKey extends string = string, I extends Message & InputMessage<InputKey> = Message & InputMessage<InputKey>, O extends Message = Message> extends Agent<I, O> {
     /**
      * Create an AIAgent with the specified options
      *
@@ -139,13 +146,13 @@ export declare class AIAgent<I extends Message = Message, O extends Message = Me
      * AI agent with custom instructions:
      * {@includeCode ../../test/agents/ai-agent.test.ts#example-ai-agent-instructions}
      */
-    static from<I extends Message, O extends Message>(options: AIAgentOptions<I, O>): AIAgent<I, O>;
+    static from<InputKey extends string, I extends Message & InputMessage<InputKey>, O extends Message>(options: AIAgentOptions<InputKey, I, O>): AIAgent<InputKey, I, O>;
     /**
      * Create an AIAgent instance
      *
      * @param options Configuration options for the AI agent
      */
-    constructor(options: AIAgentOptions<I, O>);
+    constructor(options: AIAgentOptions<InputKey, I, O>);
     /**
      * The language model used by this agent
      *
@@ -163,6 +170,10 @@ export declare class AIAgent<I extends Message = Message, O extends Message = Me
      * {@includeCode ../../test/agents/ai-agent.test.ts#example-ai-agent-prompt-builder}
      */
     instructions: PromptBuilder;
+    /**
+     * Pick a message from input to use as the user's message
+     */
+    inputKey?: InputKey;
     /**
      * Custom key to use for text output in the response
      *
@@ -170,7 +181,7 @@ export declare class AIAgent<I extends Message = Message, O extends Message = Me
      * Setting a custom output key:
      * {@includeCode ../../test/agents/ai-agent.test.ts#example-ai-agent-custom-output-key}
      */
-    outputKey?: string;
+    outputKey: string;
     /**
      * Controls how the agent uses tools during execution
      *
@@ -224,3 +235,4 @@ export declare class AIAgent<I extends Message = Message, O extends Message = Me
      */
     _processRouter(input: I, model: ChatModel, modelInput: ChatModelInput, options: AgentInvokeOptions, toolsMap: Map<string, Agent>): AgentProcessAsyncGenerator<O>;
 }
+export {};

package/lib/esm/agents/ai-agent.js

@@ -1,10 +1,11 @@
 import { z } from "zod";
-import { MESSAGE_KEY, PromptBuilder } from "../prompt/prompt-builder.js";
+import { PromptBuilder } from "../prompt/prompt-builder.js";
 import { AgentMessageTemplate, ToolMessageTemplate } from "../prompt/template.js";
 import { checkArguments, isEmpty } from "../utils/type-utils.js";
 import { Agent, agentOptionsSchema, isAgentResponseDelta, } from "./agent.js";
 import { ChatModel, } from "./chat-model.js";
 import { isTransferAgentOutput } from "./types.js";
+const DEFAULT_OUTPUT_KEY = "message";
 /**
  * Tool choice options for AI agents
  *
@@ -49,6 +50,7 @@ export const aiAgentToolChoiceSchema = z.union([z.nativeEnum(AIAgentToolChoice),
 export const aiAgentOptionsSchema = agentOptionsSchema.extend({
     model: z.instanceof(ChatModel).optional(),
     instructions: z.union([z.string(), z.instanceof(PromptBuilder)]).optional(),
+    inputKey: z.string().optional(),
     outputKey: z.string().optional(),
     toolChoice: aiAgentToolChoiceSchema.optional(),
     memoryAgentsAsTools: z.boolean().optional(),
@@ -96,19 +98,23 @@ export class AIAgent extends Agent {
      * @param options Configuration options for the AI agent
      */
     constructor(options) {
-        super(options);
+        super({ ...options, inputSchema: options.inputSchema });
         checkArguments("AIAgent", aiAgentOptionsSchema, options);
         this.model = options.model;
         this.instructions =
             typeof options.instructions === "string"
                 ? PromptBuilder.from(options.instructions)
                 : (options.instructions ?? new PromptBuilder());
-        this.outputKey = options.outputKey;
+        this.inputKey = options.inputKey;
+        this.outputKey = options.outputKey || DEFAULT_OUTPUT_KEY;
         this.toolChoice = options.toolChoice;
         this.memoryAgentsAsTools = options.memoryAgentsAsTools;
         this.memoryPromptTemplate = options.memoryPromptTemplate;
         if (typeof options.catchToolsError === "boolean")
             this.catchToolsError = options.catchToolsError;
+        if (!this.inputKey && !this.instructions) {
+            throw new Error("AIAgent requires either inputKey or instructions to be set");
+        }
     }
     /**
      * The language model used by this agent
@@ -127,6 +133,10 @@ export class AIAgent extends Agent {
      * {@includeCode ../../test/agents/ai-agent.test.ts#example-ai-agent-prompt-builder}
      */
     instructions;
+    /**
+     * Pick a message from input to use as the user's message
+     */
+    inputKey;
     /**
      * Custom key to use for text output in the response
      *
@@ -191,7 +201,7 @@ export class AIAgent extends Agent {
             return yield* this._processRouter(input, model, modelInput, options, toolsMap);
         }
         const toolCallMessages = [];
-        const outputKey = this.outputKey || MESSAGE_KEY;
+        const outputKey = this.outputKey;
         for (;;) {
             const modelOutput = {};
             const stream = await options.context.invoke(model, { ...modelInput, messages: modelInput.messages.concat(toolCallMessages) }, { streaming: true });
@@ -251,7 +261,7 @@ export class AIAgent extends Agent {
         }
     }
     async onGuideRailError(error) {
-        const outputKey = this.outputKey || MESSAGE_KEY;
+        const outputKey = this.outputKey || DEFAULT_OUTPUT_KEY;
         return {
             [outputKey]: error.reason,
         };

package/lib/esm/agents/guide-rail-agent.d.ts

@@ -59,4 +59,4 @@ export interface GuideRailAgentOutput extends Message {
  * - Monitor and audit agent behavior
  */
 export type GuideRailAgent = Agent<GuideRailAgentInput, GuideRailAgentOutput>;
-export declare const guideRailAgentOptions: AgentOptions<GuideRailAgentInput, GuideRailAgentOutput>;
+export declare const guideRailAgentOptions: AgentOptions<any, GuideRailAgentOutput>;

package/lib/esm/agents/mcp-agent.d.ts

@@ -192,7 +192,7 @@ export declare abstract class MCPBase<I extends Message, O extends Message> exte
 export declare class MCPTool extends MCPBase<Message, CallToolResult> {
     process(input: Message): Promise<CallToolResult>;
 }
-export interface MCPPromptInput extends Message {
+export interface MCPPromptInput extends Record<string, unknown> {
     [key: string]: string;
 }
 export declare class MCPPrompt extends MCPBase<MCPPromptInput, GetPromptResult> {

package/lib/esm/aigne/aigne.d.ts

@@ -129,12 +129,12 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * This overload is useful when you need to track which agent was ultimately responsible for generating the response.
      *
      * @param agent - Target agent to invoke
-     * @param message - Input message to send to the agent (can be a string or a structured message object)
+     * @param message - Input message to send to the agent
      * @param options.returnActiveAgent - Must be true to return the final active agent
      * @param options.streaming - Must be false to return a response stream
      * @returns A promise resolving to a tuple containing the agent's response and the final active agent
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions<U> & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions<U> & {
         returnActiveAgent: true;
         streaming?: false;
     }): Promise<[O, Agent]>;
@@ -143,12 +143,12 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * This overload is useful when you need streaming responses while also tracking which agent provided them.
      *
      * @param agent - Target agent to invoke
-     * @param message - Input message to send to the agent (can be a string or a structured message object)
+     * @param message - Input message to send to the agent
      * @param options.returnActiveAgent - Must be true to return the final active agent
      * @param options.streaming - Must be true to return a response stream
      * @returns A promise resolving to a tuple containing the agent's response stream and a promise for the final agent
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions<U> & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions<U> & {
         returnActiveAgent: true;
         streaming: true;
     }): Promise<[AgentResponseStream<O>, Promise<Agent>]>;
@@ -157,7 +157,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * This is the standard way to invoke an agent when you only need the response.
      *
      * @param agent - Target agent to invoke
-     * @param message - Input message to send to the agent (can be a string or a structured message object)
+     * @param message - Input message to send to the agent
      * @param options - Optional configuration parameters for the invocation
      * @returns A promise resolving to the agent's complete response
      *
@@ -165,7 +165,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * Here's a simple example of how to invoke an agent:
      * {@includeCode ../../test/aigne/aigne.test.ts#example-simple}
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options?: InvokeOptions<U> & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options?: InvokeOptions<U> & {
         returnActiveAgent?: false;
         streaming?: false;
     }): Promise<O>;
@@ -174,7 +174,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * This allows processing the response incrementally as it's being generated.
      *
      * @param agent - Target agent to invoke
-     * @param message - Input message to send to the agent (can be a string or a structured message object)
+     * @param message - Input message to send to the agent
      * @param options - Configuration with streaming enabled to receive incremental response chunks
      * @returns A promise resolving to a stream of the agent's response that can be consumed incrementally
      *
@@ -182,7 +182,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * Here's an example of how to invoke an agent with streaming response:
      * {@includeCode ../../test/aigne/aigne.test.ts#example-streaming}
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions<U> & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions<U> & {
         returnActiveAgent?: false;
         streaming: true;
     }): Promise<AgentResponseStream<O>>;
@@ -196,7 +196,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * @returns Either a UserAgent (when no message provided) or a promise resolving to the agent's response
      * with optional active agent information based on the provided options
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message?: I | string, options?: InvokeOptions<U>): UserAgent<I, O> | Promise<AgentResponse<O> | [AgentResponse<O>, Agent]>;
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message?: I, options?: InvokeOptions<U>): UserAgent<I, O> | Promise<AgentResponse<O> | [AgentResponse<O>, Agent]>;
     /**
      * Publishes a message to the message queue for inter-agent communication.
      * This method broadcasts a message to all subscribers of the specified topic(s).
@@ -210,7 +210,7 @@ export declare class AIGNE<U extends UserContext = UserContext> {
      * Here's an example of how to publish a message:
      * {@includeCode ../../test/aigne/aigne.test.ts#example-publish-message}
      */
-    publish(topic: string | string[], payload: Omit<MessagePayload, "context"> | Message | string, options?: InvokeOptions<U>): void;
+    publish(topic: string | string[], payload: Omit<MessagePayload, "context"> | Message, options?: InvokeOptions<U>): void;
     /**
      * Subscribes to receive the next message on a specific topic.
      * This overload returns a Promise that resolves with the next message published to the topic.

package/lib/esm/aigne/context.d.ts

@@ -42,6 +42,7 @@ export type ContextEmitEventMap = {
 export interface InvokeOptions<U extends UserContext = UserContext> extends Partial<Omit<AgentInvokeOptions<U>, "context">> {
     returnActiveAgent?: boolean;
     returnProgressChunks?: boolean;
+    returnMetadata?: boolean;
     disableTransfer?: boolean;
     sourceAgent?: Agent;
 }
@@ -77,11 +78,11 @@ export interface Context<U extends UserContext = UserContext> extends TypedEvent
      * @param options.streaming return a stream of the output
      * @returns the output of the agent and the final active agent
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions & {
         returnActiveAgent: true;
         streaming?: false;
     }): Promise<[O, Agent]>;
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions & {
         returnActiveAgent: true;
         streaming: true;
     }): Promise<[AgentResponseStream<O>, Promise<Agent>]>;
@@ -91,19 +92,19 @@ export interface Context<U extends UserContext = UserContext> extends TypedEvent
      * @param message Message to pass to the agent
      * @returns the output of the agent
      */
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options?: InvokeOptions & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options?: InvokeOptions & {
         streaming?: false;
     }): Promise<O>;
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I | string, options: InvokeOptions & {
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message: I, options: InvokeOptions & {
         streaming: true;
     }): Promise<AgentResponseStream<O>>;
-    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message?: I | string, options?: InvokeOptions): UserAgent<I, O> | Promise<AgentResponse<O> | [AgentResponse<O>, Agent]>;
+    invoke<I extends Message, O extends Message>(agent: Agent<I, O>, message?: I, options?: InvokeOptions): UserAgent<I, O> | Promise<AgentResponse<O> | [AgentResponse<O>, Agent]>;
     /**
      * Publish a message to a topic, the aigne will invoke the listeners of the topic
      * @param topic topic name, or an array of topic names
      * @param payload message to publish
      */
-    publish(topic: string | string[], payload: Omit<MessagePayload, "context"> | Message | string, options?: InvokeOptions): void;
+    publish(topic: string | string[], payload: Omit<MessagePayload, "context"> | Message, options?: InvokeOptions): void;
     subscribe(topic: string | string[], listener?: undefined): Promise<MessagePayload>;
     subscribe(topic: string | string[], listener: MessageQueueListener): Unsubscribe;
     subscribe(topic: string | string[], listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
@@ -142,6 +143,7 @@ export declare class AIGNEContext implements Context {
         reset?: boolean;
     }): AIGNEContext;
     invoke: Context["invoke"];
+    private onInvocationResult;
     publish: Context["publish"];
     subscribe: Context["subscribe"];
     unsubscribe: Context["unsubscribe"];

package/lib/esm/aigne/context.js

@@ -5,11 +5,10 @@ import { z } from "zod";
 import { Agent, isAgentResponseDelta, isEmptyChunk, } from "../agents/agent.js";
 import { isTransferAgentOutput, transferAgentOutputKey, } from "../agents/types.js";
 import { UserAgent } from "../agents/user-agent.js";
-import { createMessage } from "../prompt/prompt-builder.js";
 import { AgentResponseProgressStream } from "../utils/event-stream.js";
 import { promiseWithResolvers } from "../utils/promise.js";
 import { agentResponseStreamToObject, asyncGeneratorToReadableStream, mergeReadableStreams, onAgentResponseStreamEnd, } from "../utils/stream-utils.js";
-import { checkArguments, isEmpty, isNil, omitBy, } from "../utils/type-utils.js";
+import { checkArguments, isEmpty, isNil, omit, } from "../utils/type-utils.js";
 import { MessageQueue, toMessagePayload, } from "./message-queue.js";
 import { newEmptyContextUsage } from "./usage.js";
 /**
@@ -81,10 +80,10 @@ export class AIGNEContext {
             });
         }
         const newContext = this.newContext();
-        const msg = createMessage(message);
-        return Promise.resolve(newContext.internal.invoke(agent, msg, newContext, options)).then(async (response) => {
+        return Promise.resolve(newContext.internal.invoke(agent, message, newContext, options)).then(async (response) => {
             if (!options?.streaming) {
-                const { __activeAgent__: activeAgent, ...output } = await agentResponseStreamToObject(response);
+                let { __activeAgent__: activeAgent, ...output } = await agentResponseStreamToObject(response);
+                output = await this.onInvocationResult(output, options);
                 if (options?.returnActiveAgent) {
                     return [output, activeAgent];
                 }
@@ -98,13 +97,14 @@ export class AIGNEContext {
                             ...chunk,
                             delta: {
                                 ...chunk.delta,
-                                json: omitBy(chunk.delta.json, (_, k) => k === "__activeAgent__"),
+                                json: omit(chunk.delta.json, "__activeAgent__"),
                             },
                         };
                     }
                 },
-                onResult({ __activeAgent__: activeAgent }) {
-                    activeAgentPromise.resolve(activeAgent);
+                onResult: async (output) => {
+                    activeAgentPromise.resolve(output.__activeAgent__);
+                    return await this.onInvocationResult(output, options);
                 },
             });
             const finalStream = !options.returnProgressChunks
@@ -116,6 +116,18 @@ export class AIGNEContext {
             return finalStream;
         });
     });
+    async onInvocationResult(output, options) {
+        if (!options?.returnMetadata) {
+            return output;
+        }
+        return {
+            ...output,
+            $meta: {
+                ...output.$meta,
+                usage: this.usage,
+            },
+        };
+    }
     publish = ((topic, payload, options) => {
         if (options?.userContext) {
             Object.assign(this.userContext, options.userContext);
@@ -201,46 +213,54 @@ class AIGNEContextShared {
         return withAbortSignal(this.abortController.signal, new Error("AIGNEContext is timeout"), () => this.invokeAgent(agent, input, context, options));
     }
     async *invokeAgent(agent, input, context, options) {
-        let activeAgent = agent;
-        for (;;) {
-            const result = {};
-            if (options?.sourceAgent && activeAgent !== options.sourceAgent) {
-                options.sourceAgent.hooks.onHandoff?.({
-                    context,
-                    source: options.sourceAgent,
-                    target: activeAgent,
-                    input,
-                });
-            }
-            const stream = await activeAgent.invoke(input, { ...options, context, streaming: true });
-            for await (const value of stream) {
-                if (isAgentResponseDelta(value)) {
-                    if (value.delta.json) {
-                        value.delta.json = omitExistsProperties(result, value.delta.json);
-                        Object.assign(result, value.delta.json);
+        const startedAt = Date.now();
+        try {
+            let activeAgent = agent;
+            for (;;) {
+                const result = {};
+                if (options?.sourceAgent && activeAgent !== options.sourceAgent) {
+                    options.sourceAgent.hooks.onHandoff?.({
+                        context,
+                        source: options.sourceAgent,
+                        target: activeAgent,
+                        input,
+                    });
+                }
+                const stream = await activeAgent.invoke(input, { ...options, context, streaming: true });
+                for await (const value of stream) {
+                    if (isAgentResponseDelta(value)) {
+                        if (value.delta.json) {
+                            value.delta.json = omitExistsProperties(result, value.delta.json);
+                            Object.assign(result, value.delta.json);
+                        }
+                        delete value.delta.json?.[transferAgentOutputKey];
                     }
-                    delete value.delta.json?.[transferAgentOutputKey];
+                    if (isEmptyChunk(value))
+                        continue;
+                    yield value;
                 }
-                if (isEmptyChunk(value))
-                    continue;
-                yield value;
-            }
-            if (!options?.disableTransfer) {
-                const transferToAgent = isTransferAgentOutput(result)
-                    ? result[transferAgentOutputKey].agent
-                    : undefined;
-                if (transferToAgent) {
-                    activeAgent = transferToAgent;
-                    continue;
+                if (!options?.disableTransfer) {
+                    const transferToAgent = isTransferAgentOutput(result)
+                        ? result[transferAgentOutputKey].agent
+                        : undefined;
+                    if (transferToAgent) {
+                        activeAgent = transferToAgent;
+                        continue;
+                    }
                 }
+                break;
             }
-            break;
+            yield {
+                delta: {
+                    json: { __activeAgent__: activeAgent },
+                },
+            };
+        }
+        finally {
+            const endedAt = Date.now();
+            const duration = endedAt - startedAt;
+            this.usage.duration += duration;
         }
-        yield {
-            delta: {
-                json: { __activeAgent__: activeAgent },
-            },
-        };
     }
 }
 function omitExistsProperties(result, { ...delta }) {

package/lib/esm/aigne/message-queue.d.ts

@@ -21,7 +21,7 @@ export interface MessagePayload {
 /**
  * @hidden
  */
-export declare function toMessagePayload(payload: Omit<MessagePayload, "context"> | string | Message, options?: Partial<Pick<MessagePayload, "role" | "source">>): Omit<MessagePayload, "context">;
+export declare function toMessagePayload(payload: Omit<MessagePayload, "context"> | Message, options?: Partial<Pick<MessagePayload, "role" | "source">>): Omit<MessagePayload, "context">;
 /**
  * @hidden
  */

package/lib/esm/aigne/message-queue.js

@@ -1,6 +1,5 @@
 import { Emitter } from "strict-event-emitter";
 import { z } from "zod";
-import { createMessage } from "../prompt/prompt-builder.js";
 import { checkArguments, isNil, orArrayToArray } from "../utils/type-utils.js";
 /**
  * @hidden
@@ -24,12 +23,12 @@ function isMessagePayload(payload) {
  */
 export function toMessagePayload(payload, options) {
     if (isMessagePayload(payload)) {
-        return { ...payload, message: createMessage(payload.message), ...options };
+        return { ...payload, ...options };
     }
     return {
         role: options?.role || "user",
         source: options?.source,
-        message: createMessage(payload),
+        message: payload,
     };
 }
 /**

package/lib/esm/aigne/usage.d.ts

@@ -5,6 +5,7 @@ export interface ContextUsage {
     inputTokens: number;
     outputTokens: number;
     agentCalls: number;
+    duration: number;
 }
 /**
  * @hidden

package/lib/esm/aigne/usage.js

@@ -6,5 +6,6 @@ export function newEmptyContextUsage() {
         inputTokens: 0,
         outputTokens: 0,
         agentCalls: 0,
+        duration: 0,
     };
 }