diff --git a/prisma/migrations/20260720161314_add_active_runs_to_shared_conversation/migration.sql b/prisma/migrations/20260720161314_add_active_runs_to_shared_conversation/migration.sql new file mode 100644 index 0000000000..6dacf185c0 --- /dev/null +++ b/prisma/migrations/20260720161314_add_active_runs_to_shared_conversation/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "shared_conversations" ADD COLUMN "active_runs" JSONB; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 08bf97deb0..f33bc2034d 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -653,6 +653,10 @@ model SharedConversation { ownerSeenAt DateTime? @map("owner_seen_at") source String? settings Json? @default("{}") + // Map of requestId → { requestId, workspaceId, startedAt, abortRequested? } for in-flight + // repo_agent runs. Nullable so rows without active runs carry no overhead. + // abortIntents: { [turnId]: { turnId, requestedAt, expiresAt } } for start-race cancellation. + activeRuns Json? @map("active_runs") user User? @relation(fields: [userId], references: [id], onDelete: Cascade) workspace Workspace? @relation(fields: [workspaceId], references: [id], onDelete: Cascade) sourceControlOrg SourceControlOrg? @relation(fields: [sourceControlOrgId], references: [id], onDelete: Cascade) diff --git a/src/__tests__/unit/lib/ai/repoAgent.test.ts b/src/__tests__/unit/lib/ai/repoAgent.test.ts new file mode 100644 index 0000000000..7e193ec050 --- /dev/null +++ b/src/__tests__/unit/lib/ai/repoAgent.test.ts @@ -0,0 +1,268 @@ +/** + * Unit tests for `repoAgent()` in `src/lib/ai/askTools.ts`. + * + * Coverage: + * 1. completed-with-result after abort → returns real result (Req 5) + * 2. failed after abort → returns cancelled marker (not a throw) + * 3. aborted status after abort → returns cancelled marker (not a throw) + * 4. grace-window exit fires when swarm never reports a terminal status + * 5. non-aborted completed path → returns result normally + * 6. non-aborted failed path → throws + * 7. non-aborted timeout path → throws + * 8. onRequestId hook is called with the request_id immediately + * 9. isAbortRequested is called each poll cycle + * 10. repoAgent module has no db/Prisma import + */ + +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// ── Mocks ──────────────────────────────────────────────────────────────────── +// We mock fetch globally so we can control initiate + progress responses. + +// Also need to mock the deep-import deps pulled in by askTools at module level. +vi.mock("gitsee/server", () => ({ RepoAnalyzer: vi.fn() })); +vi.mock("@/lib/ai/provider", () => ({ getProviderTool: vi.fn() })); +vi.mock("@ai-sdk/mcp", () => ({ createMCPClient: vi.fn() })); +vi.mock("@/lib/ai/mcpTimeout", () => ({ + withMcpTimeout: vi.fn(), + isMcpTimeout: vi.fn(), +})); +vi.mock("@/lib/mcp/mcpTools", () => ({ + mcpListFeatures: vi.fn(), + mcpReadFeature: vi.fn(), + mcpListTasks: vi.fn(), + mcpReadTask: vi.fn(), + mcpCheckStatus: vi.fn(), + findWorkspaceUser: vi.fn(), +})); +vi.mock("@/services/bifrost/orchestrator", () => ({ + getBifrostForLLM: vi.fn().mockResolvedValue(undefined), +})); +vi.mock("@/lib/ai/concepts", () => ({ swarmFetch: vi.fn() })); +vi.mock("@/lib/ai/mcpResult", () => ({ + mcpText: vi.fn(), + capMcpResult: vi.fn(), +})); + +// ── Import after mocks ──────────────────────────────────────────────────────── +import { repoAgent, REPO_AGENT_CANCELLED } from "@/lib/ai/askTools"; + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +/** Mock a successful initiate response. */ +function mockInitiate(requestId = "req-1") { + return { + ok: true, + json: () => Promise.resolve({ request_id: requestId }), + text: () => Promise.resolve(""), + }; +} + +/** Mock a progress response with a given status and optional result/error. */ +function mockProgress(status: string, result?: Record, error?: string) { + return { + ok: true, + json: () => Promise.resolve({ status, result: result ?? null, error }), + }; +} + +const SWARM_URL = "https://swarm.test"; +const SWARM_KEY = "test-key"; +const BASE_PARAMS = { prompt: "What does the code do?" }; + +/** + * Build a fetch mock sequence: initiate response followed by progress responses. + */ +function buildFetch(progressSequence: Array>) { + let call = 0; + return vi.fn().mockImplementation(() => { + if (call === 0) { + call++; + return Promise.resolve(mockInitiate()); + } + const idx = call - 1; // progress call index + call++; + return Promise.resolve(progressSequence[Math.min(idx, progressSequence.length - 1)]); + }); +} + +beforeEach(() => { + vi.useFakeTimers(); +}); + +afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); +}); + +/** + * Advance timers in tight loop to let poll `setTimeout` resolve. + * Returns a promise that resolves once the repoAgent promise settles. + */ +async function drainPolls(agentPromise: Promise, maxCycles = 10): Promise { + for (let i = 0; i < maxCycles; i++) { + // Advance past the 5-second poll interval. + await vi.advanceTimersByTimeAsync(5100); + } + return agentPromise; +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +describe("repoAgent()", () => { + it("1. returns the real result when run completes after abort was requested (Req 5)", async () => { + let abortFlag = false; + const fetchMock = buildFetch([ + // First two polls: still running, abort is requested on poll 2 + mockProgress("running"), + mockProgress("completed", { content: "real answer" }), + ]); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { + isAbortRequested: async () => { + // Returns true starting from the second poll invocation + abortFlag = true; + return abortFlag; + }, + }); + + const result = await drainPolls(promise, 3); + expect(result).toEqual({ content: "real answer" }); + expect(result).not.toBe(REPO_AGENT_CANCELLED); + }); + + it("2. returns cancelled marker (not a throw) when run fails after abort", async () => { + const fetchMock = buildFetch([ + mockProgress("running"), + mockProgress("failed", undefined, "some error"), + ]); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { + isAbortRequested: async () => true, // abort from the start + }); + + const result = await drainPolls(promise, 3); + expect(result).toBe(REPO_AGENT_CANCELLED); + }); + + it("3. returns cancelled marker (not a throw) when swarm reports 'aborted' status", async () => { + const fetchMock = buildFetch([ + mockProgress("aborted"), + ]); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { + isAbortRequested: async () => true, + }); + + const result = await drainPolls(promise, 2); + expect(result).toBe(REPO_AGENT_CANCELLED); + }); + + it("4. grace-window exit fires when swarm never returns a terminal status", async () => { + // Always returns "running"; abort is requested from cycle 1. + const fetchMock = buildFetch(Array(20).fill(mockProgress("running"))); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { + isAbortRequested: async () => true, + }); + + // ABORT_GRACE_POLL_CYCLES = 3, so after 3 cycles we expect exit. + const result = await drainPolls(promise, 5); + expect(result).toBe(REPO_AGENT_CANCELLED); + }); + + it("5. non-aborted completed path returns result normally", async () => { + const fetchMock = buildFetch([ + mockProgress("running"), + mockProgress("completed", { content: "normal answer" }), + ]); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS); + const result = await drainPolls(promise, 3); + expect(result).toEqual({ content: "normal answer" }); + }); + + it("6. non-aborted failed path throws an error", async () => { + const fetchMock = buildFetch([ + mockProgress("failed", undefined, "execution error"), + ]); + vi.stubGlobal("fetch", fetchMock); + + // Attach the rejection handler before advancing timers so there's no window + // where the promise is rejected but not yet observed (unhandled rejection). + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS); + const settled = promise.catch((e) => ({ threw: true, message: (e as Error).message })); + await vi.advanceTimersByTimeAsync(5100); + const result = await settled; + expect(result).toMatchObject({ threw: true, message: expect.stringContaining("execution error") }); + }); + + it("7. non-aborted timeout path throws after max attempts", async () => { + // All polls return "running" — never terminates normally. + const fetchMock = buildFetch(Array(130).fill(mockProgress("running"))); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS); + // Attach handler before advancing to avoid unhandled rejection warnings. + const settled = promise.catch((e) => ({ threw: true, message: (e as Error).message })); + // Advance past 120 × 5s = 600 s + await vi.advanceTimersByTimeAsync(120 * 5100); + const result = await settled; + expect(result).toMatchObject({ threw: true, message: expect.stringContaining("timed out") }); + }); + + it("8. onRequestId hook is called with the request_id right after initiate", async () => { + const fetchMock = buildFetch([mockProgress("completed", { content: "ok" })]); + vi.stubGlobal("fetch", fetchMock); + + const onRequestId = vi.fn().mockResolvedValue(undefined); + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { onRequestId }); + await drainPolls(promise, 2); + + expect(onRequestId).toHaveBeenCalledOnce(); + expect(onRequestId).toHaveBeenCalledWith("req-1"); + }); + + it("9. isAbortRequested is called each poll cycle", async () => { + const fetchMock = buildFetch([ + mockProgress("running"), + mockProgress("running"), + mockProgress("completed", { content: "done" }), + ]); + vi.stubGlobal("fetch", fetchMock); + + const isAbortRequested = vi.fn().mockResolvedValue(false); + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { isAbortRequested }); + await drainPolls(promise, 4); + + // Should have been called at least once per poll cycle (3 progress polls = 3 calls). + expect(isAbortRequested.mock.calls.length).toBeGreaterThanOrEqual(3); + }); + + it("10. repoAgent module imports no db/Prisma", async () => { + // Dynamic import of the module source — check it doesn't pull in Prisma. + // We verify by checking what's imported at the top of askTools.ts. + // The simplest runtime check: no `db` property on the module. + const mod = await import("@/lib/ai/askTools"); + expect((mod as Record).db).toBeUndefined(); + expect((mod as Record).prisma).toBeUndefined(); + }); + + it("completed with empty result after abort returns cancelled marker", async () => { + const fetchMock = buildFetch([ + mockProgress("completed", {}), // empty result + ]); + vi.stubGlobal("fetch", fetchMock); + + const promise = repoAgent(SWARM_URL, SWARM_KEY, BASE_PARAMS, undefined, { + isAbortRequested: async () => true, + }); + const result = await drainPolls(promise, 2); + expect(result).toBe(REPO_AGENT_CANCELLED); + }); +}); diff --git a/src/__tests__/unit/services/canvas-active-runs.test.ts b/src/__tests__/unit/services/canvas-active-runs.test.ts new file mode 100644 index 0000000000..e9a7a99b74 --- /dev/null +++ b/src/__tests__/unit/services/canvas-active-runs.test.ts @@ -0,0 +1,415 @@ +/** + * Unit tests for `src/services/canvas-active-runs.ts`. + * + * Coverage: + * 1. setActiveRun writes the entry and concurrent calls don't clobber each other + * 2. clearActiveRun removes only its own key; sibling keys survive + * 3. requestAbortForRuns sets abortRequested on targeted keys only + * 4. getActiveRuns filters out stale entries (past TTL) + * 5. isRunAbortRequested returns true/false correctly + * 6. setPendingAbortIntent + consumePendingAbortIntent (happy path) + * 7. consumePendingAbortIntent is idempotent (second call → null) + * 8. consumePendingAbortIntent returns null for a different turnId + * 9. Stale entries are pruned on next write + * 10. hasActiveRun returns false when only stale entries remain + */ + +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// ── Prisma mock ────────────────────────────────────────────────────────────── +vi.mock("@/lib/db", () => ({ + db: { + $transaction: vi.fn(), + sharedConversation: { + findFirst: vi.fn(), + update: vi.fn(), + }, + }, +})); + +import { db } from "@/lib/db"; +import { + setActiveRun, + clearActiveRun, + requestAbortForRuns, + getActiveRuns, + isRunAbortRequested, + setPendingAbortIntent, + consumePendingAbortIntent, + hasActiveRun, + ACTIVE_RUN_TTL_MS, + ABORT_INTENT_TTL_MS, +} from "@/services/canvas-active-runs"; + +// ── Type aliases for mock cast ──────────────────────────────────────────────── +const txn = db.$transaction as ReturnType; +const findFirst = db.sharedConversation.findFirst as ReturnType; +const update = db.sharedConversation.update as ReturnType; + +const CONV_ID = "conv-test-1"; + +// ── Transaction simulator ──────────────────────────────────────────────────── +/** + * Simulates the row-locking transaction: the callback receives a tx object + * whose `$queryRaw` returns `[{ active_runs: currentValue }]`. + * Updates are captured in `updates`. + */ +function setupTxn(currentActiveRuns: unknown) { + const updates: unknown[] = []; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: currentActiveRuns }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + updates.push(data.activeRuns); + // Update the "DB" so the next read sees the latest value. + currentActiveRuns = data.activeRuns; + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + return { updates, getCurrentValue: () => currentActiveRuns }; +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +// ── Tests ───────────────────────────────────────────────────────────────────── + +describe("setActiveRun", () => { + it("writes an entry under the requestId key", async () => { + const { updates } = setupTxn(null); + + await setActiveRun(CONV_ID, { + requestId: "req-1", + workspaceId: "ws-1", + startedAt: new Date().toISOString(), + }); + + expect(updates).toHaveLength(1); + const written = updates[0] as { runs: Record }; + expect(written.runs["req-1"]).toBeDefined(); + expect((written.runs["req-1"] as { requestId: string }).requestId).toBe("req-1"); + }); + + it("two concurrent setActiveRun calls each persist their key without clobbering", async () => { + // Simulate concurrent calls by running both through the same txn helper, + // but we need to serialize them since the txn mock is sequential. + const now = new Date().toISOString(); + let state: unknown = null; + + // Make txn actually accumulate state. + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const currentState = state; + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: currentState }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + state = data.activeRuns; + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + + await setActiveRun(CONV_ID, { requestId: "req-A", workspaceId: "ws-1", startedAt: now }); + await setActiveRun(CONV_ID, { requestId: "req-B", workspaceId: "ws-2", startedAt: now }); + + const col = state as { runs: Record }; + expect(col.runs["req-A"]).toBeDefined(); + expect(col.runs["req-B"]).toBeDefined(); + }); +}); + +describe("clearActiveRun", () => { + it("removes only its own key; sibling key survives", async () => { + const now = new Date().toISOString(); + const initial = { + runs: { + "req-1": { requestId: "req-1", workspaceId: "ws-1", startedAt: now }, + "req-2": { requestId: "req-2", workspaceId: "ws-2", startedAt: now }, + }, + abortIntents: {}, + }; + const { updates } = setupTxn(initial); + + await clearActiveRun(CONV_ID, "req-1"); + + const written = updates[0] as { runs: Record }; + expect(written.runs["req-1"]).toBeUndefined(); + expect(written.runs["req-2"]).toBeDefined(); + }); +}); + +describe("requestAbortForRuns", () => { + it("sets abortRequested on targeted keys only", async () => { + const now = new Date().toISOString(); + const initial = { + runs: { + "req-1": { requestId: "req-1", workspaceId: "ws-1", startedAt: now }, + "req-2": { requestId: "req-2", workspaceId: "ws-2", startedAt: now }, + }, + abortIntents: {}, + }; + const { updates } = setupTxn(initial); + + await requestAbortForRuns(CONV_ID, ["req-1"]); + + const written = updates[0] as { runs: Record }; + expect(written.runs["req-1"]?.abortRequested).toBe(true); + expect(written.runs["req-2"]?.abortRequested).toBeUndefined(); + }); + + it("silently skips missing keys (already cleared)", async () => { + const { updates } = setupTxn({ runs: {}, abortIntents: {} }); + await expect(requestAbortForRuns(CONV_ID, ["req-ghost"])).resolves.not.toThrow(); + expect(updates).toHaveLength(1); + }); +}); + +describe("getActiveRuns", () => { + it("returns non-stale entries", async () => { + const now = new Date().toISOString(); + findFirst.mockResolvedValue({ + activeRuns: { + runs: { "req-1": { requestId: "req-1", workspaceId: "ws-1", startedAt: now } }, + abortIntents: {}, + }, + }); + + const runs = await getActiveRuns(CONV_ID); + expect(runs).toHaveLength(1); + expect(runs[0].requestId).toBe("req-1"); + }); + + it("filters out stale entries (past TTL)", async () => { + const staleTime = new Date(Date.now() - ACTIVE_RUN_TTL_MS - 1000).toISOString(); + findFirst.mockResolvedValue({ + activeRuns: { + runs: { + "req-stale": { requestId: "req-stale", workspaceId: "ws-1", startedAt: staleTime }, + }, + abortIntents: {}, + }, + }); + + const runs = await getActiveRuns(CONV_ID); + expect(runs).toHaveLength(0); + }); + + it("returns empty array when no activeRuns column exists", async () => { + findFirst.mockResolvedValue({ activeRuns: null }); + const runs = await getActiveRuns(CONV_ID); + expect(runs).toEqual([]); + }); +}); + +describe("isRunAbortRequested", () => { + it("returns true when the flag is set", async () => { + const now = new Date().toISOString(); + findFirst.mockResolvedValue({ + activeRuns: { + runs: { + "req-1": { + requestId: "req-1", + workspaceId: "ws-1", + startedAt: now, + abortRequested: true, + }, + }, + abortIntents: {}, + }, + }); + + expect(await isRunAbortRequested(CONV_ID, "req-1")).toBe(true); + }); + + it("returns false when the flag is not set", async () => { + const now = new Date().toISOString(); + findFirst.mockResolvedValue({ + activeRuns: { + runs: { "req-1": { requestId: "req-1", workspaceId: "ws-1", startedAt: now } }, + abortIntents: {}, + }, + }); + + expect(await isRunAbortRequested(CONV_ID, "req-1")).toBe(false); + }); + + it("returns false when the entry is gone (already cleared)", async () => { + findFirst.mockResolvedValue({ activeRuns: { runs: {}, abortIntents: {} } }); + expect(await isRunAbortRequested(CONV_ID, "req-gone")).toBe(false); + }); +}); + +describe("setPendingAbortIntent / consumePendingAbortIntent", () => { + it("stores and then consumes the intent", async () => { + let state: unknown = { runs: {}, abortIntents: {} }; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: state }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + state = data.activeRuns; + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + + await setPendingAbortIntent(CONV_ID, "turn-1"); + const intent = await consumePendingAbortIntent(CONV_ID, "turn-1"); + + expect(intent).not.toBeNull(); + expect(intent?.turnId).toBe("turn-1"); + }); + + it("consumePendingAbortIntent is idempotent — second call returns null", async () => { + let state: unknown = { runs: {}, abortIntents: {} }; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: state }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + state = data.activeRuns; + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + + await setPendingAbortIntent(CONV_ID, "turn-2"); + await consumePendingAbortIntent(CONV_ID, "turn-2"); // first consume + const second = await consumePendingAbortIntent(CONV_ID, "turn-2"); // second consume + expect(second).toBeNull(); + }); + + it("does not consume an intent for a different turnId", async () => { + let state: unknown = { runs: {}, abortIntents: {} }; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: state }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + state = data.activeRuns; + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + + await setPendingAbortIntent(CONV_ID, "turn-A"); + const result = await consumePendingAbortIntent(CONV_ID, "turn-B"); // different turn + expect(result).toBeNull(); + + // The intent for turn-A should still be there. + const remaining = await consumePendingAbortIntent(CONV_ID, "turn-A"); + expect(remaining).not.toBeNull(); + expect(remaining?.turnId).toBe("turn-A"); + }); + + it("returns null for an expired intent", async () => { + const expiredTime = new Date(Date.now() - ABORT_INTENT_TTL_MS - 1000).toISOString(); + const state = { + runs: {}, + abortIntents: { + "turn-old": { + turnId: "turn-old", + requestedAt: expiredTime, + expiresAt: expiredTime, // already expired + }, + }, + }; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: state }]), + sharedConversation: { update: vi.fn().mockResolvedValue({}) }, + }; + await cb(tx); + }); + + const result = await consumePendingAbortIntent(CONV_ID, "turn-old"); + expect(result).toBeNull(); + }); +}); + +describe("hasActiveRun", () => { + it("returns true when a fresh run exists", async () => { + findFirst.mockResolvedValue({ + activeRuns: { + runs: { + "req-1": { + requestId: "req-1", + workspaceId: "ws-1", + startedAt: new Date().toISOString(), + }, + }, + abortIntents: {}, + }, + }); + + expect(await hasActiveRun(CONV_ID)).toBe(true); + }); + + it("returns false when only stale entries remain", async () => { + const staleTime = new Date(Date.now() - ACTIVE_RUN_TTL_MS - 1000).toISOString(); + findFirst.mockResolvedValue({ + activeRuns: { + runs: { + "req-stale": { requestId: "req-stale", workspaceId: "ws-1", startedAt: staleTime }, + }, + abortIntents: {}, + }, + }); + + expect(await hasActiveRun(CONV_ID)).toBe(false); + }); + + it("returns false when there are no runs at all", async () => { + findFirst.mockResolvedValue({ activeRuns: null }); + expect(await hasActiveRun(CONV_ID)).toBe(false); + }); +}); + +describe("stale entry pruning on write", () => { + it("stale entries are removed during the next setActiveRun write", async () => { + const staleTime = new Date(Date.now() - ACTIVE_RUN_TTL_MS - 5000).toISOString(); + const initial = { + runs: { + "req-stale": { requestId: "req-stale", workspaceId: "ws-1", startedAt: staleTime }, + }, + abortIntents: {}, + }; + + const updates: unknown[] = []; + txn.mockImplementation(async (cb: (tx: unknown) => Promise) => { + const tx = { + $queryRaw: vi.fn().mockResolvedValue([{ active_runs: initial }]), + sharedConversation: { + update: vi.fn().mockImplementation(({ data }: { data: { activeRuns: unknown } }) => { + updates.push(data.activeRuns); + return Promise.resolve({}); + }), + }, + }; + await cb(tx); + }); + + await setActiveRun(CONV_ID, { + requestId: "req-fresh", + workspaceId: "ws-2", + startedAt: new Date().toISOString(), + }); + + const written = updates[0] as { runs: Record }; + expect(written.runs["req-stale"]).toBeUndefined(); + expect(written.runs["req-fresh"]).toBeDefined(); + }); +}); diff --git a/src/app/api/learnings/diagrams/create/route.ts b/src/app/api/learnings/diagrams/create/route.ts index 2290538812..1d0647014c 100644 --- a/src/app/api/learnings/diagrams/create/route.ts +++ b/src/app/api/learnings/diagrams/create/route.ts @@ -2,7 +2,7 @@ import { NextRequest, NextResponse } from "next/server"; import { db } from "@/lib/db"; import { getSwarmConfig } from "../../utils"; import { getAllRepositories, joinRepoUrls } from "@/lib/helpers/repository"; -import { repoAgent } from "@/lib/ai/askTools"; +import { repoAgent, REPO_AGENT_CANCELLED } from "@/lib/ai/askTools"; import { getGithubUsernameAndPAT } from "@/lib/auth/nextauth"; import { resolveWorkspaceAccess, requireMemberAccess } from "@/lib/auth/workspace-access"; import { extractMermaidBody } from "@/lib/diagrams/mermaid-parser"; @@ -86,6 +86,10 @@ export async function POST(request: NextRequest) { bifrost, ); + if (agentResult === REPO_AGENT_CANCELLED) { + return NextResponse.json({ error: "Agent run was cancelled" }, { status: 499 }); + } + const responseContent = agentResult?.content ?? JSON.stringify(agentResult); const extractedBody = extractMermaidBody(responseContent); diff --git a/src/app/api/learnings/diagrams/edit/route.ts b/src/app/api/learnings/diagrams/edit/route.ts index f26607e00e..35bb51cd77 100644 --- a/src/app/api/learnings/diagrams/edit/route.ts +++ b/src/app/api/learnings/diagrams/edit/route.ts @@ -2,7 +2,7 @@ import { NextRequest, NextResponse } from "next/server"; import { db } from "@/lib/db"; import { getSwarmConfig } from "../../utils"; import { getAllRepositories, joinRepoUrls } from "@/lib/helpers/repository"; -import { repoAgent } from "@/lib/ai/askTools"; +import { repoAgent, REPO_AGENT_CANCELLED } from "@/lib/ai/askTools"; import { getGithubUsernameAndPAT } from "@/lib/auth/nextauth"; import { resolveWorkspaceAccess, requireMemberAccess } from "@/lib/auth/workspace-access"; import { extractMermaidBody } from "@/lib/diagrams/mermaid-parser"; @@ -93,6 +93,10 @@ export async function POST(request: NextRequest) { bifrost, ); + if (agentResult === REPO_AGENT_CANCELLED) { + return NextResponse.json({ error: "Agent run was cancelled" }, { status: 499 }); + } + const responseContent = agentResult?.content ?? JSON.stringify(agentResult); const extractedBody = extractMermaidBody(responseContent); diff --git a/src/lib/ai/askTools.ts b/src/lib/ai/askTools.ts index 0cca7468bd..60300307c9 100644 --- a/src/lib/ai/askTools.ts +++ b/src/lib/ai/askTools.ts @@ -61,6 +61,20 @@ export interface SubAgent { timeoutSeconds?: number; } +/** + * Sentinel value returned (never thrown) when the user cancels a run via the + * Stop control. Callers should surface this as a normal, non-error completion. + */ +export const REPO_AGENT_CANCELLED = "REPO_AGENT_CANCELLED" as const; +export type RepoAgentCancelledMarker = typeof REPO_AGENT_CANCELLED; + +/** + * Number of poll cycles to wait after `isAbortRequested` returns `true` + * before giving up and returning the cancelled marker locally (grace window). + * At the default 5-second poll interval this is ~10-15 seconds. + */ +const ABORT_GRACE_POLL_CYCLES = 3; + export async function repoAgent( swarmUrl: string, swarmApiKey: string, @@ -124,7 +138,20 @@ export async function repoAgent( baseUrl: string; headers?: Record; }, -): Promise> { + hooks?: { + /** + * Fired immediately after the initiate POST returns `request_id`. + * DB-free: the caller supplies this; `repoAgent` itself has no Prisma dep. + */ + onRequestId?: (id: string) => Promise; + /** + * Called each poll cycle. When it returns `true`, the run is treated as + * aborted by the user (see grace-window logic below). + * DB-free: supplied by the caller. + */ + isAbortRequested?: () => Promise; + }, +): Promise | RepoAgentCancelledMarker> { const body: Record = { ...params }; if (bifrost) { body.apiKey = bifrost.apiKey; @@ -146,7 +173,7 @@ export async function repoAgent( if (!initiateResponse.ok) { const errorText = await initiateResponse.text(); - console.error(`Repo agent initiation error: ${initiateResponse.status} - ${errorText}`); + console.error(`[repoAgent] Repo agent initiation error: ${initiateResponse.status} - ${errorText}`); throw new Error("Failed to initiate repo agent"); } @@ -157,12 +184,27 @@ export async function repoAgent( throw new Error("No request_id returned from repo agent"); } + // Notify caller so it can persist the run entry and check for pending-abort intent. + if (hooks?.onRequestId) { + await hooks.onRequestId(requestId); + } + const maxAttempts = 120; const pollInterval = 5000; + let abortDetectedCycles = 0; // how many cycles since we first saw abortRequested + for (let attempt = 0; attempt < maxAttempts; attempt++) { await new Promise((resolve) => setTimeout(resolve, pollInterval)); + // ── Check abort flag ────────────────────────────────────────── + const abortRequested = + hooks?.isAbortRequested ? await hooks.isAbortRequested() : false; + + if (abortRequested) { + abortDetectedCycles++; + } + const progressUrl = `${swarmUrl}/progress?request_id=${encodeURIComponent(requestId)}`; const progressResponse = await fetch(progressUrl, { method: "GET", @@ -173,17 +215,50 @@ export async function repoAgent( }); if (!progressResponse.ok) { - console.error(`Progress check failed: ${progressResponse.status}`); + console.error(`[repoAgent] Progress check failed: ${progressResponse.status}`); + // Still check grace-window even on poll failure. + if (abortRequested && abortDetectedCycles >= ABORT_GRACE_POLL_CYCLES) { + console.warn(`[repoAgent] Grace window elapsed after abort (requestId=${requestId}); returning cancelled marker`); + return REPO_AGENT_CANCELLED; + } continue; } const progressData = await progressResponse.json(); + // ── Completed ───────────────────────────────────────────────── if (progressData.status === "completed") { - return progressData.result || {}; - } else if (progressData.status === "failed") { + // Requirement 5: a run that truly completes — even microseconds after + // Stop — returns its real result, not the cancelled marker. + const result = progressData.result; + if (result && Object.keys(result).length > 0) { + return result; + } + // completed but no usable result — if aborting, treat as cancelled. + if (abortRequested) { + console.warn(`[repoAgent] Run completed without result after abort (requestId=${requestId}); returning cancelled marker`); + return REPO_AGENT_CANCELLED; + } + return result || {}; + } + + // ── Failed / aborted ───────────────────────────────────────── + if (progressData.status === "failed" || progressData.status === "aborted") { + if (abortRequested) { + // cancelled ≠ error: non-throwing marker. + console.warn(`[repoAgent] Run ${progressData.status} after abort request (requestId=${requestId}); returning cancelled marker`); + return REPO_AGENT_CANCELLED; + } throw new Error(progressData.error || "Repo agent execution failed"); } + + // ── Grace-window exit ───────────────────────────────────────── + // Swarm is neither completing nor acknowledging the abort — exit locally + // rather than waiting up to 120×5s. + if (abortRequested && abortDetectedCycles >= ABORT_GRACE_POLL_CYCLES) { + console.warn(`[repoAgent] Grace window elapsed after abort (requestId=${requestId}); returning cancelled marker`); + return REPO_AGENT_CANCELLED; + } } throw new Error("Repo agent execution timed out. Please try again."); @@ -343,6 +418,7 @@ export function askTools(swarmUrl: string, swarmApiKey: string, repoUrls: string }, bifrost, ); + if (rr === REPO_AGENT_CANCELLED) return "Agent run was cancelled"; return rr.content; } catch (e) { console.error("Error executing repo agent:", e); diff --git a/src/lib/ai/askToolsMulti.ts b/src/lib/ai/askToolsMulti.ts index 81b65a8b41..a69b4880e5 100644 --- a/src/lib/ai/askToolsMulti.ts +++ b/src/lib/ai/askToolsMulti.ts @@ -3,7 +3,7 @@ import { z } from "zod"; import { createMCPClient } from "@ai-sdk/mcp"; import { withMcpTimeout, isMcpTimeout } from './mcpTimeout'; import { WorkspaceConfig } from "./types"; -import { listConcepts, repoAgent } from "./askTools"; +import { listConcepts, repoAgent, REPO_AGENT_CANCELLED } from "./askTools"; import { buildCourtlistenerTools } from "@/lib/ai/courtlistenerTools"; import { LEGAL_SLUGS } from "@/lib/eval-capture-slugs"; // Deep import — see comment in services/task-workflow.ts. @@ -301,6 +301,7 @@ export function askToolsMulti( }, bifrost, ); + if (rr === REPO_AGENT_CANCELLED) return "Agent run was cancelled"; return rr.content; } catch (e) { console.error(`Error executing repo agent for ${ws.slug}:`, e); diff --git a/src/lib/ai/workflowExplorerTools.ts b/src/lib/ai/workflowExplorerTools.ts index 39c874df40..88df85bf58 100644 --- a/src/lib/ai/workflowExplorerTools.ts +++ b/src/lib/ai/workflowExplorerTools.ts @@ -20,7 +20,7 @@ import { tool, type ToolSet } from "ai"; import { z } from "zod"; import { db } from "@/lib/db"; import { EncryptionService } from "@/lib/encryption"; -import { repoAgent } from "./askTools"; +import { repoAgent, REPO_AGENT_CANCELLED } from "./askTools"; /** The workspace whose swarm hosts the Jarvis workflow-library graph. */ const WORKFLOW_LIBRARY_WORKSPACE_SLUG = "stakwork"; @@ -120,6 +120,7 @@ export function buildWorkflowExplorerTools(): ToolSet { mode: "workflow", stakworkApiKey, }); + if (rr === REPO_AGENT_CANCELLED) return "Agent run was cancelled"; return rr.content; } catch (e) { console.error("Error executing workflow explorer agent:", e); diff --git a/src/services/canvas-active-runs.ts b/src/services/canvas-active-runs.ts new file mode 100644 index 0000000000..44d9b7b37a --- /dev/null +++ b/src/services/canvas-active-runs.ts @@ -0,0 +1,265 @@ +/** + * Atomic helpers for tracking in-flight `repo_agent` runs against a + * `SharedConversation` row. + * + * ## Shape + * The `active_runs` JSON column holds: + * ``` + * { + * runs: { [requestId]: ActiveRunEntry }, + * abortIntents: { [turnId]: AbortIntentEntry } + * } + * ``` + * + * ## Atomicity + * All mutations use a `db.$transaction` + `SELECT … FOR UPDATE` (via + * `$queryRaw`) to serialize concurrent writers. Reads are lock-free. + * + * ## Staleness reaper + * Any run entry older than ACTIVE_RUN_TTL_MS (~10 min) is considered + * dead and excluded from "is a run active" derivations. Stale entries + * are pruned opportunistically on the next write. + * + * ## Pending-abort intent + * Stored under `abortIntents[turnId]`; consumed atomically only by an + * `execute` belonging to that `turnId`, so it cannot cancel an + * unrelated later run. + */ + +import { db } from "@/lib/db"; +import { Prisma } from "@prisma/client"; + +// ─── Constants ──────────────────────────────────────────────────────────────── + +/** Maximum age of a run entry before it is treated as a crashed-lambda phantom. */ +export const ACTIVE_RUN_TTL_MS = 10 * 60 * 1000; // 10 minutes + +/** How long a pending-abort intent remains valid (start-race window). */ +export const ABORT_INTENT_TTL_MS = 60 * 1000; // 60 seconds + +// ─── Types ──────────────────────────────────────────────────────────────────── + +export interface ActiveRunEntry { + requestId: string; + /** Re-resolved at abort time to fetch swarm creds — never stored with secrets. */ + workspaceId: string | null; + startedAt: string; // ISO-8601 + abortRequested?: boolean; +} + +export interface AbortIntentEntry { + turnId: string; + requestedAt: string; // ISO-8601 + expiresAt: string; // ISO-8601 +} + +interface ActiveRunsColumn { + runs: Record; + abortIntents: Record; +} + +// ─── Internal helpers ───────────────────────────────────────────────────────── + +function emptyColumn(): ActiveRunsColumn { + return { runs: {}, abortIntents: {} }; +} + +function parseColumn(raw: unknown): ActiveRunsColumn { + if (!raw || typeof raw !== "object") return emptyColumn(); + const col = raw as Partial; + return { + runs: col.runs ?? {}, + abortIntents: col.abortIntents ?? {}, + }; +} + +function isStale(entry: ActiveRunEntry): boolean { + const age = Date.now() - new Date(entry.startedAt).getTime(); + return age > ACTIVE_RUN_TTL_MS; +} + +function pruneStaleRuns(runs: Record): Record { + const out: Record = {}; + for (const [k, v] of Object.entries(runs)) { + if (!isStale(v)) out[k] = v; + } + return out; +} + +function pruneExpiredIntents( + intents: Record, +): Record { + const now = Date.now(); + const out: Record = {}; + for (const [k, v] of Object.entries(intents)) { + if (new Date(v.expiresAt).getTime() > now) out[k] = v; + } + return out; +} + +/** + * Execute `mutate` inside a row-locking transaction on `SharedConversation`. + * The `FOR UPDATE` lock serializes concurrent writers on this row. + */ +async function withLockedRow( + conversationId: string, + mutate: (current: ActiveRunsColumn) => ActiveRunsColumn, +): Promise { + await db.$transaction(async (tx) => { + // Lock the row for the duration of this transaction. + const rows = await tx.$queryRaw<{ active_runs: unknown }[]>` + SELECT active_runs + FROM shared_conversations + WHERE id = ${conversationId} + FOR UPDATE + `; + + const current = parseColumn(rows[0]?.active_runs ?? null); + const next = mutate(current); + + await tx.sharedConversation.update({ + where: { id: conversationId }, + data: { activeRuns: next as unknown as Prisma.InputJsonValue }, + }); + }); +} + +// ─── Public API ─────────────────────────────────────────────────────────────── + +/** + * Record a newly-started run. Also opportunistically prunes stale entries. + */ +export async function setActiveRun( + conversationId: string, + entry: Omit, +): Promise { + await withLockedRow(conversationId, (col) => { + const runs = pruneStaleRuns(col.runs); + runs[entry.requestId] = { ...entry, startedAt: entry.startedAt }; + return { ...col, runs }; + }); +} + +/** + * Remove a single run entry by `requestId` (called in `finally`). + * Only removes the named key — sibling keys are untouched. + */ +export async function clearActiveRun( + conversationId: string, + requestId: string, +): Promise { + await withLockedRow(conversationId, (col) => { + const runs = pruneStaleRuns(col.runs); + delete runs[requestId]; + return { ...col, runs }; + }); +} + +/** + * Flag one or more runs as `abortRequested`. Missing keys are silently skipped + * (race: the run may have already cleared in `finally`). + */ +export async function requestAbortForRuns( + conversationId: string, + requestIds: string[], +): Promise { + await withLockedRow(conversationId, (col) => { + const runs = pruneStaleRuns(col.runs); + for (const id of requestIds) { + if (runs[id]) { + runs[id] = { ...runs[id], abortRequested: true }; + } + } + return { ...col, runs }; + }); +} + +/** + * Lock-free read of all currently-active run entries (stale entries filtered out). + */ +export async function getActiveRuns( + conversationId: string, +): Promise { + const row = await db.sharedConversation.findFirst({ + where: { id: conversationId }, + select: { activeRuns: true }, + }); + const col = parseColumn(row?.activeRuns ?? null); + return Object.values(col.runs).filter((r) => !isStale(r)); +} + +/** + * Lock-free read of the `abortRequested` flag for a single run. + * Returns `false` if the entry is gone (already cleared). + */ +export async function isRunAbortRequested( + conversationId: string, + requestId: string, +): Promise { + const row = await db.sharedConversation.findFirst({ + where: { id: conversationId }, + select: { activeRuns: true }, + }); + const col = parseColumn(row?.activeRuns ?? null); + return col.runs[requestId]?.abortRequested === true; +} + +// ─── Pending-abort intent ──────────────────────────────────────────────────── + +/** + * Store a turn-scoped abort intent (for the start-race: Stop pressed before + * the run's `request_id` was persisted). + */ +export async function setPendingAbortIntent( + conversationId: string, + turnId: string, +): Promise { + const now = new Date(); + const intent: AbortIntentEntry = { + turnId, + requestedAt: now.toISOString(), + expiresAt: new Date(now.getTime() + ABORT_INTENT_TTL_MS).toISOString(), + }; + + await withLockedRow(conversationId, (col) => { + const abortIntents = pruneExpiredIntents(col.abortIntents); + abortIntents[turnId] = intent; + return { ...col, abortIntents }; + }); +} + +/** + * Atomically consume the pending-abort intent for `turnId` (if it exists and + * has not expired). Returns the intent if it was present; `null` otherwise. + * + * "Consume" means remove — calling this twice for the same `turnId` returns + * `null` on the second call, preventing double-cancellation. + */ +export async function consumePendingAbortIntent( + conversationId: string, + turnId: string, +): Promise { + let found: AbortIntentEntry | null = null; + + await withLockedRow(conversationId, (col) => { + const abortIntents = pruneExpiredIntents(col.abortIntents); + const intent = abortIntents[turnId]; + if (intent) { + found = intent; + delete abortIntents[turnId]; + } + return { ...col, abortIntents }; + }); + + return found; +} + +/** + * Derive a single boolean: "does this conversation have at least one + * non-stale, non-aborted active run?" Used by the Stop-button visibility + * logic and the Pusher broadcast. + */ +export async function hasActiveRun(conversationId: string): Promise { + const runs = await getActiveRuns(conversationId); + return runs.length > 0; +}