feat(obs): basic agent instrumentation

Adds a new `observability` property to Agents that receives events during the Agent's exeuction that can be used for whatever. The default observability implementation prints the events to console, but in the future this could be hooked up to a UI or some other observability tooling.

Author: zebp

packages/agents/package.json

@@ -58,6 +58,11 @@
       "import": "./dist/mcp/do-oauth-client-provider.js",
       "require": "./dist/mcp/do-oauth-client-provider.js"
     },
+    "./observability": {
+      "import": "./dist/observability/index.js",
+      "require": "./dist/observability/index.js",
+      "types": "./dist/observability/index.d.ts"
+    },
     "./react": {
       "types": "./dist/react.d.ts",
       "import": "./dist/react.js",

packages/agents/scripts/build.ts

@@ -11,6 +11,7 @@ async function main() {
       "src/mcp/index.ts",
       "src/mcp/client.ts",
       "src/mcp/do-oauth-client-provider.ts",
+      "src/observability/index.ts",
     ],
     external: [
       "cloudflare:workers",

packages/agents/src/ai-chat-agent.ts

@@ -76,8 +76,23 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
           },
           [connection.id]
         );
+
+        const incomingMessages = this._messagesNotAlreadyInAgent(messages);
         await this.persistMessages(messages, [connection.id]);
 
+        this.observability?.emit(
+          {
+            displayMessage: "Chat message request",
+            id: data.id,
+            payload: {
+              message: incomingMessages,
+            },
+            timestamp: Date.now(),
+            type: "message:request",
+          },
+          this.ctx
+        );
+
         const chatMessageId = data.id;
         const abortSignal = this._getAbortSignal(chatMessageId);
 
@@ -89,8 +104,23 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
                 responseMessages: response.messages,
               });
 
+              const outgoingMessages =
+                this._messagesNotAlreadyInAgent(finalMessages);
               await this.persistMessages(finalMessages, [connection.id]);
               this._removeAbortController(chatMessageId);
+
+              this.observability?.emit(
+                {
+                  displayMessage: "Chat message response",
+                  id: data.id,
+                  payload: {
+                    message: outgoingMessages,
+                  },
+                  timestamp: Date.now(),
+                  type: "message:response",
+                },
+                this.ctx
+              );
             },
             abortSignal ? { abortSignal } : undefined
           );
@@ -219,6 +249,11 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
     );
   }
 
+  private _messagesNotAlreadyInAgent(messages: ChatMessage[]) {
+    const existingIds = new Set(this.messages.map((message) => message.id));
+    return messages.filter((message) => !existingIds.has(message.id));
+  }
+
   private async _reply(id: string, response: Response) {
     // now take chunks out from dataStreamResponse and send them to the client
     return this._tryCatchChat(async () => {

packages/agents/src/index.ts

@@ -23,6 +23,7 @@ import { camelCaseToKebabCase } from "./client";
 import { MCPClientManager } from "./mcp/client";
 // import type { MCPClientConnection } from "./mcp/client-connection";
 import { DurableObjectOAuthClientProvider } from "./mcp/do-oauth-client-provider";
+import { genericObservability, type Observability } from "./observability";
 
 export type { Connection, ConnectionContext, WSMessage } from "partyserver";
 
@@ -316,6 +317,11 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     hibernate: true, // default to hibernate
   };
 
+  /**
+   * The observability implementation to use for the Agent
+   */
+  observability?: Observability = genericObservability;
+
   /**
    * Execute SQL queries against the Agent's database
    * @template T Type of the returned rows
@@ -460,6 +466,23 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
               // For regular methods, execute and send response
               const result = await methodFn.apply(this, args);
+
+              this.observability?.emit(
+                {
+                  displayMessage: `RPC call to ${method}`,
+                  id: nanoid(),
+                  payload: {
+                    args,
+                    method,
+                    streaming: metadata?.streaming,
+                    success: true,
+                  },
+                  timestamp: Date.now(),
+                  type: "rpc",
+                },
+                this.ctx
+              );
+
               const response: RPCResponse = {
                 done: true,
                 id,
@@ -512,6 +535,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
               })
             );
 
+            this.observability?.emit(
+              {
+                displayMessage: "Connection established",
+                id: nanoid(),
+                payload: {
+                  connectionId: connection.id,
+                },
+                timestamp: Date.now(),
+                type: "connect",
+              },
+              this.ctx
+            );
             return this._tryCatch(() => _onConnect(connection, ctx));
           }, 20);
         }
@@ -561,6 +596,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     state: State,
     source: Connection | "server" = "server"
   ) {
+    const previousState = this._state;
     this._state = state;
     this.sql`
     INSERT OR REPLACE INTO cf_agents_state (id, state)
@@ -582,6 +618,19 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       return agentContext.run(
         { agent: this, connection, request },
         async () => {
+          this.observability?.emit(
+            {
+              displayMessage: "State updated",
+              id: nanoid(),
+              payload: {
+                previousState,
+                state,
+              },
+              timestamp: Date.now(),
+              type: "state:update",
+            },
+            this.ctx
+          );
           return this.onStateUpdate(state, source);
         }
       );
@@ -677,6 +726,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   ): Promise<Schedule<T>> {
     const id = nanoid(9);
 
+    const emitScheduleCreate = (schedule: Schedule<T>) =>
+      this.observability?.emit(
+        {
+          displayMessage: `Schedule ${schedule.id} created`,
+          id: nanoid(),
+          payload: schedule,
+          timestamp: Date.now(),
+          type: "schedule:create",
+        },
+        this.ctx
+      );
+
     if (typeof callback !== "string") {
       throw new Error("Callback must be a string");
     }
@@ -696,13 +757,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this._scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         callback: callback,
         id,
         payload: payload as T,
         time: timestamp,
         type: "scheduled",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     if (typeof when === "number") {
       const time = new Date(Date.now() + when * 1000);
@@ -717,14 +782,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this._scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         callback: callback,
         delayInSeconds: when,
         id,
         payload: payload as T,
         time: timestamp,
         type: "delayed",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     if (typeof when === "string") {
       const nextExecutionTime = getNextCronTime(when);
@@ -739,14 +808,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this._scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         callback: callback,
         cron: when,
         id,
         payload: payload as T,
         time: timestamp,
         type: "cron",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     throw new Error("Invalid schedule type");
   }
@@ -822,6 +895,19 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * @returns true if the task was cancelled, false otherwise
    */
   async cancelSchedule(id: string): Promise<boolean> {
+    const schedule = await this.getSchedule(id);
+    if (schedule) {
+      this.observability?.emit(
+        {
+          displayMessage: `Schedule ${id} cancelled`,
+          id: nanoid(),
+          payload: schedule,
+          timestamp: Date.now(),
+          type: "schedule:cancel",
+        },
+        this.ctx
+      );
+    }
     this.sql`DELETE FROM cf_agents_schedules WHERE id = ${id}`;
 
     await this._scheduleNextAlarm();
@@ -870,6 +956,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         { agent: this, connection: undefined, request: undefined },
         async () => {
           try {
+            this.observability?.emit(
+              {
+                displayMessage: `Schedule ${row.id} executed`,
+                id: nanoid(),
+                payload: row,
+                timestamp: Date.now(),
+                type: "schedule:execute",
+              },
+              this.ctx
+            );
+
             await (
               callback as (
                 payload: unknown,
@@ -914,6 +1011,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     await this.ctx.storage.deleteAlarm();
     await this.ctx.storage.deleteAll();
     this.ctx.abort("destroyed"); // enforce that the agent is evicted
+
+    this.observability?.emit(
+      {
+        displayMessage: "Agent destroyed",
+        id: nanoid(),
+        payload: {},
+        timestamp: Date.now(),
+        type: "destroy",
+      },
+      this.ctx
+    );
   }
 
   /**