From 66417b9f9b58ab141cc7aef4d78ceb05aabe6f37 Mon Sep 17 00:00:00 2001 From: Amruth Pillai Date: Wed, 13 May 2026 22:37:45 +0200 Subject: [PATCH] fix(agent): abort run on archive and verify ownership before deleting thread - threads.archive: before flipping status, abort any in-flight run controller and clear the active-run state on the thread; cleanup failures are logged but do not block the status update. - threads.delete: assert thread ownership via getThread before destructive work so an authenticated user cannot wipe another user's attachment rows by passing a foreign threadId. Adds focused tests for both behaviors. --- packages/api/src/services/agent.test.ts | 177 +++++++++++++++++++++++- packages/api/src/services/agent.ts | 21 +++ 2 files changed, 195 insertions(+), 3 deletions(-) diff --git a/packages/api/src/services/agent.test.ts b/packages/api/src/services/agent.test.ts index df4697009..31903bcc5 100644 --- a/packages/api/src/services/agent.test.ts +++ b/packages/api/src/services/agent.test.ts @@ -8,6 +8,14 @@ const dbMock = { delete: vi.fn(), }; +const clearActiveAgentRunIfCurrentMock = vi.fn(); +const claimActiveAgentRunMock = vi.fn(); +const storageServiceMock = { + delete: vi.fn(), + write: vi.fn(), + read: vi.fn(), +}; + const resumeServiceMock = { getById: vi.fn(), }; @@ -77,15 +85,15 @@ vi.mock("./ai", () => ({ getAgentModel: vi.fn() })); vi.mock("./ai-credentials", () => ({ assertAgentEnvironment: vi.fn() })); vi.mock("./ai-providers", () => ({ aiProvidersService: aiProvidersServiceMock })); vi.mock("./resume", () => ({ resumeService: resumeServiceMock })); -vi.mock("./storage", () => ({ getStorageService: vi.fn(), inferContentType: vi.fn() })); +vi.mock("./storage", () => ({ getStorageService: vi.fn(() => storageServiceMock), inferContentType: vi.fn() })); vi.mock("./agent-patches", () => ({ createInverseResumePatches: vi.fn() })); vi.mock("./agent-resume", () => ({ buildAgentDraftResumeName: vi.fn(), buildUniqueAgentDraftSlug: vi.fn(), })); vi.mock("./agent-run-state", () => ({ - claimActiveAgentRun: vi.fn(), - clearActiveAgentRunIfCurrent: vi.fn(), + claimActiveAgentRun: claimActiveAgentRunMock, + clearActiveAgentRunIfCurrent: clearActiveAgentRunIfCurrentMock, })); vi.mock("./agent-streams", () => ({ agentStreamLifecycle: { create: vi.fn(), resume: vi.fn() }, @@ -196,3 +204,166 @@ describe("agentService.messages.send", () => { expect(aiProvidersServiceMock.getRunnableById).not.toHaveBeenCalled(); }); }); + +describe("agentService.threads.archive", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("skips active-run cleanup when there is no active run", async () => { + const idleThread = buildArchivedThread({ + status: "active", + activeRunId: null, + activeStreamId: null, + archivedAt: null, + }); + + dbMock.select.mockImplementation(() => { + const limit = vi.fn(async () => [idleThread]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + return { from }; + }); + + const updateWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updateWhere })); + dbMock.update.mockReturnValue({ set: updateSet }); + + const { agentService } = await import("./agent"); + + await agentService.threads.archive({ id: "thread-1", userId: "user-1" }); + + expect(clearActiveAgentRunIfCurrentMock).not.toHaveBeenCalled(); + expect(updateSet).toHaveBeenCalledWith(expect.objectContaining({ status: "archived" })); + expect(updateWhere).toHaveBeenCalled(); + }); + + it("clears active-run state when an active run is present, then flips status", async () => { + const activeThread = buildArchivedThread({ + status: "active", + activeRunId: "run-1", + activeStreamId: "stream-1", + archivedAt: null, + }); + + dbMock.select.mockImplementation(() => { + const limit = vi.fn(async () => [activeThread]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + return { from }; + }); + + const updateWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updateWhere })); + dbMock.update.mockReturnValue({ set: updateSet }); + + clearActiveAgentRunIfCurrentMock.mockResolvedValue(undefined); + + const { agentService } = await import("./agent"); + + await agentService.threads.archive({ id: "thread-1", userId: "user-1" }); + + expect(clearActiveAgentRunIfCurrentMock).toHaveBeenCalledWith({ + threadId: "thread-1", + userId: "user-1", + runId: "run-1", + streamId: "stream-1", + }); + expect(updateSet).toHaveBeenCalledWith(expect.objectContaining({ status: "archived" })); + expect(updateWhere).toHaveBeenCalled(); + }); + + it("still flips status when clearActiveAgentRunIfCurrent throws", async () => { + const activeThread = buildArchivedThread({ + status: "active", + activeRunId: "run-2", + activeStreamId: "stream-2", + archivedAt: null, + }); + + dbMock.select.mockImplementation(() => { + const limit = vi.fn(async () => [activeThread]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + return { from }; + }); + + const updateWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updateWhere })); + dbMock.update.mockReturnValue({ set: updateSet }); + + clearActiveAgentRunIfCurrentMock.mockRejectedValue(new Error("boom")); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => undefined); + + const { agentService } = await import("./agent"); + + await agentService.threads.archive({ id: "thread-1", userId: "user-1" }); + + expect(clearActiveAgentRunIfCurrentMock).toHaveBeenCalled(); + expect(updateSet).toHaveBeenCalledWith(expect.objectContaining({ status: "archived" })); + expect(consoleSpy).toHaveBeenCalled(); + + consoleSpy.mockRestore(); + }); +}); + +describe("agentService.threads.delete", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("throws NOT_FOUND when the thread does not belong to the user and skips destructive work", async () => { + dbMock.select.mockImplementation(() => { + const limit = vi.fn(async () => []); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + return { from }; + }); + + const { agentService } = await import("./agent"); + + const deleting = agentService.threads.delete({ id: "thread-x", userId: "user-y" }); + + await expect(deleting).rejects.toBeInstanceOf(ORPCError); + await expect(deleting).rejects.toMatchObject({ code: "NOT_FOUND" }); + + expect(storageServiceMock.delete).not.toHaveBeenCalled(); + expect(dbMock.delete).not.toHaveBeenCalled(); + expect(dbMock.update).not.toHaveBeenCalled(); + }); + + it("proceeds with cleanup when the thread is owned by the user", async () => { + const ownedThread = buildArchivedThread({ + id: "thread-own", + userId: "user-own", + status: "active", + activeRunId: null, + activeStreamId: null, + archivedAt: null, + }); + + dbMock.select.mockImplementation(() => { + const limit = vi.fn(async () => [ownedThread]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + return { from }; + }); + + const deleteWhere = vi.fn(async () => undefined); + dbMock.delete.mockReturnValue({ where: deleteWhere }); + + const updateWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updateWhere })); + dbMock.update.mockReturnValue({ set: updateSet }); + + storageServiceMock.delete.mockResolvedValue(undefined); + + const { agentService } = await import("./agent"); + + await agentService.threads.delete({ id: "thread-own", userId: "user-own" }); + + expect(storageServiceMock.delete).toHaveBeenCalledWith("uploads/user-own/agent/thread-own"); + expect(dbMock.delete).toHaveBeenCalled(); + expect(updateSet).toHaveBeenCalledWith(expect.objectContaining({ status: "deleted" })); + }); +}); diff --git a/packages/api/src/services/agent.ts b/packages/api/src/services/agent.ts index 2a92c5fba..93ac9e169 100644 --- a/packages/api/src/services/agent.ts +++ b/packages/api/src/services/agent.ts @@ -482,6 +482,25 @@ export const agentService = { archive: async (input: { id: string; userId: string }) => { assertAgentEnvironment(); + const thread = await getThread({ id: input.id, userId: input.userId }); + const activeRunId = thread.activeRunId; + const activeStreamId = thread.activeStreamId; + + if (activeRunId) { + activeRunControllers.get(activeRunId)?.abort("USER_ARCHIVED"); + activeRunControllers.delete(activeRunId); + try { + await clearActiveAgentRunIfCurrent({ + threadId: input.id, + userId: input.userId, + runId: activeRunId, + streamId: activeStreamId, + }); + } catch (error) { + console.error("[agent] Failed to clear active run during archive", error); + } + } + await db .update(schema.agentThread) .set({ status: "archived", archivedAt: new Date() }) @@ -491,6 +510,8 @@ export const agentService = { delete: async (input: { id: string; userId: string }) => { assertAgentEnvironment(); + await getThread({ id: input.id, userId: input.userId }); + await getStorageService().delete(`uploads/${input.userId}/agent/${input.id}`); await db.delete(schema.agentAttachment).where(eq(schema.agentAttachment.threadId, input.id)); await db