agents 0.0.0-7e0777b → 0.0.0-7f4616c

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (75) hide show
  1. package/README.md +255 -27
  2. package/dist/ai-chat-agent.d.ts +20 -17
  3. package/dist/ai-chat-agent.js +532 -230
  4. package/dist/ai-chat-agent.js.map +1 -1
  5. package/dist/ai-chat-v5-migration-gdyLiTd8.js +155 -0
  6. package/dist/ai-chat-v5-migration-gdyLiTd8.js.map +1 -0
  7. package/dist/ai-chat-v5-migration.d.ts +155 -0
  8. package/dist/ai-chat-v5-migration.js +3 -0
  9. package/dist/ai-react.d.ts +73 -85
  10. package/dist/ai-react.js +261 -199
  11. package/dist/ai-react.js.map +1 -1
  12. package/dist/ai-types-BWW4umHY.d.ts +95 -0
  13. package/dist/ai-types-UZlfLOYP.js +20 -0
  14. package/dist/ai-types-UZlfLOYP.js.map +1 -0
  15. package/dist/ai-types.d.ts +6 -74
  16. package/dist/ai-types.js +3 -1
  17. package/dist/client-CZBVDDoO.js +786 -0
  18. package/dist/client-CZBVDDoO.js.map +1 -0
  19. package/dist/client-CmMi85Sj.d.ts +104 -0
  20. package/dist/client-CrWcaPgn.d.ts +5313 -0
  21. package/dist/client-DjR-lC16.js +117 -0
  22. package/dist/client-DjR-lC16.js.map +1 -0
  23. package/dist/client.d.ts +12 -93
  24. package/dist/client.js +4 -11
  25. package/dist/codemode/ai.d.ts +27 -0
  26. package/dist/codemode/ai.js +151 -0
  27. package/dist/codemode/ai.js.map +1 -0
  28. package/dist/do-oauth-client-provider-B2jr6UNq.js +93 -0
  29. package/dist/do-oauth-client-provider-B2jr6UNq.js.map +1 -0
  30. package/dist/do-oauth-client-provider-CCwGwnrA.d.ts +55 -0
  31. package/dist/index-DpH9o0ao.d.ts +568 -0
  32. package/dist/index-W4JUkafc.d.ts +54 -0
  33. package/dist/index.d.ts +61 -402
  34. package/dist/index.js +7 -22
  35. package/dist/mcp/client.d.ts +4 -783
  36. package/dist/mcp/client.js +3 -9
  37. package/dist/mcp/do-oauth-client-provider.d.ts +2 -41
  38. package/dist/mcp/do-oauth-client-provider.js +3 -7
  39. package/dist/mcp/index.d.ts +73 -82
  40. package/dist/mcp/index.js +836 -772
  41. package/dist/mcp/index.js.map +1 -1
  42. package/dist/mcp/x402.d.ts +34 -0
  43. package/dist/mcp/x402.js +194 -0
  44. package/dist/mcp/x402.js.map +1 -0
  45. package/dist/mcp-BEwaCsxO.d.ts +61 -0
  46. package/dist/observability/index.d.ts +3 -0
  47. package/dist/observability/index.js +7 -0
  48. package/dist/react-LfPKBVtU.d.ts +113 -0
  49. package/dist/react.d.ts +10 -119
  50. package/dist/react.js +183 -110
  51. package/dist/react.js.map +1 -1
  52. package/dist/schedule.d.ts +87 -10
  53. package/dist/schedule.js +46 -21
  54. package/dist/schedule.js.map +1 -1
  55. package/dist/serializable-gtr9YMhp.d.ts +34 -0
  56. package/dist/serializable.d.ts +7 -32
  57. package/dist/serializable.js +1 -1
  58. package/dist/src-L3cHuAag.js +1231 -0
  59. package/dist/src-L3cHuAag.js.map +1 -0
  60. package/package.json +72 -32
  61. package/src/index.ts +1101 -189
  62. package/dist/ai-types.js.map +0 -1
  63. package/dist/chunk-767EASBA.js +0 -106
  64. package/dist/chunk-767EASBA.js.map +0 -1
  65. package/dist/chunk-E3LCYPCB.js +0 -469
  66. package/dist/chunk-E3LCYPCB.js.map +0 -1
  67. package/dist/chunk-NKZZ66QY.js +0 -116
  68. package/dist/chunk-NKZZ66QY.js.map +0 -1
  69. package/dist/chunk-ZRRXJUAA.js +0 -788
  70. package/dist/chunk-ZRRXJUAA.js.map +0 -1
  71. package/dist/client.js.map +0 -1
  72. package/dist/index.js.map +0 -1
  73. package/dist/mcp/client.js.map +0 -1
  74. package/dist/mcp/do-oauth-client-provider.js.map +0 -1
  75. package/dist/serializable.js.map +0 -1
package/src/index.ts CHANGED
@@ -1,3 +1,4 @@
1
+ import type { env } from "cloudflare:workers";
1
2
  import { AsyncLocalStorage } from "node:async_hooks";
2
3
  import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
3
4
  import type { SSEClientTransportOptions } from "@modelcontextprotocol/sdk/client/sse.js";
@@ -6,23 +7,28 @@ import type {
6
7
  Prompt,
7
8
  Resource,
8
9
  ServerCapabilities,
9
- Tool,
10
+ Tool
10
11
  } from "@modelcontextprotocol/sdk/types.js";
11
12
  import { parseCronExpression } from "cron-schedule";
12
13
  import { nanoid } from "nanoid";
14
+ import { EmailMessage } from "cloudflare:email";
13
15
  import {
14
16
  type Connection,
15
17
  type ConnectionContext,
16
- getServerByName,
17
18
  type PartyServerOptions,
18
- routePartykitRequest,
19
19
  Server,
20
20
  type WSMessage,
21
+ getServerByName,
22
+ routePartykitRequest
21
23
  } from "partyserver";
22
24
  import { camelCaseToKebabCase } from "./client";
23
- import { MCPClientManager } from "./mcp/client";
24
- // import type { MCPClientConnection } from "./mcp/client-connection";
25
+ import { MCPClientManager, type MCPClientOAuthResult } from "./mcp/client";
26
+ import type { MCPConnectionState } from "./mcp/client-connection";
25
27
  import { DurableObjectOAuthClientProvider } from "./mcp/do-oauth-client-provider";
28
+ import type { TransportType } from "./mcp/types";
29
+ import { genericObservability, type Observability } from "./observability";
30
+ import { DisposableStore } from "./core/events";
31
+ import { MessageType } from "./ai-types";
26
32
 
27
33
  export type { Connection, ConnectionContext, WSMessage } from "partyserver";
28
34
 
@@ -40,7 +46,7 @@ export type RPCRequest = {
40
46
  * State update message from client
41
47
  */
42
48
  export type StateUpdateMessage = {
43
- type: "cf_agent_state";
49
+ type: MessageType.CF_AGENT_STATE;
44
50
  state: unknown;
45
51
  };
46
52
 
@@ -48,7 +54,7 @@ export type StateUpdateMessage = {
48
54
  * RPC response message to client
49
55
  */
50
56
  export type RPCResponse = {
51
- type: "rpc";
57
+ type: MessageType.RPC;
52
58
  id: string;
53
59
  } & (
54
60
  | {
@@ -75,7 +81,7 @@ function isRPCRequest(msg: unknown): msg is RPCRequest {
75
81
  typeof msg === "object" &&
76
82
  msg !== null &&
77
83
  "type" in msg &&
78
- msg.type === "rpc" &&
84
+ msg.type === MessageType.RPC &&
79
85
  "id" in msg &&
80
86
  typeof msg.id === "string" &&
81
87
  "method" in msg &&
@@ -93,7 +99,7 @@ function isStateUpdateMessage(msg: unknown): msg is StateUpdateMessage {
93
99
  typeof msg === "object" &&
94
100
  msg !== null &&
95
101
  "type" in msg &&
96
- msg.type === "cf_agent_state" &&
102
+ msg.type === MessageType.CF_AGENT_STATE &&
97
103
  "state" in msg
98
104
  );
99
105
  }
@@ -114,7 +120,7 @@ const callableMetadata = new Map<Function, CallableMetadata>();
114
120
  * Decorator that marks a method as callable by clients
115
121
  * @param metadata Optional metadata about the callable method
116
122
  */
117
- export function unstable_callable(metadata: CallableMetadata = {}) {
123
+ export function callable(metadata: CallableMetadata = {}) {
118
124
  return function callableDecorator<This, Args extends unknown[], Return>(
119
125
  target: (this: This, ...args: Args) => Return,
120
126
  // biome-ignore lint/correctness/noUnusedFunctionParameters: later
@@ -128,6 +134,30 @@ export function unstable_callable(metadata: CallableMetadata = {}) {
128
134
  };
129
135
  }
130
136
 
137
+ let didWarnAboutUnstableCallable = false;
138
+
139
+ /**
140
+ * Decorator that marks a method as callable by clients
141
+ * @deprecated this has been renamed to callable, and unstable_callable will be removed in the next major version
142
+ * @param metadata Optional metadata about the callable method
143
+ */
144
+ export const unstable_callable = (metadata: CallableMetadata = {}) => {
145
+ if (!didWarnAboutUnstableCallable) {
146
+ didWarnAboutUnstableCallable = true;
147
+ console.warn(
148
+ "unstable_callable is deprecated, use callable instead. unstable_callable will be removed in the next major version."
149
+ );
150
+ }
151
+ callable(metadata);
152
+ };
153
+
154
+ export type QueueItem<T = string> = {
155
+ id: string;
156
+ payload: T;
157
+ callback: keyof Agent<unknown>;
158
+ created_at: number;
159
+ };
160
+
131
161
  /**
132
162
  * Represents a scheduled task within an Agent
133
163
  * @template T Type of the payload data
@@ -169,11 +199,13 @@ function getNextCronTime(cron: string) {
169
199
  return interval.getNextDate();
170
200
  }
171
201
 
202
+ export type { TransportType } from "./mcp/types";
203
+
172
204
  /**
173
205
  * MCP Server state update message from server -> Client
174
206
  */
175
207
  export type MCPServerMessage = {
176
- type: "cf_agent_mcp_servers";
208
+ type: MessageType.CF_AGENT_MCP_SERVERS;
177
209
  mcp: MCPServersState;
178
210
  };
179
211
 
@@ -193,7 +225,7 @@ export type MCPServer = {
193
225
  // This state is specifically about the temporary process of getting a token (if needed).
194
226
  // Scope outside of that can't be relied upon because when the DO sleeps, there's no way
195
227
  // to communicate a change to a non-ready state.
196
- state: "authenticating" | "connecting" | "ready" | "discovering" | "failed";
228
+ state: MCPConnectionState;
197
229
  instructions: string | null;
198
230
  capabilities: ServerCapabilities | null;
199
231
  };
@@ -217,23 +249,26 @@ const STATE_WAS_CHANGED = "cf_state_was_changed";
217
249
  const DEFAULT_STATE = {} as unknown;
218
250
 
219
251
  const agentContext = new AsyncLocalStorage<{
220
- agent: Agent<unknown>;
252
+ agent: Agent<unknown, unknown>;
221
253
  connection: Connection | undefined;
222
254
  request: Request | undefined;
255
+ email: AgentEmail | undefined;
223
256
  }>();
224
257
 
225
258
  export function getCurrentAgent<
226
- T extends Agent<unknown, unknown> = Agent<unknown, unknown>,
259
+ T extends Agent<unknown, unknown> = Agent<unknown, unknown>
227
260
  >(): {
228
261
  agent: T | undefined;
229
262
  connection: Connection | undefined;
230
- request: Request<unknown, CfProperties<unknown>> | undefined;
263
+ request: Request | undefined;
264
+ email: AgentEmail | undefined;
231
265
  } {
232
266
  const store = agentContext.getStore() as
233
267
  | {
234
268
  agent: T;
235
269
  connection: Connection | undefined;
236
- request: Request<unknown, CfProperties<unknown>> | undefined;
270
+ request: Request | undefined;
271
+ email: AgentEmail | undefined;
237
272
  }
238
273
  | undefined;
239
274
  if (!store) {
@@ -241,23 +276,57 @@ export function getCurrentAgent<
241
276
  agent: undefined,
242
277
  connection: undefined,
243
278
  request: undefined,
279
+ email: undefined
244
280
  };
245
281
  }
246
282
  return store;
247
283
  }
248
284
 
285
+ /**
286
+ * Wraps a method to run within the agent context, ensuring getCurrentAgent() works properly
287
+ * @param agent The agent instance
288
+ * @param method The method to wrap
289
+ * @returns A wrapped method that runs within the agent context
290
+ */
291
+
292
+ // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
293
+ function withAgentContext<T extends (...args: any[]) => any>(
294
+ method: T
295
+ ): (this: Agent<unknown, unknown>, ...args: Parameters<T>) => ReturnType<T> {
296
+ return function (...args: Parameters<T>): ReturnType<T> {
297
+ const { connection, request, email, agent } = getCurrentAgent();
298
+
299
+ if (agent === this) {
300
+ // already wrapped, so we can just call the method
301
+ return method.apply(this, args);
302
+ }
303
+ // not wrapped, so we need to wrap it
304
+ return agentContext.run({ agent: this, connection, request, email }, () => {
305
+ return method.apply(this, args);
306
+ });
307
+ };
308
+ }
309
+
249
310
  /**
250
311
  * Base class for creating Agent implementations
251
312
  * @template Env Environment type containing bindings
252
313
  * @template State State type to store within the Agent
253
314
  */
254
- export class Agent<Env, State = unknown> extends Server<Env> {
315
+ export class Agent<
316
+ Env = typeof env,
317
+ State = unknown,
318
+ Props extends Record<string, unknown> = Record<string, unknown>
319
+ > extends Server<Env, Props> {
255
320
  private _state = DEFAULT_STATE as State;
321
+ private _disposables = new DisposableStore();
256
322
 
257
323
  private _ParentClass: typeof Agent<Env, State> =
258
324
  Object.getPrototypeOf(this).constructor;
259
325
 
260
- mcp: MCPClientManager = new MCPClientManager(this._ParentClass.name, "0.0.1");
326
+ readonly mcp: MCPClientManager = new MCPClientManager(
327
+ this._ParentClass.name,
328
+ "0.0.1"
329
+ );
261
330
 
262
331
  /**
263
332
  * Initial state for the Agent
@@ -313,9 +382,14 @@ export class Agent<Env, State = unknown> extends Server<Env> {
313
382
  */
314
383
  static options = {
315
384
  /** Whether the Agent should hibernate when inactive */
316
- hibernate: true, // default to hibernate
385
+ hibernate: true // default to hibernate
317
386
  };
318
387
 
388
+ /**
389
+ * The observability implementation to use for the Agent
390
+ */
391
+ observability?: Observability = genericObservability;
392
+
319
393
  /**
320
394
  * Execute SQL queries against the Agent's database
321
395
  * @template T Type of the returned rows
@@ -345,6 +419,26 @@ export class Agent<Env, State = unknown> extends Server<Env> {
345
419
  constructor(ctx: AgentContext, env: Env) {
346
420
  super(ctx, env);
347
421
 
422
+ if (!wrappedClasses.has(this.constructor)) {
423
+ // Auto-wrap custom methods with agent context
424
+ this._autoWrapCustomMethods();
425
+ wrappedClasses.add(this.constructor);
426
+ }
427
+
428
+ // Broadcast server state after background connects (for OAuth servers)
429
+ this._disposables.add(
430
+ this.mcp.onConnected(async () => {
431
+ this.broadcastMcpServers();
432
+ })
433
+ );
434
+
435
+ // Emit MCP observability events
436
+ this._disposables.add(
437
+ this.mcp.onObservabilityEvent((event) => {
438
+ this.observability?.emit(event);
439
+ })
440
+ );
441
+
348
442
  this.sql`
349
443
  CREATE TABLE IF NOT EXISTS cf_agents_state (
350
444
  id TEXT PRIMARY KEY NOT NULL,
@@ -352,6 +446,15 @@ export class Agent<Env, State = unknown> extends Server<Env> {
352
446
  )
353
447
  `;
354
448
 
449
+ this.sql`
450
+ CREATE TABLE IF NOT EXISTS cf_agents_queues (
451
+ id TEXT PRIMARY KEY NOT NULL,
452
+ payload TEXT,
453
+ callback TEXT,
454
+ created_at INTEGER DEFAULT (unixepoch())
455
+ )
456
+ `;
457
+
355
458
  void this.ctx.blockConcurrencyWhile(async () => {
356
459
  return this._tryCatch(async () => {
357
460
  // Create alarms table if it doesn't exist
@@ -388,24 +491,13 @@ export class Agent<Env, State = unknown> extends Server<Env> {
388
491
  const _onRequest = this.onRequest.bind(this);
389
492
  this.onRequest = (request: Request) => {
390
493
  return agentContext.run(
391
- { agent: this, connection: undefined, request },
494
+ { agent: this, connection: undefined, request, email: undefined },
392
495
  async () => {
393
- if (this.mcp.isCallbackRequest(request)) {
394
- await this.mcp.handleCallbackRequest(request);
395
-
396
- // after the MCP connection handshake, we can send updated mcp state
397
- this.broadcast(
398
- JSON.stringify({
399
- mcp: this.getMcpServers(),
400
- type: "cf_agent_mcp_servers",
401
- })
402
- );
403
-
404
- // We probably should let the user configure this response/redirect, but this is fine for now.
405
- return new Response("<script>window.close();</script>", {
406
- headers: { "content-type": "text/html" },
407
- status: 200,
408
- });
496
+ // Check if this is an OAuth callback and restore state if needed
497
+ const callbackResult =
498
+ await this._handlePotentialOAuthCallback(request);
499
+ if (callbackResult) {
500
+ return callbackResult;
409
501
  }
410
502
 
411
503
  return this._tryCatch(() => _onRequest(request));
@@ -416,7 +508,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
416
508
  const _onMessage = this.onMessage.bind(this);
417
509
  this.onMessage = async (connection: Connection, message: WSMessage) => {
418
510
  return agentContext.run(
419
- { agent: this, connection, request: undefined },
511
+ { agent: this, connection, request: undefined, email: undefined },
420
512
  async () => {
421
513
  if (typeof message !== "string") {
422
514
  return this._tryCatch(() => _onMessage(connection, message));
@@ -460,12 +552,27 @@ export class Agent<Env, State = unknown> extends Server<Env> {
460
552
 
461
553
  // For regular methods, execute and send response
462
554
  const result = await methodFn.apply(this, args);
555
+
556
+ this.observability?.emit(
557
+ {
558
+ displayMessage: `RPC call to ${method}`,
559
+ id: nanoid(),
560
+ payload: {
561
+ method,
562
+ streaming: metadata?.streaming
563
+ },
564
+ timestamp: Date.now(),
565
+ type: "rpc"
566
+ },
567
+ this.ctx
568
+ );
569
+
463
570
  const response: RPCResponse = {
464
571
  done: true,
465
572
  id,
466
573
  result,
467
574
  success: true,
468
- type: "rpc",
575
+ type: MessageType.RPC
469
576
  };
470
577
  connection.send(JSON.stringify(response));
471
578
  } catch (e) {
@@ -475,7 +582,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
475
582
  e instanceof Error ? e.message : "Unknown error occurred",
476
583
  id: parsed.id,
477
584
  success: false,
478
- type: "rpc",
585
+ type: MessageType.RPC
479
586
  };
480
587
  connection.send(JSON.stringify(response));
481
588
  console.error("RPC error:", e);
@@ -493,46 +600,72 @@ export class Agent<Env, State = unknown> extends Server<Env> {
493
600
  // TODO: This is a hack to ensure the state is sent after the connection is established
494
601
  // must fix this
495
602
  return agentContext.run(
496
- { agent: this, connection, request: ctx.request },
497
- async () => {
498
- setTimeout(() => {
499
- if (this.state) {
500
- connection.send(
501
- JSON.stringify({
502
- state: this.state,
503
- type: "cf_agent_state",
504
- })
505
- );
506
- }
507
-
603
+ { agent: this, connection, request: ctx.request, email: undefined },
604
+ () => {
605
+ if (this.state) {
508
606
  connection.send(
509
607
  JSON.stringify({
510
- mcp: this.getMcpServers(),
511
- type: "cf_agent_mcp_servers",
608
+ state: this.state,
609
+ type: MessageType.CF_AGENT_STATE
512
610
  })
513
611
  );
612
+ }
613
+
614
+ connection.send(
615
+ JSON.stringify({
616
+ mcp: this.getMcpServers(),
617
+ type: MessageType.CF_AGENT_MCP_SERVERS
618
+ })
619
+ );
514
620
 
515
- return this._tryCatch(() => _onConnect(connection, ctx));
516
- }, 20);
621
+ this.observability?.emit(
622
+ {
623
+ displayMessage: "Connection established",
624
+ id: nanoid(),
625
+ payload: {
626
+ connectionId: connection.id
627
+ },
628
+ timestamp: Date.now(),
629
+ type: "connect"
630
+ },
631
+ this.ctx
632
+ );
633
+ return this._tryCatch(() => _onConnect(connection, ctx));
517
634
  }
518
635
  );
519
636
  };
520
637
 
521
638
  const _onStart = this.onStart.bind(this);
522
- this.onStart = async () => {
639
+ this.onStart = async (props?: Props) => {
523
640
  return agentContext.run(
524
- { agent: this, connection: undefined, request: undefined },
641
+ {
642
+ agent: this,
643
+ connection: undefined,
644
+ request: undefined,
645
+ email: undefined
646
+ },
525
647
  async () => {
526
- const servers = this.sql<MCPServerRow>`
648
+ await this._tryCatch(() => {
649
+ const servers = this.sql<MCPServerRow>`
527
650
  SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
528
651
  `;
529
652
 
530
- // from DO storage, reconnect to all servers not currently in the oauth flow using our saved auth information
531
- await Promise.allSettled(
532
- servers
533
- .filter((server) => server.auth_url === null)
534
- .map((server) => {
535
- return this._connectToMcpServerInternal(
653
+ this.broadcastMcpServers();
654
+
655
+ // from DO storage, reconnect to all servers not currently in the oauth flow using our saved auth information
656
+ if (servers && Array.isArray(servers) && servers.length > 0) {
657
+ // Restore callback URLs for OAuth-enabled servers
658
+ servers.forEach((server) => {
659
+ if (server.callback_url) {
660
+ // Register the full redirect URL including serverId to avoid ambiguous matches
661
+ this.mcp.registerCallbackUrl(
662
+ `${server.callback_url}/${server.id}`
663
+ );
664
+ }
665
+ });
666
+
667
+ servers.forEach((server) => {
668
+ this._connectToMcpServerInternal(
536
669
  server.name,
537
670
  server.server_url,
538
671
  server.callback_url,
@@ -541,20 +674,25 @@ export class Agent<Env, State = unknown> extends Server<Env> {
541
674
  : undefined,
542
675
  {
543
676
  id: server.id,
544
- oauthClientId: server.client_id ?? undefined,
677
+ oauthClientId: server.client_id ?? undefined
545
678
  }
546
- );
547
- })
548
- );
549
-
550
- this.broadcast(
551
- JSON.stringify({
552
- mcp: this.getMcpServers(),
553
- type: "cf_agent_mcp_servers",
554
- })
555
- );
556
-
557
- await this._tryCatch(() => _onStart());
679
+ )
680
+ .then(() => {
681
+ // Broadcast updated MCP servers state after each server connects
682
+ this.broadcastMcpServers();
683
+ })
684
+ .catch((error) => {
685
+ console.error(
686
+ `Error connecting to MCP server: ${server.name} (${server.server_url})`,
687
+ error
688
+ );
689
+ // Still broadcast even if connection fails, so clients know about the failure
690
+ this.broadcastMcpServers();
691
+ });
692
+ });
693
+ }
694
+ return _onStart(props);
695
+ });
558
696
  }
559
697
  );
560
698
  };
@@ -576,15 +714,25 @@ export class Agent<Env, State = unknown> extends Server<Env> {
576
714
  this.broadcast(
577
715
  JSON.stringify({
578
716
  state: state,
579
- type: "cf_agent_state",
717
+ type: MessageType.CF_AGENT_STATE
580
718
  }),
581
719
  source !== "server" ? [source.id] : []
582
720
  );
583
721
  return this._tryCatch(() => {
584
- const { connection, request } = agentContext.getStore() || {};
722
+ const { connection, request, email } = agentContext.getStore() || {};
585
723
  return agentContext.run(
586
- { agent: this, connection, request },
724
+ { agent: this, connection, request, email },
587
725
  async () => {
726
+ this.observability?.emit(
727
+ {
728
+ displayMessage: "State updated",
729
+ id: nanoid(),
730
+ payload: {},
731
+ timestamp: Date.now(),
732
+ type: "state:update"
733
+ },
734
+ this.ctx
735
+ );
588
736
  return this.onStateUpdate(state, source);
589
737
  }
590
738
  );
@@ -610,19 +758,83 @@ export class Agent<Env, State = unknown> extends Server<Env> {
610
758
  }
611
759
 
612
760
  /**
613
- * Called when the Agent receives an email
761
+ * Called when the Agent receives an email via routeAgentEmail()
762
+ * Override this method to handle incoming emails
614
763
  * @param email Email message to process
615
764
  */
616
- // biome-ignore lint/correctness/noUnusedFunctionParameters: overridden later
617
- onEmail(email: ForwardableEmailMessage) {
765
+ async _onEmail(email: AgentEmail) {
766
+ // nb: we use this roundabout way of getting to onEmail
767
+ // because of https://github.com/cloudflare/workerd/issues/4499
618
768
  return agentContext.run(
619
- { agent: this, connection: undefined, request: undefined },
769
+ { agent: this, connection: undefined, request: undefined, email: email },
620
770
  async () => {
621
- console.error("onEmail not implemented");
771
+ if ("onEmail" in this && typeof this.onEmail === "function") {
772
+ return this._tryCatch(() =>
773
+ (this.onEmail as (email: AgentEmail) => Promise<void>)(email)
774
+ );
775
+ } else {
776
+ console.log("Received email from:", email.from, "to:", email.to);
777
+ console.log("Subject:", email.headers.get("subject"));
778
+ console.log(
779
+ "Implement onEmail(email: AgentEmail): Promise<void> in your agent to process emails"
780
+ );
781
+ }
622
782
  }
623
783
  );
624
784
  }
625
785
 
786
+ /**
787
+ * Reply to an email
788
+ * @param email The email to reply to
789
+ * @param options Options for the reply
790
+ * @returns void
791
+ */
792
+ async replyToEmail(
793
+ email: AgentEmail,
794
+ options: {
795
+ fromName: string;
796
+ subject?: string | undefined;
797
+ body: string;
798
+ contentType?: string;
799
+ headers?: Record<string, string>;
800
+ }
801
+ ): Promise<void> {
802
+ return this._tryCatch(async () => {
803
+ const agentName = camelCaseToKebabCase(this._ParentClass.name);
804
+ const agentId = this.name;
805
+
806
+ const { createMimeMessage } = await import("mimetext");
807
+ const msg = createMimeMessage();
808
+ msg.setSender({ addr: email.to, name: options.fromName });
809
+ msg.setRecipient(email.from);
810
+ msg.setSubject(
811
+ options.subject || `Re: ${email.headers.get("subject")}` || "No subject"
812
+ );
813
+ msg.addMessage({
814
+ contentType: options.contentType || "text/plain",
815
+ data: options.body
816
+ });
817
+
818
+ const domain = email.from.split("@")[1];
819
+ const messageId = `<${agentId}@${domain}>`;
820
+ msg.setHeader("In-Reply-To", email.headers.get("Message-ID")!);
821
+ msg.setHeader("Message-ID", messageId);
822
+ msg.setHeader("X-Agent-Name", agentName);
823
+ msg.setHeader("X-Agent-ID", agentId);
824
+
825
+ if (options.headers) {
826
+ for (const [key, value] of Object.entries(options.headers)) {
827
+ msg.setHeader(key, value);
828
+ }
829
+ }
830
+ await email.reply({
831
+ from: email.to,
832
+ raw: msg.asRaw(),
833
+ to: email.from
834
+ });
835
+ });
836
+ }
837
+
626
838
  private async _tryCatch<T>(fn: () => T | Promise<T>) {
627
839
  try {
628
840
  return await fn();
@@ -631,6 +843,68 @@ export class Agent<Env, State = unknown> extends Server<Env> {
631
843
  }
632
844
  }
633
845
 
846
+ /**
847
+ * Automatically wrap custom methods with agent context
848
+ * This ensures getCurrentAgent() works in all custom methods without decorators
849
+ */
850
+ private _autoWrapCustomMethods() {
851
+ // Collect all methods from base prototypes (Agent and Server)
852
+ const basePrototypes = [Agent.prototype, Server.prototype];
853
+ const baseMethods = new Set<string>();
854
+ for (const baseProto of basePrototypes) {
855
+ let proto = baseProto;
856
+ while (proto && proto !== Object.prototype) {
857
+ const methodNames = Object.getOwnPropertyNames(proto);
858
+ for (const methodName of methodNames) {
859
+ baseMethods.add(methodName);
860
+ }
861
+ proto = Object.getPrototypeOf(proto);
862
+ }
863
+ }
864
+ // Get all methods from the current instance's prototype chain
865
+ let proto = Object.getPrototypeOf(this);
866
+ let depth = 0;
867
+ while (proto && proto !== Object.prototype && depth < 10) {
868
+ const methodNames = Object.getOwnPropertyNames(proto);
869
+ for (const methodName of methodNames) {
870
+ const descriptor = Object.getOwnPropertyDescriptor(proto, methodName);
871
+
872
+ // Skip if it's a private method, a base method, a getter, or not a function,
873
+ if (
874
+ baseMethods.has(methodName) ||
875
+ methodName.startsWith("_") ||
876
+ !descriptor ||
877
+ !!descriptor.get ||
878
+ typeof descriptor.value !== "function"
879
+ ) {
880
+ continue;
881
+ }
882
+
883
+ // Now, methodName is confirmed to be a custom method/function
884
+ // Wrap the custom method with context
885
+ const wrappedFunction = withAgentContext(
886
+ // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
887
+ this[methodName as keyof this] as (...args: any[]) => any
888
+ // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
889
+ ) as any;
890
+
891
+ // if the method is callable, copy the metadata from the original method
892
+ if (this._isCallable(methodName)) {
893
+ callableMetadata.set(
894
+ wrappedFunction,
895
+ callableMetadata.get(this[methodName as keyof this] as Function)!
896
+ );
897
+ }
898
+
899
+ // set the wrapped function on the prototype
900
+ this.constructor.prototype[methodName as keyof this] = wrappedFunction;
901
+ }
902
+
903
+ proto = Object.getPrototypeOf(proto);
904
+ depth++;
905
+ }
906
+ }
907
+
634
908
  override onError(
635
909
  connection: Connection,
636
910
  error: unknown
@@ -665,6 +939,131 @@ export class Agent<Env, State = unknown> extends Server<Env> {
665
939
  throw new Error("Not implemented");
666
940
  }
667
941
 
942
+ /**
943
+ * Queue a task to be executed in the future
944
+ * @param payload Payload to pass to the callback
945
+ * @param callback Name of the method to call
946
+ * @returns The ID of the queued task
947
+ */
948
+ async queue<T = unknown>(callback: keyof this, payload: T): Promise<string> {
949
+ const id = nanoid(9);
950
+ if (typeof callback !== "string") {
951
+ throw new Error("Callback must be a string");
952
+ }
953
+
954
+ if (typeof this[callback] !== "function") {
955
+ throw new Error(`this.${callback} is not a function`);
956
+ }
957
+
958
+ this.sql`
959
+ INSERT OR REPLACE INTO cf_agents_queues (id, payload, callback)
960
+ VALUES (${id}, ${JSON.stringify(payload)}, ${callback})
961
+ `;
962
+
963
+ void this._flushQueue().catch((e) => {
964
+ console.error("Error flushing queue:", e);
965
+ });
966
+
967
+ return id;
968
+ }
969
+
970
+ private _flushingQueue = false;
971
+
972
+ private async _flushQueue() {
973
+ if (this._flushingQueue) {
974
+ return;
975
+ }
976
+ this._flushingQueue = true;
977
+ while (true) {
978
+ const result = this.sql<QueueItem<string>>`
979
+ SELECT * FROM cf_agents_queues
980
+ ORDER BY created_at ASC
981
+ `;
982
+
983
+ if (!result || result.length === 0) {
984
+ break;
985
+ }
986
+
987
+ for (const row of result || []) {
988
+ const callback = this[row.callback as keyof Agent<Env>];
989
+ if (!callback) {
990
+ console.error(`callback ${row.callback} not found`);
991
+ continue;
992
+ }
993
+ const { connection, request, email } = agentContext.getStore() || {};
994
+ await agentContext.run(
995
+ {
996
+ agent: this,
997
+ connection,
998
+ request,
999
+ email
1000
+ },
1001
+ async () => {
1002
+ // TODO: add retries and backoff
1003
+ await (
1004
+ callback as (
1005
+ payload: unknown,
1006
+ queueItem: QueueItem<string>
1007
+ ) => Promise<void>
1008
+ ).bind(this)(JSON.parse(row.payload as string), row);
1009
+ await this.dequeue(row.id);
1010
+ }
1011
+ );
1012
+ }
1013
+ }
1014
+ this._flushingQueue = false;
1015
+ }
1016
+
1017
+ /**
1018
+ * Dequeue a task by ID
1019
+ * @param id ID of the task to dequeue
1020
+ */
1021
+ async dequeue(id: string) {
1022
+ this.sql`DELETE FROM cf_agents_queues WHERE id = ${id}`;
1023
+ }
1024
+
1025
+ /**
1026
+ * Dequeue all tasks
1027
+ */
1028
+ async dequeueAll() {
1029
+ this.sql`DELETE FROM cf_agents_queues`;
1030
+ }
1031
+
1032
+ /**
1033
+ * Dequeue all tasks by callback
1034
+ * @param callback Name of the callback to dequeue
1035
+ */
1036
+ async dequeueAllByCallback(callback: string) {
1037
+ this.sql`DELETE FROM cf_agents_queues WHERE callback = ${callback}`;
1038
+ }
1039
+
1040
+ /**
1041
+ * Get a queued task by ID
1042
+ * @param id ID of the task to get
1043
+ * @returns The task or undefined if not found
1044
+ */
1045
+ async getQueue(id: string): Promise<QueueItem<string> | undefined> {
1046
+ const result = this.sql<QueueItem<string>>`
1047
+ SELECT * FROM cf_agents_queues WHERE id = ${id}
1048
+ `;
1049
+ return result
1050
+ ? { ...result[0], payload: JSON.parse(result[0].payload) }
1051
+ : undefined;
1052
+ }
1053
+
1054
+ /**
1055
+ * Get all queues by key and value
1056
+ * @param key Key to filter by
1057
+ * @param value Value to filter by
1058
+ * @returns Array of matching QueueItem objects
1059
+ */
1060
+ async getQueues(key: string, value: string): Promise<QueueItem<string>[]> {
1061
+ const result = this.sql<QueueItem<string>>`
1062
+ SELECT * FROM cf_agents_queues
1063
+ `;
1064
+ return result.filter((row) => JSON.parse(row.payload)[key] === value);
1065
+ }
1066
+
668
1067
  /**
669
1068
  * Schedule a task to be executed in the future
670
1069
  * @template T Type of the payload data
@@ -680,6 +1079,21 @@ export class Agent<Env, State = unknown> extends Server<Env> {
680
1079
  ): Promise<Schedule<T>> {
681
1080
  const id = nanoid(9);
682
1081
 
1082
+ const emitScheduleCreate = (schedule: Schedule<T>) =>
1083
+ this.observability?.emit(
1084
+ {
1085
+ displayMessage: `Schedule ${schedule.id} created`,
1086
+ id: nanoid(),
1087
+ payload: {
1088
+ callback: callback as string,
1089
+ id: id
1090
+ },
1091
+ timestamp: Date.now(),
1092
+ type: "schedule:create"
1093
+ },
1094
+ this.ctx
1095
+ );
1096
+
683
1097
  if (typeof callback !== "string") {
684
1098
  throw new Error("Callback must be a string");
685
1099
  }
@@ -699,13 +1113,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
699
1113
 
700
1114
  await this._scheduleNextAlarm();
701
1115
 
702
- return {
1116
+ const schedule: Schedule<T> = {
703
1117
  callback: callback,
704
1118
  id,
705
1119
  payload: payload as T,
706
1120
  time: timestamp,
707
- type: "scheduled",
1121
+ type: "scheduled"
708
1122
  };
1123
+
1124
+ emitScheduleCreate(schedule);
1125
+
1126
+ return schedule;
709
1127
  }
710
1128
  if (typeof when === "number") {
711
1129
  const time = new Date(Date.now() + when * 1000);
@@ -720,14 +1138,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
720
1138
 
721
1139
  await this._scheduleNextAlarm();
722
1140
 
723
- return {
1141
+ const schedule: Schedule<T> = {
724
1142
  callback: callback,
725
1143
  delayInSeconds: when,
726
1144
  id,
727
1145
  payload: payload as T,
728
1146
  time: timestamp,
729
- type: "delayed",
1147
+ type: "delayed"
730
1148
  };
1149
+
1150
+ emitScheduleCreate(schedule);
1151
+
1152
+ return schedule;
731
1153
  }
732
1154
  if (typeof when === "string") {
733
1155
  const nextExecutionTime = getNextCronTime(when);
@@ -742,14 +1164,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
742
1164
 
743
1165
  await this._scheduleNextAlarm();
744
1166
 
745
- return {
1167
+ const schedule: Schedule<T> = {
746
1168
  callback: callback,
747
1169
  cron: when,
748
1170
  id,
749
1171
  payload: payload as T,
750
1172
  time: timestamp,
751
- type: "cron",
1173
+ type: "cron"
752
1174
  };
1175
+
1176
+ emitScheduleCreate(schedule);
1177
+
1178
+ return schedule;
753
1179
  }
754
1180
  throw new Error("Invalid schedule type");
755
1181
  }
@@ -813,7 +1239,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
813
1239
  .toArray()
814
1240
  .map((row) => ({
815
1241
  ...row,
816
- payload: JSON.parse(row.payload as string) as T,
1242
+ payload: JSON.parse(row.payload as string) as T
817
1243
  })) as Schedule<T>[];
818
1244
 
819
1245
  return result;
@@ -825,6 +1251,22 @@ export class Agent<Env, State = unknown> extends Server<Env> {
825
1251
  * @returns true if the task was cancelled, false otherwise
826
1252
  */
827
1253
  async cancelSchedule(id: string): Promise<boolean> {
1254
+ const schedule = await this.getSchedule(id);
1255
+ if (schedule) {
1256
+ this.observability?.emit(
1257
+ {
1258
+ displayMessage: `Schedule ${id} cancelled`,
1259
+ id: nanoid(),
1260
+ payload: {
1261
+ callback: schedule.callback,
1262
+ id: schedule.id
1263
+ },
1264
+ timestamp: Date.now(),
1265
+ type: "schedule:cancel"
1266
+ },
1267
+ this.ctx
1268
+ );
1269
+ }
828
1270
  this.sql`DELETE FROM cf_agents_schedules WHERE id = ${id}`;
829
1271
 
830
1272
  await this._scheduleNextAlarm();
@@ -834,9 +1276,9 @@ export class Agent<Env, State = unknown> extends Server<Env> {
834
1276
  private async _scheduleNextAlarm() {
835
1277
  // Find the next schedule that needs to be executed
836
1278
  const result = this.sql`
837
- SELECT time FROM cf_agents_schedules
1279
+ SELECT time FROM cf_agents_schedules
838
1280
  WHERE time > ${Math.floor(Date.now() / 1000)}
839
- ORDER BY time ASC
1281
+ ORDER BY time ASC
840
1282
  LIMIT 1
841
1283
  `;
842
1284
  if (!result) return;
@@ -863,40 +1305,61 @@ export class Agent<Env, State = unknown> extends Server<Env> {
863
1305
  SELECT * FROM cf_agents_schedules WHERE time <= ${now}
864
1306
  `;
865
1307
 
866
- for (const row of result || []) {
867
- const callback = this[row.callback as keyof Agent<Env>];
868
- if (!callback) {
869
- console.error(`callback ${row.callback} not found`);
870
- continue;
871
- }
872
- await agentContext.run(
873
- { agent: this, connection: undefined, request: undefined },
874
- async () => {
875
- try {
876
- await (
877
- callback as (
878
- payload: unknown,
879
- schedule: Schedule<unknown>
880
- ) => Promise<void>
881
- ).bind(this)(JSON.parse(row.payload as string), row);
882
- } catch (e) {
883
- console.error(`error executing callback "${row.callback}"`, e);
884
- }
1308
+ if (result && Array.isArray(result)) {
1309
+ for (const row of result) {
1310
+ const callback = this[row.callback as keyof Agent<Env>];
1311
+ if (!callback) {
1312
+ console.error(`callback ${row.callback} not found`);
1313
+ continue;
885
1314
  }
886
- );
887
- if (row.type === "cron") {
888
- // Update next execution time for cron schedules
889
- const nextExecutionTime = getNextCronTime(row.cron);
890
- const nextTimestamp = Math.floor(nextExecutionTime.getTime() / 1000);
1315
+ await agentContext.run(
1316
+ {
1317
+ agent: this,
1318
+ connection: undefined,
1319
+ request: undefined,
1320
+ email: undefined
1321
+ },
1322
+ async () => {
1323
+ try {
1324
+ this.observability?.emit(
1325
+ {
1326
+ displayMessage: `Schedule ${row.id} executed`,
1327
+ id: nanoid(),
1328
+ payload: {
1329
+ callback: row.callback,
1330
+ id: row.id
1331
+ },
1332
+ timestamp: Date.now(),
1333
+ type: "schedule:execute"
1334
+ },
1335
+ this.ctx
1336
+ );
891
1337
 
892
- this.sql`
1338
+ await (
1339
+ callback as (
1340
+ payload: unknown,
1341
+ schedule: Schedule<unknown>
1342
+ ) => Promise<void>
1343
+ ).bind(this)(JSON.parse(row.payload as string), row);
1344
+ } catch (e) {
1345
+ console.error(`error executing callback "${row.callback}"`, e);
1346
+ }
1347
+ }
1348
+ );
1349
+ if (row.type === "cron") {
1350
+ // Update next execution time for cron schedules
1351
+ const nextExecutionTime = getNextCronTime(row.cron);
1352
+ const nextTimestamp = Math.floor(nextExecutionTime.getTime() / 1000);
1353
+
1354
+ this.sql`
893
1355
  UPDATE cf_agents_schedules SET time = ${nextTimestamp} WHERE id = ${row.id}
894
1356
  `;
895
- } else {
896
- // Delete one-time schedules after execution
897
- this.sql`
1357
+ } else {
1358
+ // Delete one-time schedules after execution
1359
+ this.sql`
898
1360
  DELETE FROM cf_agents_schedules WHERE id = ${row.id}
899
1361
  `;
1362
+ }
900
1363
  }
901
1364
  }
902
1365
 
@@ -912,11 +1375,25 @@ export class Agent<Env, State = unknown> extends Server<Env> {
912
1375
  this.sql`DROP TABLE IF EXISTS cf_agents_state`;
913
1376
  this.sql`DROP TABLE IF EXISTS cf_agents_schedules`;
914
1377
  this.sql`DROP TABLE IF EXISTS cf_agents_mcp_servers`;
1378
+ this.sql`DROP TABLE IF EXISTS cf_agents_queues`;
915
1379
 
916
1380
  // delete all alarms
917
1381
  await this.ctx.storage.deleteAlarm();
918
1382
  await this.ctx.storage.deleteAll();
1383
+ this._disposables.dispose();
1384
+ await this.mcp.dispose?.();
919
1385
  this.ctx.abort("destroyed"); // enforce that the agent is evicted
1386
+
1387
+ this.observability?.emit(
1388
+ {
1389
+ displayMessage: "Agent destroyed",
1390
+ id: nanoid(),
1391
+ payload: {},
1392
+ timestamp: Date.now(),
1393
+ type: "destroy"
1394
+ },
1395
+ this.ctx
1396
+ );
920
1397
  }
921
1398
 
922
1399
  /**
@@ -930,45 +1407,221 @@ export class Agent<Env, State = unknown> extends Server<Env> {
930
1407
  /**
931
1408
  * Connect to a new MCP Server
932
1409
  *
1410
+ * @param serverName Name of the MCP server
933
1411
  * @param url MCP Server SSE URL
934
- * @param callbackHost Base host for the agent, used for the redirect URI.
1412
+ * @param callbackHost Base host for the agent, used for the redirect URI. If not provided, will be derived from the current request.
935
1413
  * @param agentsPrefix agents routing prefix if not using `agents`
936
- * @param options MCP client and transport (header) options
1414
+ * @param options MCP client and transport options
937
1415
  * @returns authUrl
938
1416
  */
939
1417
  async addMcpServer(
940
1418
  serverName: string,
941
1419
  url: string,
942
- callbackHost: string,
1420
+ callbackHost?: string,
943
1421
  agentsPrefix = "agents",
944
1422
  options?: {
945
1423
  client?: ConstructorParameters<typeof Client>[1];
946
1424
  transport?: {
947
- headers: HeadersInit;
1425
+ headers?: HeadersInit;
1426
+ type?: TransportType;
948
1427
  };
949
1428
  }
950
1429
  ): Promise<{ id: string; authUrl: string | undefined }> {
951
- const callbackUrl = `${callbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
1430
+ // If callbackHost is not provided, derive it from the current request
1431
+ let resolvedCallbackHost = callbackHost;
1432
+ if (!resolvedCallbackHost) {
1433
+ const { request } = getCurrentAgent();
1434
+ if (!request) {
1435
+ throw new Error(
1436
+ "callbackHost is required when not called within a request context"
1437
+ );
1438
+ }
1439
+
1440
+ // Extract the origin from the request
1441
+ const requestUrl = new URL(request.url);
1442
+ resolvedCallbackHost = `${requestUrl.protocol}//${requestUrl.host}`;
1443
+ }
1444
+
1445
+ const callbackUrl = `${resolvedCallbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
1446
+
1447
+ // Generate a serverId upfront
1448
+ const serverId = nanoid(8);
1449
+
1450
+ // Persist to database BEFORE starting OAuth flow to survive DO hibernation
1451
+ this.sql`
1452
+ INSERT OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
1453
+ VALUES (
1454
+ ${serverId},
1455
+ ${serverName},
1456
+ ${url},
1457
+ ${null},
1458
+ ${null},
1459
+ ${callbackUrl},
1460
+ ${options ? JSON.stringify(options) : null}
1461
+ );
1462
+ `;
952
1463
 
1464
+ // _connectToMcpServerInternal will call mcp.connect which registers the callback URL
953
1465
  const result = await this._connectToMcpServerInternal(
954
1466
  serverName,
955
1467
  url,
956
1468
  callbackUrl,
957
- options
1469
+ options,
1470
+ { id: serverId }
958
1471
  );
959
1472
 
960
- this.broadcast(
961
- JSON.stringify({
962
- mcp: this.getMcpServers(),
963
- type: "cf_agent_mcp_servers",
964
- })
965
- );
1473
+ // Update database with OAuth client info if auth is required
1474
+ if (result.clientId || result.authUrl) {
1475
+ this.sql`
1476
+ UPDATE cf_agents_mcp_servers
1477
+ SET client_id = ${result.clientId ?? null}, auth_url = ${result.authUrl ?? null}
1478
+ WHERE id = ${serverId}
1479
+ `;
1480
+ }
1481
+
1482
+ this.broadcastMcpServers();
966
1483
 
967
1484
  return result;
968
1485
  }
969
1486
 
970
- async _connectToMcpServerInternal(
971
- serverName: string,
1487
+ /**
1488
+ * Handle potential OAuth callback requests after DO hibernation.
1489
+ * Detects OAuth callbacks, restores state from database, and processes the callback.
1490
+ * Returns a Response if this was an OAuth callback, otherwise returns undefined.
1491
+ */
1492
+ private async _handlePotentialOAuthCallback(
1493
+ request: Request
1494
+ ): Promise<Response | undefined> {
1495
+ // Quick check: must be GET with callback pattern and code parameter
1496
+ if (request.method !== "GET") {
1497
+ return undefined;
1498
+ }
1499
+
1500
+ const url = new URL(request.url);
1501
+ const hasCallbackPattern =
1502
+ url.pathname.includes("/callback/") && url.searchParams.has("code");
1503
+
1504
+ if (!hasCallbackPattern) {
1505
+ return undefined;
1506
+ }
1507
+
1508
+ // Extract serverId from callback URL
1509
+ const pathParts = url.pathname.split("/");
1510
+ const callbackIndex = pathParts.indexOf("callback");
1511
+ const serverId = callbackIndex !== -1 ? pathParts[callbackIndex + 1] : null;
1512
+
1513
+ if (!serverId) {
1514
+ return new Response("Invalid callback URL: missing serverId", {
1515
+ status: 400
1516
+ });
1517
+ }
1518
+
1519
+ // Check if callback is already registered AND connection exists (not hibernated)
1520
+ if (
1521
+ this.mcp.isCallbackRequest(request) &&
1522
+ this.mcp.mcpConnections[serverId]
1523
+ ) {
1524
+ // State already restored, handle normally
1525
+ return this._processOAuthCallback(request);
1526
+ }
1527
+
1528
+ // Need to restore from database after hibernation
1529
+ try {
1530
+ const server = this.sql<MCPServerRow>`
1531
+ SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
1532
+ FROM cf_agents_mcp_servers
1533
+ WHERE id = ${serverId}
1534
+ `.find((s) => s.id === serverId);
1535
+
1536
+ if (!server) {
1537
+ return new Response(
1538
+ `OAuth callback failed: Server ${serverId} not found in database`,
1539
+ { status: 404 }
1540
+ );
1541
+ }
1542
+
1543
+ // Register callback URL (restores it after hibernation)
1544
+ if (!server.callback_url) {
1545
+ return new Response(
1546
+ `OAuth callback failed: No callback URL stored for server ${serverId}`,
1547
+ { status: 500 }
1548
+ );
1549
+ }
1550
+
1551
+ this.mcp.registerCallbackUrl(`${server.callback_url}/${server.id}`);
1552
+
1553
+ // Restore connection if not in memory
1554
+ if (!this.mcp.mcpConnections[serverId]) {
1555
+ let parsedOptions:
1556
+ | {
1557
+ client?: ConstructorParameters<typeof Client>[1];
1558
+ transport?: {
1559
+ headers?: HeadersInit;
1560
+ type?: TransportType;
1561
+ };
1562
+ }
1563
+ | undefined;
1564
+ try {
1565
+ parsedOptions = server.server_options
1566
+ ? JSON.parse(server.server_options)
1567
+ : undefined;
1568
+ } catch {
1569
+ return new Response(
1570
+ `OAuth callback failed: Invalid server options in database for ${serverId}`,
1571
+ { status: 500 }
1572
+ );
1573
+ }
1574
+
1575
+ await this._connectToMcpServerInternal(
1576
+ server.name,
1577
+ server.server_url,
1578
+ server.callback_url,
1579
+ parsedOptions,
1580
+ {
1581
+ id: server.id,
1582
+ oauthClientId: server.client_id ?? undefined
1583
+ }
1584
+ );
1585
+ }
1586
+
1587
+ // Now process the OAuth callback
1588
+ return this._processOAuthCallback(request);
1589
+ } catch (error) {
1590
+ const errorMsg = error instanceof Error ? error.message : "Unknown error";
1591
+ console.error(`Failed to restore MCP state for ${serverId}:`, error);
1592
+ return new Response(
1593
+ `OAuth callback failed during state restoration: ${errorMsg}`,
1594
+ { status: 500 }
1595
+ );
1596
+ }
1597
+ }
1598
+
1599
+ /**
1600
+ * Process an OAuth callback request (assumes state is already restored)
1601
+ */
1602
+ private async _processOAuthCallback(request: Request): Promise<Response> {
1603
+ const result = await this.mcp.handleCallbackRequest(request);
1604
+ this.broadcastMcpServers();
1605
+
1606
+ if (result.authSuccess) {
1607
+ // Start background connection if auth was successful
1608
+ this.mcp
1609
+ .establishConnection(result.serverId)
1610
+ .catch((error) => {
1611
+ console.error("Background connection failed:", error);
1612
+ })
1613
+ .finally(() => {
1614
+ // Broadcast after background connection resolves (success/failure)
1615
+ this.broadcastMcpServers();
1616
+ });
1617
+ }
1618
+
1619
+ // Handle OAuth callback response using MCPClientManager configuration
1620
+ return this.handleOAuthCallbackResponse(result, request);
1621
+ }
1622
+
1623
+ private async _connectToMcpServerInternal(
1624
+ _serverName: string,
972
1625
  url: string,
973
1626
  callbackUrl: string,
974
1627
  // it's important that any options here are serializable because we put them into our sqlite DB for reconnection purposes
@@ -983,13 +1636,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
983
1636
  */
984
1637
  transport?: {
985
1638
  headers?: HeadersInit;
1639
+ type?: TransportType;
986
1640
  };
987
1641
  },
988
1642
  reconnect?: {
989
1643
  id: string;
990
1644
  oauthClientId?: string;
991
1645
  }
992
- ): Promise<{ id: string; authUrl: string | undefined }> {
1646
+ ): Promise<{
1647
+ id: string;
1648
+ authUrl: string | undefined;
1649
+ clientId: string | undefined;
1650
+ }> {
993
1651
  const authProvider = new DurableObjectOAuthClientProvider(
994
1652
  this.ctx.storage,
995
1653
  this.name,
@@ -1003,6 +1661,9 @@ export class Agent<Env, State = unknown> extends Server<Env> {
1003
1661
  }
1004
1662
  }
1005
1663
 
1664
+ // Use the transport type specified in options, or default to "auto"
1665
+ const transportType: TransportType = options?.transport?.type ?? "auto";
1666
+
1006
1667
  // allows passing through transport headers if necessary
1007
1668
  // this handles some non-standard bearer auth setups (i.e. MCP server behind CF access instead of OAuth)
1008
1669
  let headerTransportOpts: SSEClientTransportOptions = {};
@@ -1012,12 +1673,12 @@ export class Agent<Env, State = unknown> extends Server<Env> {
1012
1673
  fetch: (url, init) =>
1013
1674
  fetch(url, {
1014
1675
  ...init,
1015
- headers: options?.transport?.headers,
1016
- }),
1676
+ headers: options?.transport?.headers
1677
+ })
1017
1678
  },
1018
1679
  requestInit: {
1019
- headers: options?.transport?.headers,
1020
- },
1680
+ headers: options?.transport?.headers
1681
+ }
1021
1682
  };
1022
1683
  }
1023
1684
 
@@ -1027,39 +1688,24 @@ export class Agent<Env, State = unknown> extends Server<Env> {
1027
1688
  transport: {
1028
1689
  ...headerTransportOpts,
1029
1690
  authProvider,
1030
- },
1691
+ type: transportType
1692
+ }
1031
1693
  });
1032
1694
 
1033
- this.sql`
1034
- INSERT OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
1035
- VALUES (
1036
- ${id},
1037
- ${serverName},
1038
- ${url},
1039
- ${clientId ?? null},
1040
- ${authUrl ?? null},
1041
- ${callbackUrl},
1042
- ${options ? JSON.stringify(options) : null}
1043
- );
1044
- `;
1045
-
1046
1695
  return {
1047
1696
  authUrl,
1048
- id,
1697
+ clientId,
1698
+ id
1049
1699
  };
1050
1700
  }
1051
1701
 
1052
1702
  async removeMcpServer(id: string) {
1053
1703
  this.mcp.closeConnection(id);
1704
+ this.mcp.unregisterCallbackUrl(id);
1054
1705
  this.sql`
1055
1706
  DELETE FROM cf_agents_mcp_servers WHERE id = ${id};
1056
1707
  `;
1057
- this.broadcast(
1058
- JSON.stringify({
1059
- mcp: this.getMcpServers(),
1060
- type: "cf_agent_mcp_servers",
1061
- })
1062
- );
1708
+ this.broadcastMcpServers();
1063
1709
  }
1064
1710
 
1065
1711
  getMcpServers(): MCPServersState {
@@ -1067,30 +1713,77 @@ export class Agent<Env, State = unknown> extends Server<Env> {
1067
1713
  prompts: this.mcp.listPrompts(),
1068
1714
  resources: this.mcp.listResources(),
1069
1715
  servers: {},
1070
- tools: this.mcp.listTools(),
1716
+ tools: this.mcp.listTools()
1071
1717
  };
1072
1718
 
1073
1719
  const servers = this.sql<MCPServerRow>`
1074
1720
  SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
1075
1721
  `;
1076
1722
 
1077
- for (const server of servers) {
1078
- const serverConn = this.mcp.mcpConnections[server.id];
1079
- mcpState.servers[server.id] = {
1080
- auth_url: server.auth_url,
1081
- capabilities: serverConn?.serverCapabilities ?? null,
1082
- instructions: serverConn?.instructions ?? null,
1083
- name: server.name,
1084
- server_url: server.server_url,
1085
- // mark as "authenticating" because the server isn't automatically connected, so it's pending authenticating
1086
- state: serverConn?.connectionState ?? "authenticating",
1087
- };
1723
+ if (servers && Array.isArray(servers) && servers.length > 0) {
1724
+ for (const server of servers) {
1725
+ const serverConn = this.mcp.mcpConnections[server.id];
1726
+ mcpState.servers[server.id] = {
1727
+ auth_url: server.auth_url,
1728
+ capabilities: serverConn?.serverCapabilities ?? null,
1729
+ instructions: serverConn?.instructions ?? null,
1730
+ name: server.name,
1731
+ server_url: server.server_url,
1732
+ // mark as "authenticating" because the server isn't automatically connected, so it's pending authenticating
1733
+ state: serverConn?.connectionState ?? "authenticating"
1734
+ };
1735
+ }
1088
1736
  }
1089
1737
 
1090
1738
  return mcpState;
1091
1739
  }
1740
+
1741
+ private broadcastMcpServers() {
1742
+ this.broadcast(
1743
+ JSON.stringify({
1744
+ mcp: this.getMcpServers(),
1745
+ type: MessageType.CF_AGENT_MCP_SERVERS
1746
+ })
1747
+ );
1748
+ }
1749
+
1750
+ /**
1751
+ * Handle OAuth callback response using MCPClientManager configuration
1752
+ * @param result OAuth callback result
1753
+ * @param request The original request (needed for base URL)
1754
+ * @returns Response for the OAuth callback
1755
+ */
1756
+ private handleOAuthCallbackResponse(
1757
+ result: MCPClientOAuthResult,
1758
+ request: Request
1759
+ ): Response {
1760
+ const config = this.mcp.getOAuthCallbackConfig();
1761
+
1762
+ // Use custom handler if configured
1763
+ if (config?.customHandler) {
1764
+ return config.customHandler(result);
1765
+ }
1766
+
1767
+ // Use redirect URLs if configured
1768
+ if (config?.successRedirect && result.authSuccess) {
1769
+ return Response.redirect(config.successRedirect);
1770
+ }
1771
+
1772
+ if (config?.errorRedirect && !result.authSuccess) {
1773
+ return Response.redirect(
1774
+ `${config.errorRedirect}?error=${encodeURIComponent(result.authError || "Unknown error")}`
1775
+ );
1776
+ }
1777
+
1778
+ // Default behavior - redirect to base URL
1779
+ const baseUrl = new URL(request.url).origin;
1780
+ return Response.redirect(baseUrl);
1781
+ }
1092
1782
  }
1093
1783
 
1784
+ // A set of classes that have been wrapped with agent context
1785
+ const wrappedClasses = new Set<typeof Agent.prototype.constructor>();
1786
+
1094
1787
  /**
1095
1788
  * Namespace for creating Agent instances
1096
1789
  * @template Agentic Type of the Agent class
@@ -1131,14 +1824,14 @@ export async function routeAgentRequest<Env>(
1131
1824
  "Access-Control-Allow-Credentials": "true",
1132
1825
  "Access-Control-Allow-Methods": "GET, POST, HEAD, OPTIONS",
1133
1826
  "Access-Control-Allow-Origin": "*",
1134
- "Access-Control-Max-Age": "86400",
1827
+ "Access-Control-Max-Age": "86400"
1135
1828
  }
1136
1829
  : options?.cors;
1137
1830
 
1138
1831
  if (request.method === "OPTIONS") {
1139
1832
  if (corsHeaders) {
1140
1833
  return new Response(null, {
1141
- headers: corsHeaders,
1834
+ headers: corsHeaders
1142
1835
  });
1143
1836
  }
1144
1837
  console.warn(
@@ -1151,7 +1844,7 @@ export async function routeAgentRequest<Env>(
1151
1844
  env as Record<string, unknown>,
1152
1845
  {
1153
1846
  prefix: "agents",
1154
- ...(options as PartyServerOptions<Record<string, unknown>>),
1847
+ ...(options as PartyServerOptions<Record<string, unknown>>)
1155
1848
  }
1156
1849
  );
1157
1850
 
@@ -1164,24 +1857,238 @@ export async function routeAgentRequest<Env>(
1164
1857
  response = new Response(response.body, {
1165
1858
  headers: {
1166
1859
  ...response.headers,
1167
- ...corsHeaders,
1168
- },
1860
+ ...corsHeaders
1861
+ }
1169
1862
  });
1170
1863
  }
1171
1864
  return response;
1172
1865
  }
1173
1866
 
1867
+ export type EmailResolver<Env> = (
1868
+ email: ForwardableEmailMessage,
1869
+ env: Env
1870
+ ) => Promise<{
1871
+ agentName: string;
1872
+ agentId: string;
1873
+ } | null>;
1874
+
1875
+ /**
1876
+ * Create a resolver that uses the message-id header to determine the agent to route the email to
1877
+ * @returns A function that resolves the agent to route the email to
1878
+ */
1879
+ export function createHeaderBasedEmailResolver<Env>(): EmailResolver<Env> {
1880
+ return async (email: ForwardableEmailMessage, _env: Env) => {
1881
+ const messageId = email.headers.get("message-id");
1882
+ if (messageId) {
1883
+ const messageIdMatch = messageId.match(/<([^@]+)@([^>]+)>/);
1884
+ if (messageIdMatch) {
1885
+ const [, agentId, domain] = messageIdMatch;
1886
+ const agentName = domain.split(".")[0];
1887
+ return { agentName, agentId };
1888
+ }
1889
+ }
1890
+
1891
+ const references = email.headers.get("references");
1892
+ if (references) {
1893
+ const referencesMatch = references.match(
1894
+ /<([A-Za-z0-9+/]{43}=)@([^>]+)>/
1895
+ );
1896
+ if (referencesMatch) {
1897
+ const [, base64Id, domain] = referencesMatch;
1898
+ const agentId = Buffer.from(base64Id, "base64").toString("hex");
1899
+ const agentName = domain.split(".")[0];
1900
+ return { agentName, agentId };
1901
+ }
1902
+ }
1903
+
1904
+ const agentName = email.headers.get("x-agent-name");
1905
+ const agentId = email.headers.get("x-agent-id");
1906
+ if (agentName && agentId) {
1907
+ return { agentName, agentId };
1908
+ }
1909
+
1910
+ return null;
1911
+ };
1912
+ }
1913
+
1914
+ /**
1915
+ * Create a resolver that uses the email address to determine the agent to route the email to
1916
+ * @param defaultAgentName The default agent name to use if the email address does not contain a sub-address
1917
+ * @returns A function that resolves the agent to route the email to
1918
+ */
1919
+ export function createAddressBasedEmailResolver<Env>(
1920
+ defaultAgentName: string
1921
+ ): EmailResolver<Env> {
1922
+ return async (email: ForwardableEmailMessage, _env: Env) => {
1923
+ const emailMatch = email.to.match(/^([^+@]+)(?:\+([^@]+))?@(.+)$/);
1924
+ if (!emailMatch) {
1925
+ return null;
1926
+ }
1927
+
1928
+ const [, localPart, subAddress] = emailMatch;
1929
+
1930
+ if (subAddress) {
1931
+ return {
1932
+ agentName: localPart,
1933
+ agentId: subAddress
1934
+ };
1935
+ }
1936
+
1937
+ // Option 2: Use defaultAgentName namespace, localPart as agentId
1938
+ // Common for catch-all email routing to a single EmailAgent namespace
1939
+ return {
1940
+ agentName: defaultAgentName,
1941
+ agentId: localPart
1942
+ };
1943
+ };
1944
+ }
1945
+
1946
+ /**
1947
+ * Create a resolver that uses the agentName and agentId to determine the agent to route the email to
1948
+ * @param agentName The name of the agent to route the email to
1949
+ * @param agentId The id of the agent to route the email to
1950
+ * @returns A function that resolves the agent to route the email to
1951
+ */
1952
+ export function createCatchAllEmailResolver<Env>(
1953
+ agentName: string,
1954
+ agentId: string
1955
+ ): EmailResolver<Env> {
1956
+ return async () => ({ agentName, agentId });
1957
+ }
1958
+
1959
+ export type EmailRoutingOptions<Env> = AgentOptions<Env> & {
1960
+ resolver: EmailResolver<Env>;
1961
+ };
1962
+
1963
+ // Cache the agent namespace map for email routing
1964
+ // This maps both kebab-case and original names to namespaces
1965
+ const agentMapCache = new WeakMap<
1966
+ Record<string, unknown>,
1967
+ Record<string, unknown>
1968
+ >();
1969
+
1174
1970
  /**
1175
1971
  * Route an email to the appropriate Agent
1176
- * @param email Email message to route
1177
- * @param env Environment containing Agent bindings
1178
- * @param options Routing options
1972
+ * @param email The email to route
1973
+ * @param env The environment containing the Agent bindings
1974
+ * @param options The options for routing the email
1975
+ * @returns A promise that resolves when the email has been routed
1179
1976
  */
1180
1977
  export async function routeAgentEmail<Env>(
1181
- _email: ForwardableEmailMessage,
1182
- _env: Env,
1183
- _options?: AgentOptions<Env>
1184
- ): Promise<void> {}
1978
+ email: ForwardableEmailMessage,
1979
+ env: Env,
1980
+ options: EmailRoutingOptions<Env>
1981
+ ): Promise<void> {
1982
+ const routingInfo = await options.resolver(email, env);
1983
+
1984
+ if (!routingInfo) {
1985
+ console.warn("No routing information found for email, dropping message");
1986
+ return;
1987
+ }
1988
+
1989
+ // Build a map that includes both original names and kebab-case versions
1990
+ if (!agentMapCache.has(env as Record<string, unknown>)) {
1991
+ const map: Record<string, unknown> = {};
1992
+ for (const [key, value] of Object.entries(env as Record<string, unknown>)) {
1993
+ if (
1994
+ value &&
1995
+ typeof value === "object" &&
1996
+ "idFromName" in value &&
1997
+ typeof value.idFromName === "function"
1998
+ ) {
1999
+ // Add both the original name and kebab-case version
2000
+ map[key] = value;
2001
+ map[camelCaseToKebabCase(key)] = value;
2002
+ }
2003
+ }
2004
+ agentMapCache.set(env as Record<string, unknown>, map);
2005
+ }
2006
+
2007
+ const agentMap = agentMapCache.get(env as Record<string, unknown>)!;
2008
+ const namespace = agentMap[routingInfo.agentName];
2009
+
2010
+ if (!namespace) {
2011
+ // Provide helpful error message listing available agents
2012
+ const availableAgents = Object.keys(agentMap)
2013
+ .filter((key) => !key.includes("-")) // Show only original names, not kebab-case duplicates
2014
+ .join(", ");
2015
+ throw new Error(
2016
+ `Agent namespace '${routingInfo.agentName}' not found in environment. Available agents: ${availableAgents}`
2017
+ );
2018
+ }
2019
+
2020
+ const agent = await getAgentByName(
2021
+ namespace as unknown as AgentNamespace<Agent<Env>>,
2022
+ routingInfo.agentId
2023
+ );
2024
+
2025
+ // let's make a serialisable version of the email
2026
+ const serialisableEmail: AgentEmail = {
2027
+ getRaw: async () => {
2028
+ const reader = email.raw.getReader();
2029
+ const chunks: Uint8Array[] = [];
2030
+
2031
+ let done = false;
2032
+ while (!done) {
2033
+ const { value, done: readerDone } = await reader.read();
2034
+ done = readerDone;
2035
+ if (value) {
2036
+ chunks.push(value);
2037
+ }
2038
+ }
2039
+
2040
+ const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
2041
+ const combined = new Uint8Array(totalLength);
2042
+ let offset = 0;
2043
+ for (const chunk of chunks) {
2044
+ combined.set(chunk, offset);
2045
+ offset += chunk.length;
2046
+ }
2047
+
2048
+ return combined;
2049
+ },
2050
+ headers: email.headers,
2051
+ rawSize: email.rawSize,
2052
+ setReject: (reason: string) => {
2053
+ email.setReject(reason);
2054
+ },
2055
+ forward: (rcptTo: string, headers?: Headers) => {
2056
+ return email.forward(rcptTo, headers);
2057
+ },
2058
+ reply: (options: { from: string; to: string; raw: string }) => {
2059
+ return email.reply(
2060
+ new EmailMessage(options.from, options.to, options.raw)
2061
+ );
2062
+ },
2063
+ from: email.from,
2064
+ to: email.to
2065
+ };
2066
+
2067
+ await agent._onEmail(serialisableEmail);
2068
+ }
2069
+
2070
+ export type AgentEmail = {
2071
+ from: string;
2072
+ to: string;
2073
+ getRaw: () => Promise<Uint8Array>;
2074
+ headers: Headers;
2075
+ rawSize: number;
2076
+ setReject: (reason: string) => void;
2077
+ forward: (rcptTo: string, headers?: Headers) => Promise<void>;
2078
+ reply: (options: { from: string; to: string; raw: string }) => Promise<void>;
2079
+ };
2080
+
2081
+ export type EmailSendOptions = {
2082
+ to: string;
2083
+ subject: string;
2084
+ body: string;
2085
+ contentType?: string;
2086
+ headers?: Record<string, string>;
2087
+ includeRoutingHeaders?: boolean;
2088
+ agentName?: string;
2089
+ agentId?: string;
2090
+ domain?: string;
2091
+ };
1185
2092
 
1186
2093
  /**
1187
2094
  * Get or create an Agent by name
@@ -1192,12 +2099,17 @@ export async function routeAgentEmail<Env>(
1192
2099
  * @param options Options for Agent creation
1193
2100
  * @returns Promise resolving to an Agent instance stub
1194
2101
  */
1195
- export async function getAgentByName<Env, T extends Agent<Env>>(
2102
+ export async function getAgentByName<
2103
+ Env,
2104
+ T extends Agent<Env>,
2105
+ Props extends Record<string, unknown> = Record<string, unknown>
2106
+ >(
1196
2107
  namespace: AgentNamespace<T>,
1197
2108
  name: string,
1198
2109
  options?: {
1199
2110
  jurisdiction?: DurableObjectJurisdiction;
1200
2111
  locationHint?: DurableObjectLocationHint;
2112
+ props?: Props;
1201
2113
  }
1202
2114
  ) {
1203
2115
  return getServerByName<Env, T>(namespace, name, options);
@@ -1229,7 +2141,7 @@ export class StreamingResponse {
1229
2141
  id: this._id,
1230
2142
  result: chunk,
1231
2143
  success: true,
1232
- type: "rpc",
2144
+ type: MessageType.RPC
1233
2145
  };
1234
2146
  this._connection.send(JSON.stringify(response));
1235
2147
  }
@@ -1248,7 +2160,7 @@ export class StreamingResponse {
1248
2160
  id: this._id,
1249
2161
  result: finalChunk,
1250
2162
  success: true,
1251
- type: "rpc",
2163
+ type: MessageType.RPC
1252
2164
  };
1253
2165
  this._connection.send(JSON.stringify(response));
1254
2166
  }