diff --git a/server/src/services/remote-workspace-migration.test.ts b/server/src/services/remote-workspace-migration.test.ts index 32dd4fa4d6b..c31fad595eb 100644 --- a/server/src/services/remote-workspace-migration.test.ts +++ b/server/src/services/remote-workspace-migration.test.ts @@ -8,7 +8,7 @@ import { createServer, type Server } from "node:http"; import { pathToFileURL } from "node:url"; import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { and, eq, sql } from "drizzle-orm"; -import { activityLog, agents, environmentLeases, heartbeatRuns, nativeRunFinalizations, nativeRunResults, completionContracts, applyPendingMigrations, closeRegisteredClients, companies, createDb, executionWorkspaceSnapshots, +import { activityLog, agents, environments, environmentLeases, heartbeatRuns, nativeRunFinalizations, nativeRunResults, completionContracts, applyPendingMigrations, closeRegisteredClients, companies, createDb, executionWorkspaceSnapshots, executionWorkspaceRepositories, executionWorkspaceRuntimeLeases, executionWorkspaces, issueRecoveryActions, issues, projects, projectWorkspaces, workspaceOperations, workspaceRuntimeServices } from "@paperclipai/db"; import { readRemoteWorkspaceAuthority } from "@paperclipai/shared"; import { createLocalDiskStorageProvider } from "../storage/local-disk-provider.js"; @@ -37,6 +37,9 @@ import { isolatedRemoteSourcePath } from "./isolated-remote-workspace-source.js" import { verifiedGitCacheSource } from "./git-workspace-restoration.js"; import { withIndependentWorkspaceSourceLocks, withWorkspaceSourceLocks } from "./workspace-source-lock.js"; import { inspectOnlineWorkspaceSource } from "./online-workspace-source.js"; +import { tryLockNativeWorkspaceFinalization } from "./native-runtime/native-workspace-finalization-ownership.js"; +import { readMigratedLegacyWorkspaceSource } from "./legacy-workspace-recovery-migration.js"; +import { nativeSha256 } from "./native-runtime/canonical.js"; const execute = promisify(execFile); const git = async (cwd: string, ...args: string[]) => (await execute("git", ["-C", cwd, ...args], { env: { ...process.env, @@ -61,6 +64,7 @@ describe.skipIf(!support.supported)("remote workspace migration", () => { await Promise.all(servers.splice(0).map(server => new Promise((resolve, reject) => server.close(error => error ? reject(error) : resolve())))); vi.unstubAllEnvs(); await db.delete(activityLog); await db.delete(workspaceOperations); await db.delete(executionWorkspaceSnapshots); await db.delete(workspaceRuntimeServices); await db.delete(executionWorkspaceRuntimeLeases); + await db.delete(environmentLeases); await db.delete(environments); await db.delete(nativeRunFinalizations); await db.delete(nativeRunResults); await db.delete(heartbeatRuns); await db.delete(completionContracts); await db.delete(issues); await db.delete(agents); await db.delete(executionWorkspaces); await db.delete(projectWorkspaces); await db.delete(projects); await db.delete(companies); await fs.rm(root, { recursive: true, force: true }); @@ -378,12 +382,12 @@ describe.skipIf(!support.supported)("remote workspace migration", () => { await restoreRemoteWorkspaceArchive({ archivePath, targetDir, expectedManifestDigest: reference!.manifestDigest }); expect(await fs.readFile(path.join(targetDir, ".paperclip/worktrees/child/child.txt"), "utf8")).toBe("saved child edits\n"); }); - async function failedCopybackFixture() { - const f = await seed(), runId = randomUUID(); + async function failedCopybackFixture(existing?: Awaited>) { + const f = existing ?? await seed(), runId = randomUUID(); const [agent] = await db.insert(agents).values({ companyId: f.companyId, name: "Historical agent", adapterType: "paperclip_runner" }).returning(); const reference = { schema: "paperclip.native-workspace-sync/v3", state: "prepared", workspaceId: f.workspaceId, leaseId: randomUUID(), providerLeaseId: randomUUID(), remoteCwd: "/workspaces/task", - descriptorSha256: "12".repeat(32), baselineSha256: "34".repeat(32), finalHostSha256: null, resourceDisposition: "stop_and_retain" }; + descriptorSha256: "12".repeat(32), baselineSha256: "34".repeat(32), finalHostSha256: null, resourceDisposition: "stop_and_retain" } as const; await db.insert(heartbeatRuns).values({ id: runId, companyId: f.companyId, agentId: agent!.id, nativeIssueId: f.issueId, runtimeMode: "native", status: "failed", finishedAt: new Date(), runnerProfileJson: { nativeWorkspaceSync: reference } }); await db.insert(nativeRunFinalizations).values({ companyId: f.companyId, issueId: f.issueId, runId, phase: "terminal_failure" }); @@ -633,6 +637,122 @@ describe.skipIf(!support.supported)("remote workspace migration", () => { branchName, strategyType: "git_worktree", providerType: "git_worktree", metadata: { createdByRuntime: true, gitBranchOwnershipVersion: 1 } }).where(eq(executionWorkspaces.id, f.workspaceId)); return { ...f, cwd, base, upstream, commit, projectId, projectWorkspaceId, repoUrl, branchName }; } + async function archivedPreparedFixture() { + const base = await gitArchiveFixture(), f = { ...base, ...await failedCopybackFixture(base) }; + await fs.writeFile(path.join(f.cwd, "saved.txt"), "archived staged edit\n"); await git(f.cwd, "add", "saved.txt"); + await fs.writeFile(path.join(f.cwd, "saved.txt"), "archived unstaged edit\n"); + await fs.writeFile(path.join(f.cwd, "untracked.txt"), "archived private file\n"); + const captured = (await captureRestorePoint({ cwd: f.cwd, upstreamCommits: [f.upstream] }))!; + const id = randomUUID(), generation = 1, sealed = await sealRestorePoint(captured.payload, + { companyId: f.companyId, workspaceId: f.workspaceId, restorePointId: id, generation }); + const objectKey = `${f.companyId}/git-restore-points/${f.workspaceId}/${generation}/${id}.pcrp`; + await storage.putFile({ companyId: f.companyId, namespace: "git-restore-points", objectKey, + originalFilename: "saved.pcrp", contentType: "application/octet-stream", body: sealed.ciphertext }); + const restorePoint = { schema: "paperclip.git-restore-point.v1", id, generation, objectKey, byteSize: sealed.ciphertext.length, + sha256: sealed.sha256, key: sealed.key, head: captured.manifest.head, index: captured.manifest.index, + worktree: captured.manifest.worktree, ignored: captured.manifest.ignored, modes: captured.manifest.modes, headOnRemote: false }; + const contents = await captureGitWorkspaceContents(f.cwd, { companyId: f.companyId, workspaceId: f.workspaceId }, []); + await db.update(executionWorkspaces).set({ status: "archived", closedAt: new Date(), metadata: { + createdByRuntime: true, gitBranchOwnershipVersion: 1, gitCacheRestoration: { schema: "paperclip.git-cache-restoration.v2", + companyId: f.companyId, workspaceId: f.workspaceId, sourceIssueId: f.issueId, repoUrl: f.repoUrl, + commit: f.commit, branchName: f.branchName, sourceKind: "managed_project", projectId: f.projectId, + sourcePath: f.base, worktreePath: f.cwd, gitBranchOwnershipVersion: 1, verifiedAt: new Date().toISOString(), restorePoint, contents }, + } }).where(eq(executionWorkspaces.id, f.workspaceId)); + await git(f.base, "worktree", "remove", "--force", f.cwd); await git(f.base, "branch", "-D", f.branchName); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const [environment] = await db.insert(environments).values({ name: f.runId, driver: "daytona" }).returning(); + const [lease] = await db.insert(environmentLeases).values({ id: f.reference.leaseId, companyId: f.companyId, + environmentId: environment.id, heartbeatRunId: f.runId, issueId: f.issueId, executionWorkspaceId: f.workspaceId, + provider: "daytona", providerLeaseId: f.reference.providerLeaseId, status: "retained", leasePolicy: "retain_on_failure", + metadata: { remoteCwd: f.reference.remoteCwd } }).returning(); + const cause = "native_workspace_sync_out_unrecoverable"; + await db.update(nativeRunFinalizations).set({ failureCode: cause }).where(eq(nativeRunFinalizations.runId, f.runId)); + const [hold] = await db.insert(issueRecoveryActions).values({ companyId: f.companyId, sourceIssueId: f.issueId, + kind: "active_run_watchdog", cause, ownerType: "board", ownerAgentId: null, returnOwnerAgentId: run.agentId, + fingerprint: nativeSha256({ runId: f.runId, failureCode: cause }), nextAction: "Preserve retained remote files", evidence: { runId: f.runId } }).returning(); + const [selected] = await db.insert(executionWorkspaces).values({ companyId: f.companyId, sourceIssueId: f.issueId, + name: "Current task workspace", providerType: "local_fs", mode: "isolated_workspace", strategyType: "project_primary", cwd: path.join(root, "current-task") }).returning(); + await db.update(issues).set({ status: "cancelled", executionWorkspaceId: selected.id }).where(eq(issues.id, f.issueId)); + await db.insert(activityLog).values({ companyId: f.companyId, actorType: "user", actorId: "original owner", + action: "issue.updated", entityType: "issue", entityId: f.issueId, details: { retained: "original task history" } }); + const [task] = await db.select().from(issues).where(eq(issues.id, f.issueId)); + const [owner] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + const history = await db.select().from(activityLog).where(eq(activityLog.entityId, f.issueId)); + return { ...f, run, lease, hold, task, owner, history, before: await workspace(f.workspaceId) }; + } + it("preserves an unselected archived recipe and its pending native files without reopening the task", async () => { + const f = await archivedPreparedFixture(), put = storage.putFile.bind(storage); + let interrupt = true; + storage.putFile = async input => { const saved = await put(input); + if (interrupt) { interrupt = false; throw new Error("archived upload reply lost"); } return saved; }; + await expect(migrateRemoteWorkspace({ db, storage, ...f, apply: true })).rejects.toThrow("archived upload reply lost"); + const [interrupted] = await db.select().from(workspaceOperations).where(eq(workspaceOperations.phase, "remote_workspace_migration")); + const [pending] = await db.select().from(executionWorkspaceSnapshots); + expect(readRemoteWorkspaceAuthority((await workspace(f.workspaceId)).metadata)).toBeNull(); + expect(await migrateRemoteWorkspace({ db, storage, ...f, apply: true })).toMatchObject({ applied: true }); + const saved = await workspace(f.workspaceId), authority = readRemoteWorkspaceAuthority(saved.metadata)!; + expect(saved.status).toBe("archived"); expect(saved.closedAt).toEqual(f.before.closedAt); + expect(authority.snapshotId).toBe(pending.id); + expect((await db.select().from(workspaceOperations).where(eq(workspaceOperations.id, interrupted.id)))[0].status).toBe("succeeded"); + expect(saved.metadata?.gitCacheRestoration).toEqual(f.before.metadata?.gitCacheRestoration); + expect(saved.metadata?.remoteWorkspaceRecoverySourceRunIds).toEqual([f.runId]); + expect((await db.select().from(issues).where(eq(issues.id, f.issueId)))[0]).toEqual(f.task); + expect(await db.select().from(activityLog).where(eq(activityLog.entityId, f.issueId))).toEqual(f.history); + expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)))[0]).toEqual(f.owner); + expect((await db.select().from(environmentLeases).where(eq(environmentLeases.id, f.lease.id)))[0]).toEqual(f.lease); + expect((await db.select().from(issueRecoveryActions).where(eq(issueRecoveryActions.id, f.hold.id)))[0]).toEqual(f.hold); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const marker = { schema: "paperclip.legacy-workspace-recovery-migration.v1", operationId: expect.any(String), workspaceId: f.workspaceId }; + expect(run).toEqual({ ...f.run, updatedAt: expect.any(Date), runnerProfileJson: { ...f.run.runnerProfileJson, legacyWorkspaceRecoveryMigration: marker } }); + await expect(readMigratedLegacyWorkspaceSource(db, { companyId: f.companyId, runId: f.runId, + reference: f.reference })).resolves.toMatchObject({ runId: f.runId, reference: f.reference }); + await expect(fs.access(f.cwd)).rejects.toMatchObject({ code: "ENOENT" }); + const ref = (await getWorkspaceSnapshot(db, { ...f, snapshotId: authority.snapshotId! }))!; + const archivePath = path.join(root, "archived-native.pcws"), targetDir = path.join(root, "restored-native"); + await loadWorkspaceSnapshotArchive({ storage, reference: ref, targetPath: archivePath }); + await restoreRemoteWorkspaceArchive({ archivePath, targetDir, expectedManifestDigest: ref.manifestDigest }); + expect(await git(targetDir, "show", ":saved.txt")).toBe("archived staged edit"); + expect(await fs.readFile(path.join(targetDir, "saved.txt"), "utf8")).toBe("archived unstaged edit\n"); + expect(await fs.readFile(path.join(targetDir, "untracked.txt"), "utf8")).toBe("archived private file\n"); + await migrateRemoteWorkspace({ db, storage, ...f, apply: true }); + expect((await db.select().from(workspaceOperations).where(eq(workspaceOperations.phase, "remote_workspace_migration")))).toHaveLength(1); + }); + it.each(["recipe", "reference", "provider"])("refuses an archived prepared source changed during upload (%s)", async change => { + const f = await archivedPreparedFixture(), put = storage.putFile.bind(storage); + let uploaded = false; + storage.putFile = async input => { + const saved = await put(input); uploaded = true; + if (change === "recipe") await db.update(executionWorkspaces).set({ metadata: { ...f.before.metadata, + remoteWorkspaceMigration: (await workspace(f.workspaceId)).metadata?.remoteWorkspaceMigration, + gitCacheRestoration: { ...f.before.metadata!.gitCacheRestoration as object, commit: f.upstream } } }).where(eq(executionWorkspaces.id, f.workspaceId)); + else { + const reference = { ...f.reference, ...(change === "reference" ? { baselineSha256: "56".repeat(32) } : { providerLeaseId: randomUUID() }) }; + if (change === "provider") await db.update(environmentLeases).set({ providerLeaseId: reference.providerLeaseId }).where(eq(environmentLeases.id, f.lease.id)); + await db.update(heartbeatRuns).set({ runnerProfileJson: { ...f.run.runnerProfileJson, nativeWorkspaceSync: reference } }).where(eq(heartbeatRuns.id, f.runId)); + } + return saved; + }; + await expect(migrateRemoteWorkspace({ db, storage, ...f, apply: true })).rejects.toThrow(); + expect(uploaded).toBe(true); + expect(readRemoteWorkspaceAuthority((await workspace(f.workspaceId)).metadata)).toBeNull(); + expect((await db.select().from(issues).where(eq(issues.id, f.issueId)))[0]).toEqual(f.task); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)))[0].runnerProfileJson).not.toHaveProperty("legacyWorkspaceRecoveryMigration"); + await expect(fs.access(f.cwd)).rejects.toMatchObject({ code: "ENOENT" }); + }); + it("refuses a competing native F owner before capturing an archived prepared source", async () => { + const f = await archivedPreparedFixture(); + let release!: () => void, entered!: () => void; + const held = new Promise(done => { release = done; }), ready = new Promise(done => { entered = done; }); + const owner = db.transaction(async tx => { expect(await tryLockNativeWorkspaceFinalization(tx as unknown as typeof db, + { companyId: f.companyId, runId: f.runId })).toBe(true); entered(); await held; }); + await ready; + try { + await expect(migrateRemoteWorkspace({ db, storage, ...f, apply: true })).rejects.toThrow("remote_copyback_owned"); + expect(readRemoteWorkspaceAuthority((await workspace(f.workspaceId)).metadata)).toBeNull(); + expect(await db.select().from(executionWorkspaceSnapshots)).toHaveLength(0); + expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)))[0]).toEqual(f.owner); + } finally { release(); await owner; } + }); it.each(["legacy", "verified", "retry"])("recovers an archived Git workspace from its exact branch-deletion receipt (%s)", async receiptKind => { const f = await gitArchiveFixture(), closedAt = new Date(); await db.update(executionWorkspaces).set({ status: "archived", closedAt }).where(eq(executionWorkspaces.id, f.workspaceId)); @@ -1143,7 +1263,7 @@ describe.skipIf(!support.supported)("remote workspace migration", () => { const f = await seed(), closedAt = new Date("2026-10-20T12:00:00Z"); await db.update(executionWorkspaces).set({ status: "archived", closedAt }).where(eq(executionWorkspaces.id, f.workspaceId)); await db.update(issues).set({ status: "done" }).where(eq(issues.id, f.issueId)); - await expect(migrateRemoteWorkspace({ db, storage, ...f, apply: true })).resolves.toMatchObject({ applied: true }); + expect(await migrateRemoteWorkspace({ db, storage, ...f, apply: true })).toMatchObject({ applied: true }); const saved = await workspace(f.workspaceId), authority = readRemoteWorkspaceAuthority(saved.metadata)!; expect(saved.status).toBe("archived"); expect(saved.closedAt).toEqual(closedAt); const [issue] = await db.select().from(issues).where(eq(issues.id, f.issueId)); diff --git a/server/src/services/remote-workspace-migration.ts b/server/src/services/remote-workspace-migration.ts index 6ccbfebd609..258724ccfff 100644 --- a/server/src/services/remote-workspace-migration.ts +++ b/server/src/services/remote-workspace-migration.ts @@ -160,7 +160,8 @@ async function inspect(db: ReadDb, input: { companyId: string; workspaceId: stri sql`${heartbeatRuns.runnerProfileJson}->'nativeWorkspaceSync'->>'state' in ('prepared', 'finalized')`, )); const legacyRecoverySources: LegacyWorkspaceRecoverySource[] = []; - const onlineRecoveryAllowed = Boolean(input.onlineSource) && !input.writersStopped; + const onlineRecoveryAllowed = !input.writersStopped + && Boolean(input.onlineSource || workspace.status === "archived" && archivedGitSource?.recipe); for (const run of referencedRuns.filter(run => object(run.reference).state === "prepared").sort((a, b) => a.id.localeCompare(b.id))) { if (onlineRecoveryAllowed) { const source = await inspectOnlineLegacyWorkspaceRecoverySource(db, { workspace, runId: run.id }); @@ -414,11 +415,12 @@ type ReservedMigration = Awaited>; type MigrationPlan = { input: MigrationInput; onlineSource?: OnlineWorkspaceSource; reserved: ReservedMigration }; async function lockMigrationRecovery(lock: MigrationTransaction, plans: MigrationPlan[], cohort?: WorkspaceMigrationCohort) { - const online = plans.filter(plan => plan.onlineSource); + const online = plans.filter(plan => !plan.input.writersStopped + && (plan.onlineSource || plan.reserved.workspace.status === "archived" && plan.reserved.archivedGitSource?.recipe)); if (!online.length) return; if (!await lockOnlineLegacyWorkspaceRecoverySources(lock as unknown as Db, { companyId: online[0].input.companyId, - sources: online.flatMap(({ reserved, onlineSource }) => [...reserved.legacyRecoverySources, ...onlineSource!.companions.flatMap(proof => proof.sources)]), - holds: online.flatMap(({ reserved, onlineSource }) => [...reserved.onlineRecoveryHolds, ...onlineSource!.companions.flatMap(proof => proof.holds)]) })) + sources: online.flatMap(({ reserved, onlineSource }) => [...reserved.legacyRecoverySources, ...(onlineSource?.companions.flatMap(proof => proof.sources) ?? [])]), + holds: online.flatMap(({ reserved, onlineSource }) => [...reserved.onlineRecoveryHolds, ...(onlineSource?.companions.flatMap(proof => proof.holds) ?? [])]) })) throw new Error("remote_workspace_migration_blocked:remote_copyback_owned"); const readers = await Promise.allSettled(online.map(async ({ input, onlineSource, reserved }) => { const current = await inspect(input.db, { ...input, onlineSource, operationId: reserved.operation!.id }, cohort); @@ -531,9 +533,9 @@ async function validateMigrationPublication(tx: MigrationTransaction, input: Mig || journal.issueId !== operation.issueId || journal.cwd !== operation.cwd || canonical(journal.metadata) !== canonical(operation.metadata)) throw new Error("workspace_cohort_journal_changed"); } - await lockMigrationParents(tx, input.companyId, reserved.workspace.sourceIssueId, onlineSource ? [ + await lockMigrationParents(tx, input.companyId, reserved.workspace.sourceIssueId, onlineSource || reserved.archivedGitSource?.recipe ? [ ...reserved.legacyRecoverySources.map(source => source.issueId), - ...onlineSource.companions.flatMap(proof => proof.sources.map(source => source.issueId)), + ...(onlineSource?.companions.flatMap(proof => proof.sources.map(source => source.issueId)) ?? []), ] : []); for (const actionId of [...new Set([...reserved.finalizedRecoveryHolds, ...reserved.onlineRecoveryHolds].map(held => held.actionId))].sort()) { await tx.select({ id: issueRecoveryActions.id }).from(issueRecoveryActions).where(and( diff --git a/ui/src/components/task-chat/TaskChatComposer.test.tsx b/ui/src/components/task-chat/TaskChatComposer.test.tsx index f02cf4b6d8d..c53445468e8 100644 --- a/ui/src/components/task-chat/TaskChatComposer.test.tsx +++ b/ui/src/components/task-chat/TaskChatComposer.test.tsx @@ -13,7 +13,7 @@ import { parseRunnerGoalCommand, TaskChatComposer } from "./TaskChatComposer"; import { ComposerAddMenu } from "./ComposerAddMenu"; import { QuestionForm } from "./QuestionForm"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import type { Agent, HeartbeatRun } from "@paperclipai/shared"; +import type { Agent, HeartbeatRun, IssueWorkMode } from "@paperclipai/shared"; import { agentsApi } from "@/api/agents"; import { instanceSettingsApi } from "@/api/instanceSettings"; import { CoordinatorPresentation } from "@/components/coordinator/CoordinatorPresentation"; @@ -1201,6 +1201,143 @@ describe("TaskChatComposer", () => { expect(onAdd).toHaveBeenCalledTimes(2); }); + it("reflects an externally updated task mode after mounting", async () => { + const onAdd = vi.fn(); + const onWorkModeChange = vi.fn(); + renderComposer({ onAdd, onWorkModeChange, workMode: "standard" }); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')).toBeNull(); + + renderComposer({ onAdd, onWorkModeChange, workMode: "planning" }); + await flushAsync(); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode")) + .toBe("planning"); + }); + + it("keeps a local draft mode while persisted task data refreshes", async () => { + const onAdd = vi.fn(); + const onWorkModeChange = vi.fn(); + renderComposer({ onAdd, onWorkModeChange, workMode: "standard" }); + pressKey("Tab", { shiftKey: true }); + + renderComposer({ onAdd, onWorkModeChange, workMode: "ask" }); + await flushAsync(); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode")) + .toBe("planning"); + expect(onWorkModeChange).not.toHaveBeenCalled(); + + // Matching task data does not acknowledge an unsent local choice. + renderComposer({ onAdd, onWorkModeChange, workMode: "planning" }); + renderComposer({ onAdd, onWorkModeChange, workMode: "standard" }); + await flushAsync(); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode")) + .toBe("planning"); + }); + + it("keeps the chosen mode and message when an optimistic mode save rolls back", async () => { + const onAdd = vi.fn(); + let rejectSave!: () => void; + function OptimisticModeHarness() { + const [workMode, setWorkMode] = useState("standard"); + return { + setWorkMode(mode); + return new Promise((_resolve, reject) => { + rejectSave = () => { + setWorkMode("standard"); + reject(new Error("Mode save rejected")); + }; + }); + }} />; + } + render(); + pressKey("Tab", { shiftKey: true }); + typeText("Keep this planning message"); + pressKey("Enter"); + await flushAsync(); + await act(async () => rejectSave()); + await flushAsync(); + + expect(editable().textContent).toBe("Keep this planning message"); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode")) + .toBe("planning"); + expect(onAdd).not.toHaveBeenCalled(); + }); + + it.each([ + ["standard", "planning", "standard"], + ["ask", "standard", "ask"], + ] as const)("keeps next draft mode %s → %s → %s through successful send", async (savedMode, sentMode, nextMode) => { + let resolveSend!: () => void; + const onAdd = vi.fn().mockReturnValue(new Promise((resolve) => { resolveSend = resolve; })); + function PendingSendHarness() { + const [workMode, setWorkMode] = useState(savedMode); + return { setWorkMode(mode); }} />; + } + const selectedMode = () => container.querySelector('[data-testid="task-chat-composer-mode"]') + ?.getAttribute("data-pending-work-mode") ?? "standard"; + render(); + pressKey("Tab", { shiftKey: true }); + expect(selectedMode()).toBe(sentMode); + typeText("First message"); + pressKey("Enter"); + await flushAsync(); + expect(onAdd).toHaveBeenCalledTimes(1); + + if (nextMode === "standard") { + flushSync(() => container.querySelector('[data-testid="task-chat-composer-mode"]')!.click()); + } else { + pressKey("Tab", { shiftKey: true }); + pressKey("Tab", { shiftKey: true }); + } + expect(selectedMode()).toBe(nextMode); + typeText("Next message"); + await act(async () => resolveSend()); + await flushAsync(); + expect(selectedMode()).toBe(nextMode); + expect(editable().textContent).toBe("Next message"); + }); + + it.each([ + ["receipt", false], ["receipt", true], + ["resend", false], ["resend", true], + ] as const)("settles the sent mode through %s without clearing a newer choice (%s)", async (method, newerChoice) => { + const key = `mode-settlement-${method}-${newerChoice}`; + const onAdd = vi.fn().mockRejectedValueOnce(new CommentSubmissionUnknownError()).mockResolvedValue(undefined); + const onWorkModeChange = vi.fn().mockResolvedValue(undefined); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + client.setQueryData(queryKeys.auth.session, { user: { id: "user-1" } }); + const show = (workMode: IssueWorkMode, confirmedSubmissionIds = new Set()) => render( + , + ); + show("standard"); + await flushAsync(); + pressKey("Tab", { shiftKey: true }); + typeText("First message"); + pressKey("Enter", { metaKey: true }); + await flushAsync(); + expect(onAdd).toHaveBeenCalledTimes(1); + const intent = loadDraftSubmission(key)!; + expect(intent).not.toBeNull(); + show("planning"); + if (newerChoice) pressKey("Tab", { shiftKey: true }); + + if (method === "receipt") show("planning", new Set([intent.attemptId])); + else { + const resend = Array.from(container.querySelectorAll("button")).find(button => button.textContent === "Resend now")!; + flushSync(() => resend.click()); + } + await flushAsync(); + expect(loadDraftSubmission(key)).toBeNull(); + show("standard"); + await flushAsync(); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode") ?? "standard") + .toBe(newerChoice ? "ask" : "standard"); + expect(onAdd).toHaveBeenCalledTimes(method === "receipt" ? 1 : 2); + if (method === "resend") expect(onAdd.mock.calls[1]?.[4]).toBe(intent.attemptId); + client.clear(); + }); + it("cycles the pending mode with Shift+Tab and applies it on submit", async () => { const onAdd = vi.fn().mockResolvedValue(undefined); const onWorkModeChange = vi.fn().mockResolvedValue(undefined); @@ -1226,6 +1363,11 @@ describe("TaskChatComposer", () => { expect(onWorkModeChange).toHaveBeenCalledWith("planning"); expect(onAdd).toHaveBeenCalledWith("do the plan", undefined, undefined, undefined, expect.any(String)); + + renderComposer({ onAdd, onWorkModeChange, workMode: "ask" }); + await flushAsync(); + expect(container.querySelector('[data-testid="task-chat-composer-mode"]')?.getAttribute("data-pending-work-mode")) + .toBe("ask"); }); it("cycles Auto, Plan, and Ask modes with Cmd+Period while focused", () => { diff --git a/ui/src/components/task-chat/TaskChatComposer.tsx b/ui/src/components/task-chat/TaskChatComposer.tsx index bf392b3d6d0..9e44a506187 100644 --- a/ui/src/components/task-chat/TaskChatComposer.tsx +++ b/ui/src/components/task-chat/TaskChatComposer.tsx @@ -354,7 +354,7 @@ type ComposerAttachment = { type QueuedEditBackup = { body: string; attachments: ComposerAttachment[]; - pendingMode: IssueWorkMode; + pendingMode: IssueWorkMode | null; pendingAssignee: string | null; }; @@ -451,8 +451,10 @@ export function TaskChatComposer({ useState(null); const [takeoverControlsSlot, setTakeoverControlsSlot] = useState(null); - const [localPendingMode, setPendingMode] = useState(workMode); - const pendingMode = creation ? workMode : localPendingMode; + const [pendingModeOverride, setPendingMode] = useState(null); + const modeChoiceGeneration = useRef(0); + const sentModeChoice = useRef<{ attemptId: string; generation: number } | null>(null); + const pendingMode = creation ? workMode : pendingModeOverride ?? workMode; const [pendingAssignee, setPendingAssignee] = useState(null); const coordinator = useCoordinatorPresentation(); const [localRunSettings, setLocalRunSettings] = useState(null); @@ -556,7 +558,7 @@ export function TaskChatComposer({ queuedEditBackupRef.current = { body: bodyRef.current, attachments: attachmentsRef.current, - pendingMode, + pendingMode: pendingModeOverride, pendingAssignee: pendingAssigneeRef.current, }; if (draftKey) saveDraft(draftKey, bodyRef.current); @@ -584,7 +586,7 @@ export function TaskChatComposer({ } previousQueuedEditIdRef.current = nextId; - }, [draftKey, pendingMode, queuedEdit]); + }, [draftKey, pendingModeOverride, queuedEdit]); useEffect(() => { if (!draftKey || queuedEdit) return; @@ -704,10 +706,18 @@ export function TaskChatComposer({ } function changeMode(mode: IssueWorkMode) { + modeChoiceGeneration.current += 1; setPendingMode(mode); if (creation) void onWorkModeChange?.(mode); } + function settleSentMode(attemptId: string) { + const choice = sentModeChoice.current; + if (choice?.attemptId !== attemptId) return; + if (modeChoiceGeneration.current === choice.generation) setPendingMode(null); + sentModeChoice.current = null; + } + async function uploadImage(file: File): Promise { const id = crypto.randomUUID(); setAttachments((prev) => [ @@ -926,6 +936,7 @@ export function TaskChatComposer({ const submittedBody = bodyRef.current; const submittedAttachments = attachmentsRef.current; const submittedAssignee = pendingAssigneeRef.current; + const submittedModeGeneration = modeChoiceGeneration.current; const trimmed = submittedBody.trim(); const goalCommand = queuedEdit ? ({ matched: false } as const) @@ -1088,6 +1099,7 @@ export function TaskChatComposer({ pendingDraftRef.current = { draftKey, attemptId, submittedBody, submittedAttachmentIds: attachmentIds, request, requestOwner }; changeBody(bodyRef.current); } + sentModeChoice.current = { attemptId, generation: submittedModeGeneration }; if (runSettings) { await onAdd(fullBody, reopen, reassignment, attachmentIds.length ? attachmentIds : undefined, attemptId, runSettings); } else { @@ -1098,6 +1110,7 @@ export function TaskChatComposer({ if (draftKey) settleDraftSubmission(draftKey, attemptId, mountedTaskKey.current === draftKey ? bodyRef.current : undefined); if (mountedTaskKey.current !== draftKey) return; + settleSentMode(attemptId); const submittedIds = new Set(submittedAttachments.map((item) => item.id)); setAttachments((current) => current.filter((item) => !submittedIds.has(item.id)), @@ -1149,6 +1162,7 @@ export function TaskChatComposer({ const settlementKey = JSON.stringify([draftKey, submission.attemptId]); if (settledSubmissionIds.current.has(settlementKey)) return; settledSubmissionIds.current.add(settlementKey); + settleSentMode(submission.attemptId); const nextDraft = submission.nextDraftOffset === undefined ? "" : bodyRef.current.slice(submission.nextDraftOffset); if (draftKey) settleDraftSubmission(draftKey, submission.attemptId, nextDraft);