diff --git a/src/background/background.ts b/src/background/background.ts
index 54081de..8554a7c 100644
--- a/src/background/background.ts
+++ b/src/background/background.ts
@@ -23,6 +23,7 @@ import { toEmail, EMAIL_ATTRIBUTE_TYPE, findHtmlBody } from "../lib/utils";
import { getOrCreateLocalFolder } from "../lib/folders";
import { isPGEncrypted, wasPGEncrypted } from "../lib/detection";
import { dispatchRuntimeMessage } from "./runtime-router";
+import { handleAfterSend } from "./sent-copy";
import type {
Policy,
SerializedRecipient,
@@ -102,37 +103,13 @@ browser.runtime.onMessage.addListener(
browser.compose.onBeforeSend.addListener(handleBeforeSend);
-browser.compose.onAfterSend.addListener(async (tab, sendInfo) => {
- const state = composeTabs.get(tab.id);
- if (!state?.sentMimeData) return;
-
- try {
- for (const msg of sendInfo.messages) {
- if (await isPGEncrypted(msg.id)) {
- const localFolder = await getOrCreateLocalFolder("PostGuard Sent");
- if (localFolder) {
- const file = new File([state.sentMimeData as BlobPart], "sent.eml", {
- type: "text/plain",
- });
- const localMsg = await browser.messages.import(
- file,
- localFolder.id
- );
- await browser.messages.move([localMsg.id], msg.folder.id);
- await browser.messages.delete([msg.id], true);
- }
- }
- }
- } catch (e) {
- console.error("[PostGuard] Failed to manage sent copy:", e);
- notifyError("sentCopyError");
- } finally {
- cleanupComposeTab(tab.id);
- clearInFlightUpload(tab.id);
- persistEncryptState().catch(console.warn);
- persistInFlightUploads().catch(console.warn);
- }
-});
+browser.compose.onAfterSend.addListener((tab, sendInfo) =>
+ handleAfterSend(tab, sendInfo, {
+ notifyError,
+ isPGEncrypted,
+ getOrCreateLocalFolder,
+ }),
+);
// Clean up decryptedMessages when messages are deleted
browser.messages.onDeleted.addListener((deletedMessages) => {
diff --git a/src/background/sent-copy.ts b/src/background/sent-copy.ts
new file mode 100644
index 0000000..e04cc54
--- /dev/null
+++ b/src/background/sent-copy.ts
@@ -0,0 +1,71 @@
+///
+
+import {
+ composeTabs,
+ cleanupComposeTab,
+ clearInFlightUpload,
+ persistEncryptState,
+ persistInFlightUploads,
+} from "./state";
+
+// Pulled out of background.ts so the post-send cleanup can be exercised
+// under vitest without standing up the whole background script. Deps that
+// hit Thunderbird APIs we cannot easily mock from the browser-mock helper
+// are injected, and the rest goes through the browser-mock surface.
+
+export interface SentMessage {
+ id: number;
+ folder: { id: unknown };
+}
+
+export interface SentInfo {
+ messages: SentMessage[];
+}
+
+export interface SentCopyDeps {
+ notifyError: (messageKey: string) => void;
+ isPGEncrypted: (msgId: number) => Promise;
+ getOrCreateLocalFolder: (
+ name: string,
+ ) => Promise<{ id: unknown } | undefined>;
+}
+
+/**
+ * Post-send handler. For each message the user just sent, if it is a
+ * PostGuard ciphertext, swap the encrypted copy in Sent with the plaintext
+ * MIME stashed during encryption. State cleanup happens in `finally` so a
+ * failure mid-swap never strands the compose tab in `composeTabs`.
+ */
+export async function handleAfterSend(
+ tab: { id: number },
+ sendInfo: SentInfo,
+ deps: SentCopyDeps,
+): Promise {
+ const state = composeTabs.get(tab.id);
+ if (!state?.sentMimeData) return;
+
+ try {
+ for (const msg of sendInfo.messages) {
+ if (!(await deps.isPGEncrypted(msg.id))) continue;
+ const localFolder = await deps.getOrCreateLocalFolder("PostGuard Sent");
+ if (!localFolder) continue;
+ const file = new File([state.sentMimeData as BlobPart], "sent.eml", {
+ type: "text/plain",
+ });
+ const localMsg = await browser.messages.import(
+ file,
+ localFolder.id as any,
+ );
+ await browser.messages.move([localMsg.id], msg.folder.id as any);
+ await browser.messages.delete([msg.id], true);
+ }
+ } catch (e) {
+ console.error("[PostGuard] Failed to manage sent copy:", e);
+ deps.notifyError("sentCopyError");
+ } finally {
+ cleanupComposeTab(tab.id);
+ clearInFlightUpload(tab.id);
+ persistEncryptState().catch(console.warn);
+ persistInFlightUploads().catch(console.warn);
+ }
+}
diff --git a/tests/sent-copy.test.ts b/tests/sent-copy.test.ts
index f1b2970..82a2a33 100644
--- a/tests/sent-copy.test.ts
+++ b/tests/sent-copy.test.ts
@@ -1,22 +1,141 @@
-import { describe, it, expect } from "vitest";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { handleAfterSend, type SentInfo } from "../src/background/sent-copy";
+import {
+ composeTabs,
+ decryptedMessages,
+ inFlightUploads,
+ pendingCryptoPopups,
+ pendingPolicyEditors,
+} from "../src/background/state";
+import { installBrowserMock, type BrowserMock } from "./helpers/browser-mock";
-// These tests verify the onAfterSend sent copy management.
-// They require browser API mocks.
+let mock: BrowserMock;
+let notifyError: ReturnType;
+let isPGEncrypted: ReturnType;
+let getOrCreateLocalFolder: ReturnType;
+let importSpy: ReturnType;
+let moveSpy: ReturnType;
+let deleteSpy: ReturnType;
+
+const SENT_FOLDER = { id: "sent://Sent" };
+const LOCAL_FOLDER = { id: "local://PostGuard Sent" };
+
+function deps() {
+ return { notifyError, isPGEncrypted, getOrCreateLocalFolder };
+}
+
+function sentInfo(messageIds: number[] = [1]): SentInfo {
+ return {
+ messages: messageIds.map((id) => ({ id, folder: SENT_FOLDER })),
+ };
+}
+
+beforeEach(() => {
+ mock = installBrowserMock();
+ composeTabs.clear();
+ decryptedMessages.clear();
+ pendingCryptoPopups.clear();
+ pendingPolicyEditors.clear();
+ inFlightUploads.clear();
+
+ notifyError = vi.fn();
+ isPGEncrypted = vi.fn(async () => true);
+ getOrCreateLocalFolder = vi.fn(async () => LOCAL_FOLDER);
+
+ importSpy = vi.spyOn(browser.messages, "import");
+ moveSpy = vi.spyOn(browser.messages, "move");
+ deleteSpy = vi.spyOn(browser.messages, "delete");
+});
describe("onAfterSend — sent copy management", () => {
- it.todo("should skip when no sentMimeData is stored");
+ it("should skip when no sentMimeData is stored", async () => {
+ composeTabs.set(7, { encrypt: true });
+ await handleAfterSend({ id: 7 }, sentInfo([1]), deps());
+
+ expect(isPGEncrypted).not.toHaveBeenCalled();
+ expect(importSpy).not.toHaveBeenCalled();
+ });
+
+ it("should skip non-PG-encrypted sent messages", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+ isPGEncrypted.mockResolvedValueOnce(false);
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+ expect(isPGEncrypted).toHaveBeenCalledWith(42);
+ expect(getOrCreateLocalFolder).not.toHaveBeenCalled();
+ expect(importSpy).not.toHaveBeenCalled();
+ });
+
+ it("should import plaintext MIME into PostGuard Sent folder", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+ expect(getOrCreateLocalFolder).toHaveBeenCalledWith("PostGuard Sent");
+ expect(importSpy).toHaveBeenCalledTimes(1);
+ const [file, folder] = importSpy.mock.calls[0] as [File, unknown];
+ expect(file).toBeInstanceOf(File);
+ expect(file.name).toBe("sent.eml");
+ expect(folder).toBe(LOCAL_FOLDER.id);
+ });
+
+ it("should move imported message to original sent folder", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+ expect(moveSpy).toHaveBeenCalledWith([9999], SENT_FOLDER.id);
+ });
+
+ it("should delete the encrypted copy from sent folder", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+ expect(deleteSpy).toHaveBeenCalledWith([42], true);
+ });
+
+ it("should notify user when sent copy management fails", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+ importSpy.mockRejectedValueOnce(new Error("import boom"));
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+ expect(notifyError).toHaveBeenCalledWith("sentCopyError");
+ });
+
+ it("should clean up compose tab state after send", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+ inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() });
+
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
- it.todo("should skip non-PG-encrypted sent messages");
+ expect(composeTabs.has(7)).toBe(false);
+ expect(inFlightUploads.has(7)).toBe(false);
+ });
- it.todo("should import plaintext MIME into PostGuard Sent folder");
+ it("should clean up compose tab state even on failure", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+ inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() });
+ importSpy.mockRejectedValueOnce(new Error("nope"));
- it.todo("should move imported message to original sent folder");
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
- it.todo("should delete the encrypted copy from sent folder");
+ expect(composeTabs.has(7)).toBe(false);
+ expect(inFlightUploads.has(7)).toBe(false);
+ });
- it.todo("should notify user when sent copy management fails");
+ it("should skip the swap when no local folder is available", async () => {
+ composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+ getOrCreateLocalFolder.mockResolvedValueOnce(undefined);
- it.todo("should clean up compose tab state after send");
+ await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
- it.todo("should clean up compose tab state even on failure");
+ expect(importSpy).not.toHaveBeenCalled();
+ expect(moveSpy).not.toHaveBeenCalled();
+ expect(deleteSpy).not.toHaveBeenCalled();
+ // Cleanup still runs.
+ expect(composeTabs.has(7)).toBe(false);
+ });
});