1086 lines · 35.4 KB · TypeScript
RawDownload
import { useAtomValue } from "@effect/atom-react";import {  ModelSelection as ModelSelectionSchema,  PROVIDER_SEND_TURN_MAX_ATTACHMENTS,  ProviderInteractionMode as ProviderInteractionModeSchema,  RuntimeMode as RuntimeModeSchema,  type EnvironmentId,  type ModelSelection,  type ProviderInteractionMode,  type RuntimeMode,} from "@t3tools/contracts";import * as Schema from "effect/Schema";import { useEffect } from "react";import { Atom } from "effect/unstable/reactivity"; import { writeFileAtomically } from "../lib/atomic-file";import { DraftComposerAttachmentSchema } from "../lib/composer-image-schema";import {  composerAttachmentFileReferenceKey,  isComposerAttachmentFileRetained,  retainComposerAttachmentFile,} from "../lib/composerAttachmentFiles";import type { DraftComposerAttachment, DraftComposerFileAttachment } from "../lib/composerImages";import { SerializedAsyncQueue } from "../lib/serialized-async-queue";import { appAtomRegistry } from "./atom-registry";import { flushThreadOutbox, threadOutboxManager } from "./thread-outbox"; const COMPOSER_DRAFTS_SCHEMA_VERSION = 1;const COMPOSER_DRAFTS_DIRECTORY = "composer-drafts";const COMPOSER_DRAFTS_FILE = "drafts.json";const PERSIST_DEBOUNCE_MS = 200; export class ComposerDraftPersistenceError extends Schema.TaggedErrorClass<ComposerDraftPersistenceError>()(  "ComposerDraftPersistenceError",  {    operation: Schema.Literals(["open", "read", "decode", "encode", "write", "hydrate"]),    directory: Schema.String,    fileName: Schema.String,    cause: Schema.Defect(),  },) {  override get message(): string {    return `Composer draft persistence operation ${this.operation} failed for ${this.directory}/${this.fileName}.`;  }} export interface ComposerDraft {  readonly text: string;  readonly attachments: ReadonlyArray<DraftComposerAttachment>;  readonly importedShareIds?: ReadonlyArray<string>;  readonly modelSelection?: ModelSelection;  readonly runtimeMode?: RuntimeMode;  readonly interactionMode?: ProviderInteractionMode;  readonly workspaceSelection?: ComposerDraftWorkspaceSelection;} export interface ComposerDraftContent {  readonly text: string;  readonly attachments: ReadonlyArray<DraftComposerAttachment>;  readonly sourceShareId?: string;} export interface ComposerDraftWorkspaceSelection {  readonly mode: "local" | "worktree";  readonly branch: string | null;  readonly worktreePath: string | null;  readonly startFromOrigin?: boolean;} export type ComposerDraftSettingsUpdate = Pick<  ComposerDraft,  "modelSelection" | "runtimeMode" | "interactionMode" | "workspaceSelection">; const ComposerDraftWorkspaceSelectionSchema = Schema.Struct({  mode: Schema.Literals(["local", "worktree"]),  branch: Schema.NullOr(Schema.String),  worktreePath: Schema.NullOr(Schema.String),  startFromOrigin: Schema.optional(Schema.Boolean),}); const ComposerDraftSchema = Schema.Struct({  text: Schema.String,  attachments: Schema.Array(DraftComposerAttachmentSchema),  importedShareIds: Schema.optional(Schema.Array(Schema.String)),  modelSelection: Schema.optional(ModelSelectionSchema),  runtimeMode: Schema.optional(RuntimeModeSchema),  interactionMode: Schema.optional(ProviderInteractionModeSchema),  workspaceSelection: Schema.optional(ComposerDraftWorkspaceSelectionSchema),}); const PersistedComposerDraftsSchema = Schema.Struct({  schemaVersion: Schema.Literal(COMPOSER_DRAFTS_SCHEMA_VERSION),  drafts: Schema.Record(Schema.String, ComposerDraftSchema),  stickyModelSelection: Schema.optional(ModelSelectionSchema),}); const decodePersistedComposerDraftsDocument = Schema.decodeUnknownSync(  PersistedComposerDraftsSchema,); const EMPTY_DRAFT: ComposerDraft = {  text: "",  attachments: [],}; export const composerDraftsAtom = Atom.make<Record<string, ComposerDraft>>({}).pipe(  Atom.keepAlive,  Atom.withLabel("mobile:composer-drafts"),); export const stickyComposerModelSelectionAtom = Atom.make<ModelSelection | null>(null).pipe(  Atom.keepAlive,  Atom.withLabel("mobile:sticky-composer-model-selection"),); let loadPromise: Promise<void> | null = null;let persistTimer: ReturnType<typeof setTimeout> | null = null;const persistenceQueue = new SerializedAsyncQueue(); /** Resets module-level state between test runs. */export function resetComposerDraftsLoadState(): void {  loadPromise = null;} function normalizeDraft(draft: ComposerDraft | undefined): ComposerDraft {  if (!draft) {    return EMPTY_DRAFT;  }  return {    ...draft,    text: draft.text,    attachments: draft.attachments,  };} export function getComposerDraftSnapshot(draftKey: string): ComposerDraft {  return normalizeDraft(appAtomRegistry.get(composerDraftsAtom)[draftKey]);} export function isComposerDraftEmpty(draft: ComposerDraft): boolean {  return isEmptyDraft(draft);} function isEmptyDraft(draft: ComposerDraft): boolean {  return (    draft.text.length === 0 &&    draft.attachments.length === 0 &&    draft.modelSelection === undefined &&    draft.runtimeMode === undefined &&    draft.interactionMode === undefined &&    draft.workspaceSelection === undefined  );} export function decodePersistedComposerState(value: unknown): {  readonly drafts: Record<string, ComposerDraft>;  readonly stickyModelSelection: ModelSelection | null;} {  const parsed = decodePersistedComposerDraftsDocument(value);  return {    drafts: Object.fromEntries(      Object.entries(parsed.drafts)        .map(          ([key, draft]) =>            [              key,              // Stale new-task drafts left on disk by builds before the              // model-precedence fix carry a bare modelSelection with no              // other selector settings. Strip it so the next compose pass              // re-resolves project → sticky → provider defaults. Drafts              // with runtime/interaction/workspace settings or actual text /              // attachments were deliberately configured and are left alone.              key.startsWith("new-task:") &&              draft.modelSelection &&              draft.text.length === 0 &&              draft.attachments.length === 0 &&              draft.runtimeMode === undefined &&              draft.interactionMode === undefined &&              draft.workspaceSelection === undefined                ? { ...draft, modelSelection: undefined }                : draft,            ] as const,        )        // importedShareIds are share-import receipts: a contentless draft        // carrying one is not empty, or the same native share would be        // re-imported after restart.        .filter(([, draft]) => !isEmptyDraft(draft) || (draft.importedShareIds?.length ?? 0) > 0),    ),    stickyModelSelection: parsed.stickyModelSelection ?? null,  };} export function decodePersistedComposerDrafts(value: unknown): Record<string, ComposerDraft> {  return decodePersistedComposerState(value).drafts;} async function getComposerDraftsFile() {  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);} async function loadPersistedComposerState(): Promise<{  readonly drafts: Record<string, ComposerDraft>;  readonly stickyModelSelection: ModelSelection | null;}> {  let operation: ComposerDraftPersistenceError["operation"] = "open";  try {    const file = await getComposerDraftsFile();    if (!file.exists) {      return { drafts: {}, stickyModelSelection: null };    }    operation = "read";    const raw = await file.text();    operation = "decode";    return decodePersistedComposerState(JSON.parse(raw) as unknown);  } catch (cause) {    console.warn(      "[composer-drafts] ignored persisted draft failure",      new ComposerDraftPersistenceError({        operation,        directory: COMPOSER_DRAFTS_DIRECTORY,        fileName: COMPOSER_DRAFTS_FILE,        cause,      }),    );    return { drafts: {}, stickyModelSelection: null };  }} async function writePersistedComposerState(  drafts: Record<string, ComposerDraft>,  stickyModelSelection: ModelSelection | null,): Promise<void> {  let operation: ComposerDraftPersistenceError["operation"] = "open";  try {    const file = await getComposerDraftsFile();    operation = "encode";    const nonEmptyDrafts = Object.fromEntries(      Object.entries(drafts).filter(([, draft]) => !isEmptyDraft(draft)),    );    const document = {      schemaVersion: COMPOSER_DRAFTS_SCHEMA_VERSION,      drafts: nonEmptyDrafts,      ...(stickyModelSelection ? { stickyModelSelection } : {}),    } as const;    const encoded = JSON.stringify(document);    operation = "write";    await writeFileAtomically(file, encoded);  } catch (cause) {    throw new ComposerDraftPersistenceError({      operation,      directory: COMPOSER_DRAFTS_DIRECTORY,      fileName: COMPOSER_DRAFTS_FILE,      cause,    });  }} /** * Lands any debounced or in-flight draft write before the JS runtime is torn * down (app update restart), so the freshest draft state survives it. A write * failure propagates so the caller can decide whether the restart may proceed. */export async function flushComposerDrafts(): Promise<void> {  // Never land a pre-hydration snapshot: persisted state must merge into the  // atoms first, or this write would clobber disk with partial data.  ensureComposerDraftsLoaded();  if (loadPromise !== null) {    await loadPromise;  }  // An edit during an awaited write schedules another debounced write, so  // keep landing snapshots until no debounce is pending after a queue drain.  do {    while (persistTimer !== null) {      clearTimeout(persistTimer);      persistTimer = null;      await persistenceQueue.run(() =>        writePersistedComposerState(          appAtomRegistry.get(composerDraftsAtom),          appAtomRegistry.get(stickyComposerModelSelectionAtom),        ),      );    }    // Draining also waits for an already-fired debounce whose write is still    // gated behind its own hydration await inside the queue.    await persistenceQueue.run(() => Promise.resolve());  } while (persistTimer !== null);} function isComposerAttachmentFileReferenced(fileUri: string): boolean {  if (isComposerAttachmentFileRetained(fileUri)) {    return true;  }  const referenceKey = composerAttachmentFileReferenceKey(fileUri);  const drafts = Object.values(appAtomRegistry.get(composerDraftsAtom));  const queuedMessages = Object.values(    appAtomRegistry.get(threadOutboxManager.queuedMessagesByThreadKeyAtom),  ).flat();  return [...drafts, ...queuedMessages].some((owner) =>    owner.attachments.some(      (attachment) =>        attachment.type === "file" &&        composerAttachmentFileReferenceKey(attachment.fileUri) === referenceKey,    ),  );} function isComposerAttachmentUploadReferenced(  environmentId: EnvironmentId,  attachmentId: string,): boolean {  const drafts = Object.values(appAtomRegistry.get(composerDraftsAtom));  const queuedMessages = Object.values(    appAtomRegistry.get(threadOutboxManager.queuedMessagesByThreadKeyAtom),  ).flat();  return [...drafts, ...queuedMessages].some((owner) =>    owner.attachments.some(      (attachment) =>        attachment.type === "file" &&        attachment.uploadEnvironmentId === environmentId &&        attachment.uploadedAttachmentId === attachmentId,    ),  );} export async function releaseUnusedComposerAttachmentFiles(  attachments: ReadonlyArray<DraftComposerAttachment>,): Promise<void> {  const candidates = new Set(    attachments      .filter((attachment) => attachment.type === "file")      .map((attachment) => attachment.fileUri),  );  const uploadCandidates = new Map<EnvironmentId, Set<string>>();  for (const attachment of attachments) {    if (      attachment.type !== "file" ||      attachment.uploadEnvironmentId === undefined ||      attachment.uploadedAttachmentId === undefined    ) {      continue;    }    const ids = uploadCandidates.get(attachment.uploadEnvironmentId) ?? new Set<string>();    ids.add(attachment.uploadedAttachmentId);    uploadCandidates.set(attachment.uploadEnvironmentId, ids);  }  if (candidates.size === 0) {    return;  }   // Persisted drafts must hydrate before the reference scan. On a cold start  // the atom is still empty, and every file a persisted draft owns would look  // unused. Hydrate before flushing so a pending pre-hydration write cannot  // land an incomplete snapshot either.  await waitForComposerDraftsLoaded();  await flushComposerDrafts();  if (!(await threadOutboxManager.load())) {    // An unreadable outbox store must not look like an empty queue: deleting    // now would take bytes a persisted queued message still needs. Skip the    // sweep; the next one retries hydration.    return;  }  await flushThreadOutbox();   const allFilesReferenced = [...candidates].every(isComposerAttachmentFileReferenced);  const allUploadsReferenced = [...uploadCandidates].every(([environmentId, attachmentIds]) =>    [...attachmentIds].every((attachmentId) =>      isComposerAttachmentUploadReferenced(environmentId, attachmentId),    ),  );  if (allFilesReferenced && allUploadsReferenced) {    return;  }   let incomingShareFileUris: ReadonlySet<string>;  try {    const { loadIncomingShareDrafts } = await import("../features/sharing/incoming-share-storage");    const incomingShares = await loadIncomingShareDrafts({ strict: true });    incomingShareFileUris = new Set(      incomingShares.flatMap((share) =>        share.attachments.flatMap((attachment) =>          attachment.type === "file"            ? [composerAttachmentFileReferenceKey(attachment.fileUri)]            : [],        ),      ),    );  } catch (error) {    console.warn("[composer-attachments] could not verify incoming share ownership", error);    return;  }   const { removePersistedComposerAttachmentFile } = await import("../lib/composerImages");  for (const fileUri of candidates) {    // Re-check ownership immediately before each deletion: a restore or edit    // can re-own a file after an earlier scan decided it was unused.    if (      isComposerAttachmentFileReferenced(fileUri) ||      incomingShareFileUris.has(composerAttachmentFileReferenceKey(fileUri))    ) {      continue;    }    await removePersistedComposerAttachmentFile(fileUri);  }   if (uploadCandidates.size > 0) {    const { releasePendingAttachmentUploads } = await import("../lib/attachmentUpload");    for (const [environmentId, attachmentIds] of uploadCandidates) {      for (const attachmentId of attachmentIds) {        // A different draft or queued message can reuse the same pending        // upload with another local URI. Re-check the server-side ownership        // key immediately before deletion.        if (isComposerAttachmentUploadReferenced(environmentId, attachmentId)) {          continue;        }        try {          await releasePendingAttachmentUploads(environmentId, [attachmentId]);        } catch (error) {          // The server expires stale pending uploads. Local discard must still          // complete when the environment is disconnected or deletion fails.          console.warn("[composer-attachments] could not remove pending upload", {            environmentId,            attachmentId,            error,          });        }      }    }  }} export function scheduleUnusedComposerAttachmentCleanup(  attachments: ReadonlyArray<DraftComposerAttachment>,): void {  if (!attachments.some((attachment) => attachment.type === "file")) {    return;  }  void releaseUnusedComposerAttachmentFiles(attachments).catch((error) => {    console.warn("[composer-attachments] could not remove unused files", error);  });} /** Keeps previews usable after send/removal, then retries the normal ownership cleanup. */export function retainComposerAttachmentFileForPreview(  attachment: DraftComposerFileAttachment,): () => void {  return retainComposerAttachmentFile(attachment.fileUri, () => {    scheduleUnusedComposerAttachmentCleanup([attachment]);  });} function schedulePersistComposerState(): void {  if (persistTimer !== null) {    clearTimeout(persistTimer);  }  persistTimer = setTimeout(() => {    persistTimer = null;    ensureComposerDraftsLoaded();    // The write enters the serialization queue before waiting on hydration,    // so flushComposerDrafts' queue drain cannot resolve ahead of it.    void persistenceQueue.run(async () => {      if (loadPromise !== null) {        await loadPromise;      }      try {        await writePersistedComposerState(          appAtomRegistry.get(composerDraftsAtom),          appAtomRegistry.get(stickyComposerModelSelectionAtom),        );      } catch (error) {        console.warn("[composer-drafts] failed to persist drafts", error);        // Draft persistence is best-effort; in-memory drafts still keep working.      }    });  }, PERSIST_DEBOUNCE_MS);} export function ensureComposerDraftsLoaded(): void {  if (loadPromise !== null) {    return;  }  loadPromise = loadPersistedComposerState()    .then((persisted) => {      if (Object.keys(persisted.drafts).length > 0) {        const current = appAtomRegistry.get(composerDraftsAtom);        appAtomRegistry.set(composerDraftsAtom, {          ...persisted.drafts,          ...current,        });      }      if (        persisted.stickyModelSelection !== null &&        appAtomRegistry.get(stickyComposerModelSelectionAtom) === null      ) {        appAtomRegistry.set(stickyComposerModelSelectionAtom, persisted.stickyModelSelection);      }    })

Lines 1–500 of 1086.