@aigne/core 1.6.0 → 1.7.0

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

@@ -1,5 +1,5 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
+import { SSEClientTransport, } from "@modelcontextprotocol/sdk/client/sse.js";
 import { StdioClientTransport, getDefaultEnvironment, } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { UriTemplate } from "@modelcontextprotocol/sdk/shared/uriTemplate.js";
 import pRetry from "p-retry";
@@ -11,7 +11,6 @@ import { Agent } from "./agent.js";
 const MCP_AGENT_CLIENT_NAME = "MCPAgent";
 const MCP_AGENT_CLIENT_VERSION = "0.0.1";
 const DEFAULT_MAX_RECONNECTS = 10;
-const debug = logger.base.extend("mcp");
 function isSSEServerParameters(options) {
     return "url" in options && typeof options.url === "string";
 }
@@ -31,7 +30,7 @@ export class MCPAgent extends Agent {
     static from(options) {
         checkArguments("MCPAgent.from", mcpAgentOptionsSchema, options);
         if (isSSEServerParameters(options)) {
-            const transport = () => new SSEClientTransport(new URL(options.url));
+            const transport = () => new SSEClientTransport(new URL(options.url), options.opts);
             return MCPAgent.fromTransport(transport, options);
         }
         if (isStdioServerParameters(options)) {
@@ -53,29 +52,36 @@ export class MCPAgent extends Agent {
             version: MCP_AGENT_CLIENT_VERSION,
         }, undefined, isSSEServerParameters(options) ? { transportCreator, ...options } : undefined);
         const transport = transportCreator();
-        await debug.spinner(client.connect(transport), `Connecting to MCP server: ${getMCPServerString(options)}`);
+        logger.mcp(`Connecting to MCP server: ${getMCPServerString(options)}`);
+        await client.connect(transport);
         const mcpServer = getMCPServerName(client);
         const { tools: isToolsAvailable, prompts: isPromptsAvailable, resources: isResourcesAvailable, } = client.getServerCapabilities() ?? {};
+        logger.mcp(`Listing tools from ${mcpServer}`);
         const tools = isToolsAvailable
-            ? await debug
-                .spinner(client.listTools(), `Listing tools from ${mcpServer}`, ({ tools }) => debug("%O", tools))
-                .then(({ tools }) => tools.map((tool) => toolFromMCPTool(tool, { client })))
+            ? await client.listTools().then(({ tools }) => {
+                logger.mcp(`Listing tools from ${mcpServer} completed %O`, tools?.map((i) => i.name));
+                return tools.map((tool) => toolFromMCPTool(tool, { client }));
+            })
             : undefined;
+        logger.mcp(`Listing prompts from ${mcpServer}`);
         const prompts = isPromptsAvailable
-            ? await debug
-                .spinner(client.listPrompts(), `Listing prompts from ${mcpServer}`, ({ prompts }) => debug("%O", prompts))
-                .then(({ prompts }) => prompts.map((prompt) => promptFromMCPPrompt(prompt, { client })))
+            ? await client.listPrompts().then(({ prompts }) => {
+                logger.mcp(`Listing prompts from ${mcpServer} completed %O`, prompts?.map((i) => i.name));
+                return prompts.map((prompt) => promptFromMCPPrompt(prompt, { client }));
+            })
             : undefined;
+        logger.mcp(`Listing resources from ${mcpServer}`);
+        // TODO: should conditionally call listResourceTemplates based on the server capabilities
+        // but the capability is not correct in the current SDK version
         const resources = isResourcesAvailable
-            ? await debug
-                .spinner(
-            // TODO: should conditionally call listResourceTemplates based on the server capabilities
-            // but the capability is not correct in the current SDK version
-            Promise.all([
+            ? await Promise.all([
                 client.listResources().catch(() => ({ resources: [] })),
                 client.listResourceTemplates().catch(() => ({ resourceTemplates: [] })),
-            ]), `Listing resources from ${mcpServer}`, ([{ resources }, { resourceTemplates }]) => debug("%O\n%O", resources, resourceTemplates))
-                .then(([{ resources }, { resourceTemplates }]) => [...resources, ...resourceTemplates].map((resource) => resourceFromMCPResource(resource, { client })))
+            ]).then(([{ resources }, { resourceTemplates }]) => {
+                const result = [...resources, ...resourceTemplates].map((resource) => resourceFromMCPResource(resource, { client }));
+                logger.mcp(`Listing resources from ${mcpServer} completed %O`, result.map((i) => i.name));
+                return result;
+            })
             : undefined;
         return new MCPAgent({
             name: client.getServerVersion()?.name,
@@ -131,18 +137,22 @@ class ClientWithReconnect extends Client {
         }, {
             retries: this.reconnectOptions?.maxReconnects ?? DEFAULT_MAX_RECONNECTS,
             shouldRetry: this.shouldReconnect,
-            onFailedAttempt: (error) => debug("Reconnect attempt failed: %O", error),
+            onFailedAttempt: (error) => logger.mcp("Reconnect attempt failed: %O", error),
         });
     }
     async request(request, resultSchema, options) {
+        const mergedOptions = {
+            ...(options ?? {}),
+            timeout: options?.timeout ?? this.reconnectOptions?.timeout ?? 10000,
+        };
         try {
-            return await super.request(request, resultSchema, options);
+            return await super.request(request, resultSchema, mergedOptions);
         }
         catch (error) {
             if (this.shouldReconnect(error)) {
-                debug("Error occurred, reconnecting to MCP server: %O", error);
+                logger.mcp("Error occurred, reconnecting to MCP server: %O", error);
                 await this.reconnect();
-                return await super.request(request, resultSchema, options);
+                return await super.request(request, resultSchema, mergedOptions);
             }
             throw error;
         }
@@ -160,13 +170,13 @@ export class MCPBase extends Agent {
 }
 export class MCPTool extends MCPBase {
     async process(input) {
-        const result = await debug.spinner(this.client.callTool({ name: this.name, arguments: input }), `Call tool ${this.name} from ${this.mcpServer}`, (output) => debug("input: %O\noutput: %O", input, output));
+        const result = await this.client.callTool({ name: this.name, arguments: input });
         return result;
     }
 }
 export class MCPPrompt extends MCPBase {
     async process(input) {
-        const result = await debug.spinner(this.client.getPrompt({ name: this.name, arguments: input }), `Get prompt ${this.name} from ${this.mcpServer}`, (output) => debug("input: %O\noutput: %O", input, output));
+        const result = await this.client.getPrompt({ name: this.name, arguments: input });
         return result;
     }
 }
@@ -178,7 +188,7 @@ export class MCPResource extends MCPBase {
     uri;
     async process(input) {
         const uri = new UriTemplate(this.uri).expand(input);
-        const result = await debug.spinner(this.client.readResource({ uri }), `Read resource ${this.name} from ${this.mcpServer}`, (output) => debug("input: %O\noutput: %O", input, output));
+        const result = await this.client.readResource({ uri });
         return result;
     }
 }
@@ -197,6 +207,10 @@ const mcpAgentOptionsSchema = z.union([
     }),
     z.object({
         url: z.string(),
+        opts: z.object({}).optional(),
+        timeout: z.number().optional(),
+        maxReconnects: z.number().optional(),
+        shouldReconnect: z.function().args(z.instanceof(Error)).returns(z.boolean()).optional(),
     }),
     z.object({
         command: z.string(),

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

@@ -1,26 +1,26 @@
 import { ReadableStream } from "node:stream/web";
-import type { Context } from "../execution-engine/context.js";
-import type { MessagePayload, MessageQueueListener, Unsubscribe } from "../execution-engine/message-queue.js";
+import { type Context, type Runnable } from "../execution-engine/context.js";
+import type { MessagePayload } from "../execution-engine/message-queue.js";
 import { type PromiseOrValue } from "../utils/type-utils.js";
 import { Agent, type AgentOptions, type Message } from "./agent.js";
 export interface UserAgentOptions<I extends Message = Message, O extends Message = Message> extends AgentOptions<I, O> {
     context: Context;
     process?: (input: I, context: Context) => PromiseOrValue<O>;
+    activeAgent?: Runnable;
 }
 export declare class UserAgent<I extends Message = Message, O extends Message = Message> extends Agent<I, O> {
     static from<I extends Message, O extends Message>(options: UserAgentOptions<I, O>): UserAgent<I, O>;
     constructor(options: UserAgentOptions<I, O>);
-    readonly context: Context;
-    private get ctx();
+    context: Context;
     private _process?;
+    private activeAgent?;
     call(input: string | I, context?: Context): Promise<O>;
     process(input: I, context: Context): Promise<O>;
-    publish(topic: string | string[], message: Message | string): void;
-    subscribe(topic: string, listener?: undefined): Promise<MessagePayload>;
-    subscribe(topic: string, listener: MessageQueueListener): Unsubscribe;
-    subscribe(topic: string, listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
-    unsubscribe(topic: string, listener: MessageQueueListener): void;
+    publish: Context["publish"];
+    subscribe: Context["subscribe"];
+    unsubscribe: Context["unsubscribe"];
     get stream(): ReadableStream<MessagePayload & {
         topic: string;
     }>;
+    protected checkUsageAgentCalls(_context: Context): void;
 }

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

@@ -1,4 +1,5 @@
 import { ReadableStream } from "node:stream/web";
+import { createPublishMessage } from "../execution-engine/context.js";
 import { orArrayToArray } from "../utils/type-utils.js";
 import { Agent } from "./agent.js";
 export class UserAgent extends Agent {
@@ -6,40 +7,46 @@ export class UserAgent extends Agent {
         return new UserAgent(options);
     }
     constructor(options) {
-        super({ ...options, disableLogging: true });
+        super({ ...options, disableEvents: true });
         this._process = options.process;
         this.context = options.context;
+        this.activeAgent = options.activeAgent;
     }
     context;
-    get ctx() {
-        if (!this.context)
-            throw new Error("UserAgent must have a context");
-        return this.context;
-    }
     _process;
+    activeAgent;
     call(input, context) {
+        if (!context)
+            this.context = this.context.newContext({ reset: true });
         return super.call(input, context ?? this.context);
     }
     async process(input, context) {
         if (this._process) {
             return this._process(input, context);
         }
+        if (this.activeAgent) {
+            const [output, agent] = await context.call(this.activeAgent, input, {
+                returnActiveAgent: true,
+            });
+            this.activeAgent = agent;
+            return output;
+        }
         const publicTopic = typeof this.publishTopic === "function" ? await this.publishTopic(input) : this.publishTopic;
         if (publicTopic?.length) {
-            context.publish(publicTopic, input, this);
+            context.publish(publicTopic, createPublishMessage(input, this));
             return {};
         }
         throw new Error("UserAgent must have a process function or a publishTopic");
     }
-    publish(topic, message) {
-        return this.ctx.publish(topic, message, this);
-    }
-    subscribe(topic, listener) {
-        return this.ctx.subscribe(topic, listener);
-    }
-    unsubscribe(topic, listener) {
-        this.ctx.unsubscribe(topic, listener);
-    }
+    publish = ((...args) => {
+        return this.context.publish(...args);
+    });
+    subscribe = ((...args) => {
+        return this.context.subscribe(...args);
+    });
+    unsubscribe = ((...args) => {
+        this.context.unsubscribe(...args);
+    });
     get stream() {
         let subscriptions = [];
         return new ReadableStream({
@@ -56,4 +63,7 @@ export class UserAgent extends Agent {
             },
         });
     }
+    checkUsageAgentCalls(_context) {
+        // ignore calls usage check for UserAgent
+    }
 }

package/lib/esm/execution-engine/context.d.ts

@@ -2,19 +2,36 @@ import EventEmitter from "node:events";
 import { Agent, type FunctionAgentFn, type Message } from "../agents/agent.js";
 import { UserAgent } from "../agents/user-agent.js";
 import type { ChatModel } from "../models/chat-model.js";
+import { type OmitPropertiesFromArrayFirstElement } from "../utils/type-utils.js";
+import type { Args, Listener, TypedEventEmitter } from "../utils/typed-event-emtter.js";
 import { type MessagePayload, MessageQueue, type MessageQueueListener, type Unsubscribe } from "./message-queue.js";
+import { type ContextLimits, type ContextUsage } from "./usage.js";
 export type Runnable<I extends Message = Message, O extends Message = Message> = Agent<I, O> | FunctionAgentFn;
-export interface ContextUsage {
-    promptTokens: number;
-    completionTokens: number;
-    agentCalls: number;
+export interface AgentEvent {
+    parentContextId?: string;
+    contextId: string;
+    timestamp: number;
+    agent: Agent;
 }
-export interface ContextLimits {
-    maxTokens?: number;
-    maxAgentCalls?: number;
-    timeout?: number;
+export interface ContextEventMap {
+    agentStarted: [AgentEvent & {
+        input: Message;
+    }];
+    agentSucceed: [AgentEvent & {
+        output: Message;
+    }];
+    agentFailed: [AgentEvent & {
+        error: Error;
+    }];
 }
-export interface Context extends EventEmitter {
+export type ContextEmitEventMap = {
+    [K in keyof ContextEventMap]: OmitPropertiesFromArrayFirstElement<ContextEventMap[K], "contextId" | "parentContextId" | "timestamp">;
+};
+export interface CallOptions {
+    returnActiveAgent?: boolean;
+    disableTransfer?: boolean;
+}
+export interface Context extends TypedEventEmitter<ContextEventMap, ContextEmitEventMap> {
     model?: ChatModel;
     tools?: Agent[];
     usage: ContextUsage;
@@ -26,13 +43,6 @@ export interface Context extends EventEmitter {
      * @returns User agent
      */
     call<I extends Message, O extends Message>(agent: Runnable<I, O>): UserAgent<I, O>;
-    /**
-     * Call an agent with a message
-     * @param agent Agent to call
-     * @param message Message to pass to the agent
-     * @returns the output of the agent
-     */
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string): Promise<O>;
     /**
      * Call an agent with a message and return the output and the active agent
      * @param agent Agent to call
@@ -40,54 +50,82 @@ export interface Context extends EventEmitter {
      * @param options.returnActiveAgent return the active agent
      * @returns the output of the agent and the final active agent
      */
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: {
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: CallOptions & {
         returnActiveAgent: true;
     }): Promise<[O, Runnable]>;
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options?: {
-        returnActiveAgent?: boolean;
-    }): UserAgent<I, O> | Promise<O | [O, Runnable]>;
+    /**
+     * Call an agent with a message
+     * @param agent Agent to call
+     * @param message Message to pass to the agent
+     * @returns the output of the agent
+     */
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options?: CallOptions): Promise<O>;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message?: I | string, options?: CallOptions): UserAgent<I, O> | Promise<O | [O, Runnable]>;
     /**
      * Publish a message to a topic, the engine will call the listeners of the topic
      * @param topic topic name, or an array of topic names
-     * @param message message to publish
-     * @param from the agent who publish the message, if not provided, it will be treated as a user message
+     * @param payload message to publish
      */
-    publish(topic: string | string[], message: Message | string, from?: Agent): void;
+    publish(topic: string | string[], payload: Omit<MessagePayload, "context">): void;
     subscribe(topic: string, listener?: undefined): Promise<MessagePayload>;
     subscribe(topic: string, listener: MessageQueueListener): Unsubscribe;
     subscribe(topic: string, listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
+    subscribe(topic: string, listener?: MessageQueueListener): Unsubscribe | Promise<MessagePayload>;
     unsubscribe(topic: string, listener: MessageQueueListener): void;
+    /**
+     * Create a child context with the same configuration as the parent context.
+     * If `reset` is true, the child context will have a new state (such as: usage).
+     *
+     * @param options
+     * @param options.reset create a new context with initial state (such as: usage)
+     * @returns new context
+     */
+    newContext(options?: {
+        reset?: boolean;
+    }): Context;
+}
+export declare function createPublishMessage(message: string | Message, from?: Agent): Omit<MessagePayload, "context">;
+export declare class ExecutionContext implements Context {
+    constructor(parent?: ConstructorParameters<typeof ExecutionContextInternal>[0]);
+    parentId?: string;
+    id: string;
+    readonly internal: ExecutionContextInternal;
+    get model(): ChatModel | undefined;
+    get tools(): Agent<Message, Message>[] | undefined;
+    get limits(): ContextLimits | undefined;
+    get status(): "normal" | "timeout";
+    get usage(): ContextUsage;
+    newContext({ reset }?: {
+        reset?: boolean;
+    }): ExecutionContext;
+    call: Context["call"];
+    publish: Context["publish"];
+    subscribe: Context["subscribe"];
+    unsubscribe: Context["unsubscribe"];
+    emit<K extends keyof ContextEmitEventMap>(eventName: K, ...args: Args<K, ContextEmitEventMap>): boolean;
+    on<K extends keyof ContextEventMap>(eventName: K, listener: Listener<K, ContextEventMap>): this;
+    once<K extends keyof ContextEventMap>(eventName: K, listener: Listener<K, ContextEventMap>): this;
+    off<K extends keyof ContextEventMap>(eventName: K, listener: Listener<K, ContextEventMap>): this;
 }
-export declare class ExecutionContext extends EventEmitter implements Context {
-    private readonly engine?;
-    constructor(engine?: {
-        model?: ChatModel;
-        tools?: Agent[];
-        limits?: ContextLimits;
+declare class ExecutionContextInternal {
+    private readonly parent?;
+    constructor(parent?: (Pick<Context, "model" | "tools" | "limits"> & {
         messageQueue?: MessageQueue;
-        emit?: EventEmitter["emit"];
-    } | undefined);
-    private messageQueue;
+    }) | undefined);
+    readonly messageQueue: MessageQueue;
+    readonly events: EventEmitter<ContextEventMap>;
     get model(): ChatModel | undefined;
     get tools(): Agent<Message, Message>[] | undefined;
+    get limits(): ContextLimits | undefined;
     usage: ContextUsage;
-    limits?: ContextLimits | undefined;
     private abortController;
     private timer?;
     private initTimeout;
     get status(): "normal" | "timeout";
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>): UserAgent<I, O>;
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string): Promise<O>;
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message: I | string, options: {
-        returnActiveAgent: true;
-    }): Promise<[O, Runnable]>;
-    call<I extends Message, O extends Message>(agent: Runnable<I, O>, message?: I | string, options?: {
-        returnActiveAgent?: boolean;
-    }): UserAgent<I, O> | Promise<O | [O, Runnable]>;
-    publish(topic: string | string[], message: Message | string, from?: Agent): void;
-    subscribe(topic: string, listener?: undefined): Promise<MessagePayload>;
-    subscribe(topic: string, listener: MessageQueueListener): Unsubscribe;
-    unsubscribe(topic: string, listener: MessageQueueListener): void;
-    emit(eventName: string | symbol, ...args: unknown[]): boolean;
+    call<I extends Message, O extends Message>(agent: Runnable<I, O>, input: I, context: Context, options?: CallOptions): Promise<{
+        agent: Runnable;
+        output: O;
+    }>;
     private callAgent;
 }
+export {};