feat(obs): basic agent instrumentation

zebp · zebp · commit fc015af96a61 · 2025-04-30T15:06:46.000-04:00
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.
diff --git a/packages/agents/package.json b/packages/agents/package.json
@@ -66,6 +66,11 @@
       "types": "./dist/mcp/do-oauth-client-provider.d.ts",
       "require": "./dist/mcp/do-oauth-client-provider.js",
       "import": "./dist/mcp/do-oauth-client-provider.js"
+    },
+    "./observability": {
+      "types": "./dist/observability/index.d.ts",
+      "require": "./dist/observability/index.js",
+      "import": "./dist/observability/index.js"
     }
   },
   "keywords": [],
diff --git a/packages/agents/scripts/build.ts b/packages/agents/scripts/build.ts
@@ -9,6 +9,7 @@ async function main() {
       "src/mcp/index.ts",
       "src/mcp/client.ts",
       "src/mcp/do-oauth-client-provider.ts",
+      "src/observability/index.ts",
     ],
     splitting: true,
     sourcemap: true,
diff --git a/packages/agents/src/ai-chat-agent.ts b/packages/agents/src/ai-chat-agent.ts
@@ -73,6 +73,7 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
           // dispatcher,
           // duplex
         } = data.init;
+
         const { messages } = JSON.parse(body as string);
         this.#broadcastChatMessage(
           {
@@ -81,8 +82,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(
+          {
+            type: "message:request",
+            id: data.id,
+            displayMessage: `Incoming chat message:\n${incomingMessages.map((m) => `${m.role}: ${m.content}`).join("\n")}`,
+            timestamp: Date.now(),
+            payload: {
+              message: incomingMessages,
+            },
+          },
+          this.ctx
+        );
+
         const chatMessageId = data.id;
         const abortSignal = this.#getAbortSignal(chatMessageId);
 
@@ -94,8 +110,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(
+                {
+                  type: "message:response",
+                  id: data.id,
+                  displayMessage: `Outgoing chat message:\n${outgoingMessages.map((m) => `${m.role}: ${m.content}`).join("\n")}`,
+                  timestamp: Date.now(),
+                  payload: {
+                    message: outgoingMessages,
+                  },
+                },
+                this.ctx
+              );
             },
             abortSignal ? { abortSignal } : undefined
           );
@@ -207,6 +238,11 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
     );
   }
 
+  #messagesNotAlreadyInAgent(messages: ChatMessage[]) {
+    const existingIds = new Set(this.messages.map((message) => message.id));
+    return messages.filter((message) => !existingIds.has(message.id));
+  }
+
   async #reply(id: string, response: Response) {
     // now take chunks out from dataStreamResponse and send them to the client
     return this.#tryCatch(async () => {
diff --git a/packages/agents/src/index.ts b/packages/agents/src/index.ts
@@ -13,6 +13,7 @@ import { nanoid } from "nanoid";
 
 import { AsyncLocalStorage } from "node:async_hooks";
 import { MCPClientManager } from "./mcp/client";
+import { genericObservability, type Observability } from "./observability";
 
 export type { Connection, WSMessage, ConnectionContext } from "partyserver";
 
@@ -264,6 +265,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
@@ -369,6 +375,35 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
               // For regular methods, execute and send response
               const result = await methodFn.apply(this, args);
+
+              const displayArgs = args.map(a =>{
+                if (typeof a === "object") {
+                  if (Array.isArray(a)) {
+                    return "[...]"
+                  }
+
+                  return "{...}";
+                }
+
+                return String(a);
+              });
+
+              this.observability?.emit(
+                {
+                  id: nanoid(),
+                  type: "rpc",
+                  displayMessage: `RPC call to ${method} args: ${displayArgs.join(", ")}`,
+                  timestamp: Date.now(),
+                  payload: {
+                    method,
+                    args,
+                    success: true,
+                    streaming: metadata?.streaming,
+                  },
+                },
+                this.ctx
+              );
+
               const response: RPCResponse = {
                 type: "rpc",
                 id,
@@ -413,6 +448,19 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                 })
               );
             }
+
+            this.observability?.emit(
+              {
+                id: nanoid(),
+                type: "connect",
+                displayMessage: `Connection ${connection.id} established`,
+                timestamp: Date.now(),
+                payload: {
+                  connectionId: connection.id,
+                },
+              },
+              this.ctx
+            );
             return this.#tryCatch(() => _onConnect(connection, ctx));
           }, 20);
         }
@@ -421,6 +469,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   }
 
   #setStateInternal(state: State, source: Connection | "server" = "server") {
+    const previousState = this.#state;
     this.#state = state;
     this.sql`
     INSERT OR REPLACE INTO cf_agents_state (id, state)
@@ -442,6 +491,19 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       return agentContext.run(
         { agent: this, connection, request },
         async () => {
+          this.observability?.emit(
+            {
+              id: nanoid(),
+              type: "state:update",
+              displayMessage: "State updated",
+              timestamp: Date.now(),
+              payload: {
+                state,
+                previousState,
+              },
+            },
+            this.ctx
+          );
           return this.onStateUpdate(state, source);
         }
       );
@@ -535,6 +597,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(
+        {
+          id: nanoid(),
+          type: "schedule:create",
+          displayMessage: `Schedule ${schedule.id} created`,
+          timestamp: Date.now(),
+          payload: schedule,
+        },
+        this.ctx
+      );
+
     if (typeof callback !== "string") {
       throw new Error("Callback must be a string");
     }
@@ -554,13 +628,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this.#scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         id,
         callback: callback,
         payload: payload as T,
         time: timestamp,
         type: "scheduled",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     if (typeof when === "number") {
       const time = new Date(Date.now() + when * 1000);
@@ -575,14 +653,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this.#scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         id,
         callback: callback,
         payload: payload as T,
         delayInSeconds: when,
         time: timestamp,
         type: "delayed",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     if (typeof when === "string") {
       const nextExecutionTime = getNextCronTime(when);
@@ -597,14 +679,18 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
       await this.#scheduleNextAlarm();
 
-      return {
+      const schedule: Schedule<T> = {
         id,
         callback: callback,
         payload: payload as T,
         cron: when,
         time: timestamp,
         type: "cron",
       };
+
+      emitScheduleCreate(schedule);
+
+      return schedule;
     }
     throw new Error("Invalid schedule type");
   }
@@ -680,6 +766,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(
+        {
+          id: nanoid(),
+          type: "schedule:delete",
+          displayMessage: `Schedule ${id} deleted`,
+          timestamp: Date.now(),
+          payload: schedule,
+        },
+        this.ctx
+      );
+    }
     this.sql`DELETE FROM cf_agents_schedules WHERE id = ${id}`;
 
     await this.#scheduleNextAlarm();
@@ -724,6 +823,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         { agent: this, connection: undefined, request: undefined },
         async () => {
           try {
+            this.observability?.emit(
+              {
+                id: nanoid(),
+                type: "schedule:execute",
+                displayMessage: `Schedule ${row.id} executed`,
+                timestamp: Date.now(),
+                payload: row,
+              },
+              this.ctx
+            );
+
             await (
               callback as (
                 payload: unknown,
@@ -766,6 +876,17 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     // delete all alarms
     await this.ctx.storage.deleteAlarm();
     await this.ctx.storage.deleteAll();
+
+    this.observability?.emit(
+      {
+        id: nanoid(),
+        type: "destroy",
+        displayMessage: "Agent destroyed",
+        timestamp: Date.now(),
+        payload: {},
+      },
+      this.ctx
+    );
   }
 
   /**
diff --git a/packages/agents/src/observability/index.ts b/packages/agents/src/observability/index.ts
