agents 0.0.0-c3e8618 → 0.0.0-c8f53b8

package/src/index.ts

@@ -12,11 +12,10 @@ import { parseCronExpression } from "cron-schedule";
 import { nanoid } from "nanoid";
 
 import { AsyncLocalStorage } from "node:async_hooks";
+import { MCPClientManager } from "./mcp/client";
 
 export type { Connection, WSMessage, ConnectionContext } from "partyserver";
 
-import { WorkflowEntrypoint as CFWorkflowEntrypoint } from "cloudflare:workers";
-
 /**
  * RPC request message from client
  */
@@ -119,11 +118,6 @@ export function unstable_callable(metadata: CallableMetadata = {}) {
   };
 }
 
-/**
- * A class for creating workflow entry points that can be used with Cloudflare Workers
- */
-export class WorkflowEntrypoint extends CFWorkflowEntrypoint {}
-
 /**
  * Represents a scheduled task within an Agent
  * @template T Type of the payload data
@@ -170,12 +164,36 @@ const STATE_WAS_CHANGED = "cf_state_was_changed";
 
 const DEFAULT_STATE = {} as unknown;
 
-export const unstable_context = new AsyncLocalStorage<{
+const agentContext = new AsyncLocalStorage<{
   agent: Agent<unknown>;
   connection: Connection | undefined;
   request: Request | undefined;
 }>();
 
+export function getCurrentAgent<
+  T extends Agent<unknown, unknown> = Agent<unknown, unknown>,
+>(): {
+  agent: T | undefined;
+  connection: Connection | undefined;
+  request: Request<unknown, CfProperties<unknown>> | undefined;
+} {
+  const store = agentContext.getStore() as
+    | {
+        agent: T;
+        connection: Connection | undefined;
+        request: Request<unknown, CfProperties<unknown>> | undefined;
+      }
+    | undefined;
+  if (!store) {
+    return {
+      agent: undefined,
+      connection: undefined,
+      request: undefined,
+    };
+  }
+  return store;
+}
+
 /**
  * Base class for creating Agent implementations
  * @template Env Environment type containing bindings
@@ -184,6 +202,11 @@ export const unstable_context = new AsyncLocalStorage<{
 export class Agent<Env, State = unknown> extends Server<Env> {
   #state = DEFAULT_STATE as State;
 
+  #ParentClass: typeof Agent<Env, State> =
+    Object.getPrototypeOf(this).constructor;
+
+  mcp: MCPClientManager = new MCPClientManager(this.#ParentClass.name, "0.0.1");
+
   /**
    * Initial state for the Agent
    * Override to provide default state values
@@ -300,7 +323,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
     const _onMessage = this.onMessage.bind(this);
     this.onMessage = async (connection: Connection, message: WSMessage) => {
-      return unstable_context.run(
+      return agentContext.run(
         { agent: this, connection, request: undefined },
         async () => {
           if (typeof message !== "string") {
@@ -378,7 +401,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     this.onConnect = (connection: Connection, ctx: ConnectionContext) => {
       // TODO: This is a hack to ensure the state is sent after the connection is established
       // must fix this
-      return unstable_context.run(
+      return agentContext.run(
         { agent: this, connection, request: ctx.request },
         async () => {
           setTimeout(() => {
@@ -415,8 +438,8 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       source !== "server" ? [source.id] : []
     );
     return this.#tryCatch(() => {
-      const { connection, request } = unstable_context.getStore() || {};
-      return unstable_context.run(
+      const { connection, request } = agentContext.getStore() || {};
+      return agentContext.run(
         { agent: this, connection, request },
         async () => {
           return this.onStateUpdate(state, source);
@@ -447,7 +470,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * @param email Email message to process
    */
   onEmail(email: ForwardableEmailMessage) {
-    return unstable_context.run(
+    return agentContext.run(
       { agent: this, connection: undefined, request: undefined },
       async () => {
         console.error("onEmail not implemented");
@@ -697,7 +720,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         console.error(`callback ${row.callback} not found`);
         continue;
       }
-      await unstable_context.run(
+      await agentContext.run(
         { agent: this, connection: undefined, request: undefined },
         async () => {
           try {
@@ -856,7 +879,7 @@ export async function routeAgentEmail<Env>(
  * @param options Options for Agent creation
  * @returns Promise resolving to an Agent instance stub
  */
-export function getAgentByName<Env, T extends Agent<Env>>(
+export async function getAgentByName<Env, T extends Agent<Env>>(
   namespace: AgentNamespace<T>,
   name: string,
   options?: {