Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
172 changes: 170 additions & 2 deletions apps/mobile/src/state/use-composer-drafts.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,13 @@ import {
ThreadId,
} from "@t3tools/contracts";
import { onTestFinished, vi } from "vite-plus/test";
import { resolveNewTaskModelSelection } from "../lib/modelOptions";

const composerDraftFileMocks = vi.hoisted(() => {
let document = JSON.stringify({ schemaVersion: 1, drafts: {} });
let legacyDocument: string | null = null;
let legacyReadError: Error | null = null;
const legacyReads = vi.fn();
let readError: Error | null = null;
let writeError: Error | null = null;
let releaseRead: (() => void) | null = null;
Expand All @@ -23,6 +27,13 @@ const composerDraftFileMocks = vi.hoisted(() => {

return {
readImage,
legacyReads,
setLegacyDocument(value: string | null) {
legacyDocument = value;
},
setLegacyReadError(error: Error | null) {
legacyReadError = error;
},
blockRead() {
readBarrier = new Promise<void>((resolve) => {
releaseRead = resolve;
Expand Down Expand Up @@ -64,15 +75,32 @@ const composerDraftFileMocks = vi.hoisted(() => {
}
},
File: class {
exists = true;
readonly legacy: boolean;
constructor(_directory: unknown, fileName: string) {
this.legacy = fileName === "model-option-memory.json";
}
get exists() {
return !this.legacy || legacyDocument !== null;
}
parentDirectory = null;

create() {}

moveSync() {}
moveSync() {
if (this.legacy) throw new Error("Legacy memory must remain read-only");
}

delete() {
if (this.legacy) throw new Error("Legacy memory must remain read-only");
}

async text() {
await readBarrier;
if (this.legacy) {
legacyReads();
if (legacyReadError) throw legacyReadError;
return legacyDocument!;
}
if (readError) throw readError;
return document;
}
Expand All @@ -82,6 +110,7 @@ const composerDraftFileMocks = vi.hoisted(() => {
}

write(value: string) {
if (this.legacy) throw new Error("Legacy memory must remain read-only");
if (writeError) {
throw writeError;
}
Expand Down Expand Up @@ -206,6 +235,9 @@ afterEach(() => {
resetComposerDraftsLoadState();
composerDraftFileMocks.setDocument({ schemaVersion: 1, drafts: {} });
composerDraftFileMocks.setReadError(null);
composerDraftFileMocks.setLegacyDocument(null);
composerDraftFileMocks.setLegacyReadError(null);
composerDraftFileMocks.legacyReads.mockClear();
composerDraftFileMocks.setWriteError(null);
composerDraftFileMocks.setNextWriteBarrier(null);
composerDraftFileMocks.setOnWrite(null);
Expand Down Expand Up @@ -1956,6 +1988,142 @@ describe("mobile composer drafts", () => {
});
});

it("imports trial memory once, with per-model disk and in-flight choices winning", async () => {
vi.useFakeTimers();
const options = (value: string) => [{ id: "thinking", value }];
composerDraftFileMocks.setDocument({
schemaVersion: 1,
drafts: { "environment-1:saved": DRAFT },
modelOptionMemory: { pi: { a: options("high"), b: options("medium") } },
});
composerDraftFileMocks.setLegacyDocument(
JSON.stringify({
schemaVersion: 1,
byInstance: {
pi: { a: options("low"), b: options("low"), c: options("off"), empty: [] },
other: { a: options("max") },
},
}),
);
composerDraftFileMocks.blockRead();
ensureComposerDraftsLoaded();
appAtomRegistry.set(modelOptionMemoryAtom, { pi: { a: options("xhigh") } });
composerDraftFileMocks.releaseRead();
await waitForComposerDraftsLoaded();
const expected = {
pi: { a: options("xhigh"), b: options("medium"), c: options("off") },
other: { a: options("max") },
};
expect(appAtomRegistry.get(modelOptionMemoryAtom)).toEqual(expected);
// A failed migration write leaves saved drafts retryable.
composerDraftFileMocks.setWriteError(new Error("disk unavailable"));
await expect(flushComposerDrafts()).rejects.toMatchObject({ operation: "write" });
expect(
JSON.parse(composerDraftFileMocks.getDocument()).legacyModelOptionMemoryImported,
).toBeUndefined();
composerDraftFileMocks.setWriteError(null);
await flushComposerDrafts();
expect(JSON.parse(composerDraftFileMocks.getDocument())).toMatchObject({
drafts: { "environment-1:saved": DRAFT },
modelOptionMemory: expected,
legacyModelOptionMemoryImported: true,
});
resetComposerDraftsLoadState();
appAtomRegistry.set(modelOptionMemoryAtom, {});
composerDraftFileMocks.legacyReads.mockClear();
await waitForComposerDraftsLoaded();
expect(appAtomRegistry.get(modelOptionMemoryAtom)).toEqual(expected);
expect(composerDraftFileMocks.legacyReads).not.toHaveBeenCalled();
});

it.each(["fresh", "existing", "explicit"] as const)(
"imports legacy options without replacing the %s draft selection or sticky priority",
async (kind) => {
vi.useFakeTimers();
const selection = {
instanceId: ProviderInstanceId.make("codex"),
model: "chosen",
options: [{ id: "reasoning", value: "high" }],
};
const draft = {
text: "Keep this prompt",
attachments: [],
modelSelection: selection,
interactionMode: "plan" as const,
};
const key = kind === "existing" ? "environment-1:thread" : "new-task:chosen";
composerDraftFileMocks.setDocument({
schemaVersion: 1,
drafts: kind === "fresh" ? {} : { [key]: draft },
stickyModelSelection: selection,
});
composerDraftFileMocks.setLegacyDocument(
JSON.stringify({
schemaVersion: 1,
byInstance: { codex: { chosen: [{ id: "reasoning", value: "low" }] } },
}),
);
await waitForComposerDraftsLoaded();
const restored = appAtomRegistry.get(composerDraftsAtom)[key];
expect(restored).toEqual(kind === "fresh" ? undefined : draft);
expect(appAtomRegistry.get(stickyComposerModelSelectionAtom)).toEqual(selection);
expect(appAtomRegistry.get(modelOptionMemoryAtom)).toEqual({
codex: { chosen: [{ id: "reasoning", value: "low" }] },
});
expect(
resolveNewTaskModelSelection({
draftSelection: restored?.modelSelection ?? null,
stickySelection: appAtomRegistry.get(stickyComposerModelSelectionAtom),
projectDefaultSelection: null,
modelOptions: [],
}),
).toEqual(selection);
await flushComposerDrafts();
expect(JSON.parse(composerDraftFileMocks.getDocument()).stickyModelSelection).toEqual(
selection,
);
},
);

it.each(["json", "schema", "read"] as const)(
"preserves drafts and retries the import after a legacy %s failure",
async (failure) => {
vi.useFakeTimers();
composerDraftFileMocks.setDocument({
schemaVersion: 1,
drafts: { "environment-1:saved": DRAFT },
});
composerDraftFileMocks.setLegacyDocument(
failure === "json" ? "{" : JSON.stringify({ schemaVersion: 999, byInstance: {} }),
);
if (failure === "read") composerDraftFileMocks.setLegacyReadError(new Error("unavailable"));
setComposerDraftText("environment-1:new", "New edits");
await flushComposerDrafts();
expect(JSON.parse(composerDraftFileMocks.getDocument())).toMatchObject({
drafts: { "environment-1:saved": DRAFT, "environment-1:new": { text: "New edits" } },
});
expect(
JSON.parse(composerDraftFileMocks.getDocument()).legacyModelOptionMemoryImported,
).toBeUndefined();
resetComposerDraftsLoadState();
composerDraftFileMocks.setLegacyReadError(null);
composerDraftFileMocks.setLegacyDocument(
JSON.stringify({
schemaVersion: 1,
byInstance: { pi: { a: [{ id: "thinking", value: "max" }] } },
}),
);
await flushComposerDrafts();
expect(JSON.parse(composerDraftFileMocks.getDocument()).legacyModelOptionMemoryImported).toBe(
true,
);
expect(appAtomRegistry.get(composerDraftsAtom)["environment-1:saved"]).toEqual(DRAFT);
expect(JSON.parse(composerDraftFileMocks.getDocument()).modelOptionMemory).toEqual({
pi: { a: [{ id: "thinking", value: "max" }] },
});
},
);

it("waits for hydration before persisting the latest composer state", async () => {
vi.useFakeTimers();
composerDraftFileMocks.setDocument({
Expand Down
99 changes: 79 additions & 20 deletions apps/mobile/src/state/use-composer-drafts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -398,16 +398,21 @@ const ComposerDraftSchema = Schema.Struct({
project: Schema.optional(ComposerDraftProjectSchema),
});

const ModelOptionMemorySchema = Schema.Record(
Schema.String,
Schema.Record(Schema.String, Schema.Array(ProviderOptionSelectionSchema)),
);
const LEGACY_MODEL_OPTION_MEMORY_FILE = "model-option-memory.json";
const decodeLegacyModelOptionMemory = Schema.decodeUnknownSync(
Schema.Struct({ schemaVersion: Schema.Literal(1), byInstance: ModelOptionMemorySchema }),
);

const PersistedComposerDraftsSchema = Schema.Struct({
schemaVersion: Schema.Literal(COMPOSER_DRAFTS_SCHEMA_VERSION),
drafts: Schema.Record(Schema.String, ComposerDraftSchema),
stickyModelSelection: Schema.optional(ModelSelectionSchema),
modelOptionMemory: Schema.optional(
Schema.Record(
Schema.String,
Schema.Record(Schema.String, Schema.Array(ProviderOptionSelectionSchema)),
),
),
modelOptionMemory: Schema.optional(ModelOptionMemorySchema),
legacyModelOptionMemoryImported: Schema.optional(Schema.Literal(true)),
cloudAccountId: Schema.optional(Schema.String),
signedOutDrafts: Schema.optional(
Schema.Record(
Expand Down Expand Up @@ -466,12 +471,14 @@ export const composerCloudDraftsAtom = Atom.make<ComposerCloudDraftState>({
let loadPromise: Promise<void> | null = null;
let persistTimer: ReturnType<typeof setTimeout> | null = null;
let persistRetryNeeded = false;
let legacyModelOptionMemoryImported = false;
const persistenceQueue = new SerializedAsyncQueue();

/** Resets module-level state between test runs. */
export function resetComposerDraftsLoadState(): void {
loadPromise = null;
persistRetryNeeded = false;
legacyModelOptionMemoryImported = false;
}

function attachmentContextRecord(
Expand Down Expand Up @@ -627,6 +634,7 @@ export function decodePersistedComposerState(value: unknown): {
readonly drafts: Record<string, ComposerDraft>;
readonly stickyModelSelection: ModelSelection | null;
readonly modelOptionMemory: ModelOptionMemoryState;
readonly legacyModelOptionMemoryImported?: true;
readonly cloudDrafts: ComposerCloudDraftState;
} {
const parsed = decodePersistedComposerDraftsDocument(value);
Expand Down Expand Up @@ -666,6 +674,7 @@ export function decodePersistedComposerState(value: unknown): {
),
stickyModelSelection: parsed.stickyModelSelection ?? null,
modelOptionMemory: parsed.modelOptionMemory ?? {},
...(parsed.legacyModelOptionMemoryImported ? { legacyModelOptionMemoryImported: true } : {}),
cloudDrafts: {
accountId: parsed.cloudAccountId ?? null,
signedOut: Object.fromEntries(
Expand All @@ -687,11 +696,57 @@ export function decodePersistedComposerState(value: unknown): {
};
}

async function getComposerDraftsFile() {
async function getComposerDraftsFile(fileName = COMPOSER_DRAFTS_FILE) {
const { Directory, File, Paths } = await import("expo-file-system");
const directory = new Directory(Paths.document, COMPOSER_DRAFTS_DIRECTORY);
directory.create({ idempotent: true, intermediates: true });
return new File(directory, COMPOSER_DRAFTS_FILE);
return new File(directory, fileName);
}

function mergeModelOptionMemory(
persisted: ModelOptionMemoryState,
current: ModelOptionMemoryState,
): ModelOptionMemoryState {
return {
...persisted,
...Object.fromEntries(
Object.entries(current).map(([instanceId, models]) => [
instanceId,
{ ...persisted[instanceId], ...models },
]),
),
};
}

async function loadLegacyModelOptionMemory(): Promise<ModelOptionMemoryState | null> {
let operation: ComposerDraftPersistenceError["operation"] = "open";
try {
const file = await getComposerDraftsFile(LEGACY_MODEL_OPTION_MEMORY_FILE);
if (!file.exists) return null;
operation = "read";
const raw = await file.text();
operation = "decode";
const { byInstance } = decodeLegacyModelOptionMemory(JSON.parse(raw) as unknown);
return Object.fromEntries(
Object.entries(byInstance).map(([instanceId, models]) => [
instanceId,
Object.fromEntries(Object.entries(models).filter(([, options]) => options.length > 0)),
]),
);
} catch (cause) {
// Legacy option memory must never prevent saved drafts from hydrating. Leave
// the import unmarked on failure so a later launch can retry the old file.
console.warn(
"[composer-drafts] could not import legacy model options",
new ComposerDraftPersistenceError({
operation,
directory: COMPOSER_DRAFTS_DIRECTORY,
fileName: LEGACY_MODEL_OPTION_MEMORY_FILE,
cause,
}),
);
return null;
}
}

async function loadPersistedComposerState(): Promise<
Expand Down Expand Up @@ -737,6 +792,7 @@ async function writePersistedComposerState(
const document = {
schemaVersion: COMPOSER_DRAFTS_SCHEMA_VERSION,
drafts: nonEmptyDrafts,
...(legacyModelOptionMemoryImported ? { legacyModelOptionMemoryImported: true } : {}),
...(stickyModelSelection ? { stickyModelSelection } : {}),
...(Object.keys(appAtomRegistry.get(modelOptionMemoryAtom)).length > 0
? { modelOptionMemory: appAtomRegistry.get(modelOptionMemoryAtom) }
Expand Down Expand Up @@ -1028,7 +1084,13 @@ export function ensureComposerDraftsLoaded(): void {
if (loadPromise !== null) {
return;
}
const loading = loadPersistedComposerState().then((persisted) => {
const loading = loadPersistedComposerState().then(async (persisted) => {
const legacyMemory = persisted.legacyModelOptionMemoryImported
? null
: await loadLegacyModelOptionMemory();
const persistedMemory = mergeModelOptionMemory(legacyMemory ?? {}, persisted.modelOptionMemory);
legacyModelOptionMemoryImported =
persisted.legacyModelOptionMemoryImported === true || legacyMemory !== null;
appAtomRegistry.set(composerCloudDraftsAtom, persisted.cloudDrafts);
if (Object.keys(persisted.drafts).length > 0) {
const current = appAtomRegistry.get(composerDraftsAtom);
Expand All @@ -1043,18 +1105,15 @@ export function ensureComposerDraftsLoaded(): void {
) {
appAtomRegistry.set(stickyComposerModelSelectionAtom, persisted.stickyModelSelection);
}
if (Object.keys(persisted.modelOptionMemory).length > 0) {
const current = appAtomRegistry.get(modelOptionMemoryAtom);
appAtomRegistry.set(modelOptionMemoryAtom, {
...persisted.modelOptionMemory,
...Object.fromEntries(
Object.entries(current).map(([instanceId, models]) => [
instanceId,
{ ...(persisted.modelOptionMemory[instanceId] ?? {}), ...models },
]),
),
});
if (Object.keys(persistedMemory).length > 0) {
appAtomRegistry.set(
modelOptionMemoryAtom,
mergeModelOptionMemory(persistedMemory, appAtomRegistry.get(modelOptionMemoryAtom)),
);
}
// The marker and imported choices land atomically with drafts through the
// existing retryable writer. Never delete or write the legacy file.
if (legacyMemory !== null) schedulePersistComposerState();
});
loadPromise = loading;
// Handle fire-and-forget hook loads without swallowing failures from the
Expand Down