From 9036be74ad4234df61b5fb193b079db4d8d22527 Mon Sep 17 00:00:00 2001
From: "lichunqin.ads" <lichunqin.ads@bytedance.com>
Date: Thu, 20 Aug 2026 21:35:18 +0800
Subject: [PATCH] fix(agent): run afterAgentStep for final steps

---
 src/agent/__tests__/agent-lifecycle.test.ts | 43 +++++++++++++++++++++
 src/agent/agent.ts                          |  1 +
 2 files changed, 44 insertions(+)
 create mode 100644 src/agent/__tests__/agent-lifecycle.test.ts

diff --git a/src/agent/__tests__/agent-lifecycle.test.ts b/src/agent/__tests__/agent-lifecycle.test.ts
new file mode 100644
index 0000000..c01e3a7
--- /dev/null
+++ b/src/agent/__tests__/agent-lifecycle.test.ts
@@ -0,0 +1,43 @@
+import { expect, test } from "bun:test";
+
+import { Model, type ModelProvider, type UserMessage } from "@/foundation";
+
+import { Agent } from "../agent";
+import type { AgentMiddleware } from "../agent-middleware";
+
+const provider: ModelProvider = {
+  async invoke() {
+    return { role: "assistant", content: [{ type: "text", text: "done" }] };
+  },
+  async *stream() {
+    yield { role: "assistant", content: [{ type: "text", text: "done" }] };
+  },
+};
+
+test("runs afterAgentStep before afterAgentRun for a final step without tools", async () => {
+  const lifecycle: string[] = [];
+  const middleware: AgentMiddleware = {
+    afterAgentStep: async ({ step }) => {
+      lifecycle.push(`afterAgentStep:${step}`);
+    },
+    afterAgentRun: async () => {
+      lifecycle.push("afterAgentRun");
+    },
+  };
+  const agent = new Agent({
+    model: new Model("test-model", provider),
+    prompt: "test",
+    middlewares: [middleware],
+  });
+  const userMessage: UserMessage = {
+    role: "user",
+    content: [{ type: "text", text: "hello" }],
+  };
+
+  const stream = agent.stream(userMessage);
+  while (!(await stream.next()).done) {
+    // Drain the stream so the generator reaches the final lifecycle hooks.
+  }
+
+  expect(lifecycle).toEqual(["afterAgentStep:1", "afterAgentRun"]);
+});
diff --git a/src/agent/agent.ts b/src/agent/agent.ts
index 4535d82..b32eb83 100644
--- a/src/agent/agent.ts
+++ b/src/agent/agent.ts
@@ -156,6 +156,7 @@ export class Agent {
 
         const toolUses = this._extractToolUses(assistantMessage);
         if (toolUses.length === 0) {
+          await this._afterAgentStep(step);
           await this._afterAgentRun();
           return;
         }