CR feedback

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Author: hubertzub-db

integrations/appkit-agent/src/agent-plugin/agent.ts

@@ -239,9 +239,12 @@ export class AgentPlugin extends Plugin<IAgentConfig> {
         if (!this.agentImpl) {
           throw new Error("AgentPlugin not initialized");
         }
-        const lastUser = [...messages].reverse().find((m) => m.role === "user");
-        const input = lastUser?.content ?? "";
-        const chatHistory = messages.slice(0, -1);
+        const lastUserIdx = messages.findLastIndex((m) => m.role === "user");
+        if (lastUserIdx < 0) {
+          throw new Error("No user message found in input");
+        }
+        const input = messages[lastUserIdx].content;
+        const chatHistory = messages.slice(0, lastUserIdx);
         const items = await this.agentImpl.invoke({
           input,
           chat_history: chatHistory,
@@ -258,9 +261,12 @@ export class AgentPlugin extends Plugin<IAgentConfig> {
         if (!this.agentImpl) {
           throw new Error("AgentPlugin not initialized");
         }
-        const lastUser = [...messages].reverse().find((m) => m.role === "user");
-        const input = lastUser?.content ?? "";
-        const chatHistory = messages.slice(0, -1);
+        const lastUserIdx = messages.findLastIndex((m) => m.role === "user");
+        if (lastUserIdx < 0) {
+          throw new Error("No user message found in input");
+        }
+        const input = messages[lastUserIdx].content;
+        const chatHistory = messages.slice(0, lastUserIdx);
         yield* this.agentImpl.stream({
           input,
           chat_history: chatHistory,

integrations/appkit-agent/src/agent-plugin/invoke-handler.ts

@@ -123,7 +123,10 @@ export function createInvokeHandler(
         userInput = lastUserMessage.content as string;
       }
 
-      const chatHistory = input.slice(0, -1).map(flattenHistoryItem);
+      const lastUserIdx = input.findLastIndex(
+        (msg: any) => msg.role === "user",
+      );
+      const chatHistory = input.slice(0, lastUserIdx).map(flattenHistoryItem);
 
       const agentParams = { input: userInput, chat_history: chatHistory };
       const agent = getAgent();

integrations/appkit-agent/src/tests/agent-plugin/agent.test.ts

@@ -164,6 +164,37 @@ describe("AgentPlugin", () => {
       expect(completedEvent).toBeDefined();
     });
 
+    test("invoke history contains only messages before the last user message", async () => {
+      const spyAgent = {
+        invoke: vi.fn().mockResolvedValue([
+          {
+            id: "msg_1",
+            type: "message",
+            role: "assistant",
+            status: "completed",
+            content: [{ type: "output_text", text: "ok", annotations: [] }],
+          },
+        ]),
+        stream: vi.fn(),
+      };
+      const config: IAgentConfig = { agentInstance: spyAgent as any };
+      const plugin = new AgentPlugin(config);
+      await plugin.setup();
+
+      await plugin.exports().invoke([
+        { role: "user", content: "first" },
+        { role: "user", content: "second" },
+        { role: "assistant", content: "interrupted" },
+      ]);
+
+      expect(spyAgent.invoke).toHaveBeenCalledWith(
+        expect.objectContaining({
+          input: "second",
+          chat_history: [{ role: "user", content: "first" }],
+        }),
+      );
+    });
+
     test("throws when not initialized", async () => {
       const config: IAgentConfig = { agentInstance: new StubAgent() };
       const plugin = new AgentPlugin(config);
@@ -174,6 +205,29 @@ describe("AgentPlugin", () => {
     });
   });
 
+  describe("abortActiveOperations()", () => {
+    function injectMcpClient(
+      plugin: AgentPlugin,
+      close: ReturnType<typeof vi.fn>,
+    ) {
+      Reflect.set(plugin, "mcpClient", { getTools: vi.fn(), close });
+    }
+
+    test("closes MCP client when present", async () => {
+      const stub = new StubAgent();
+      const config: IAgentConfig = { agentInstance: stub };
+      const plugin = new AgentPlugin(config);
+      await plugin.setup();
+
+      const closeFn = vi.fn().mockResolvedValue(undefined);
+      injectMcpClient(plugin, closeFn);
+
+      await plugin.abortActiveOperations();
+
+      expect(closeFn).toHaveBeenCalledOnce();
+    });
+  });
+
   describe("addTools()", () => {
     test("throws when using agentInstance mode", async () => {
       const stub = new StubAgent();

integrations/appkit-agent/src/tests/agent-plugin/invoke-handler.test.ts

@@ -206,6 +206,48 @@ describe("createInvokeHandler", () => {
       );
     });
 
+    test("history contains only messages before the last user message", async () => {
+      const spyAgent = {
+        invoke: vi.fn().mockResolvedValue([
+          {
+            id: "msg_1",
+            type: "message",
+            role: "assistant",
+            status: "completed",
+            content: [
+              { type: "output_text", text: "response", annotations: [] },
+            ],
+          },
+        ]),
+        stream: vi.fn(),
+      };
+      const historyHandler = createInvokeHandler(() => spyAgent as any);
+
+      const req = makeReq({
+        input: [
+          { role: "user", content: "First message" },
+          { role: "user", content: "Second message" },
+          { role: "assistant", content: "Interrupted reply" },
+        ],
+        stream: false,
+      });
+      const res = makeRes();
+
+      await historyHandler(req, res, vi.fn());
+
+      expect(spyAgent.invoke).toHaveBeenCalledWith(
+        expect.objectContaining({
+          input: "Second message",
+          chat_history: [
+            expect.objectContaining({
+              role: "user",
+              content: "First message",
+            }),
+          ],
+        }),
+      );
+    });
+
     test("handles function_call items in history", async () => {
       const spyAgent = {
         invoke: vi.fn().mockResolvedValue([
@@ -278,6 +320,30 @@ describe("createInvokeHandler", () => {
       );
     });
 
+    test("returns 500 when non-streaming agent invoke throws", async () => {
+      const errorAgent = {
+        invoke: vi.fn().mockRejectedValue(new Error("invoke exploded")),
+        stream: vi.fn(),
+      };
+      const errorHandler = createInvokeHandler(() => errorAgent as any);
+
+      const req = makeReq({
+        input: [{ role: "user", content: "boom" }],
+        stream: false,
+      });
+      const res = makeRes();
+
+      await errorHandler(req, res, vi.fn());
+
+      expect(res.status).toHaveBeenCalledWith(500);
+      expect(res.json).toHaveBeenCalledWith(
+        expect.objectContaining({
+          error: "Internal server error",
+          message: "invoke exploded",
+        }),
+      );
+    });
+
     test("handles agent errors gracefully in streaming mode", async () => {
       const errorAgent = {
         invoke: vi.fn(),