1import { scopedThreadKey, scopeThreadRef } from "@t3tools/client-runtime/environment";
2import type { PendingUserInputDraftAnswer } from "@t3tools/client-runtime/state/pending-user-input";
3import {
4 onThreadRequestsResolved,
5 pruneResolvedUserInputDrafts,
6 resolvedQuestionDraftKeys,
7 userInputDraftKey,
8} from "@t3tools/client-runtime/state/thread-requests";
9import type {
10 EnvironmentId,
11 RuntimeRequestId,
12 ScopedThreadRef,
13 ThreadContextRecord,
14 ThreadId,
15} from "@t3tools/contracts";
16import { create } from "zustand";
17import { DraftId, useComposerDraftStore } from "./composerDraftStore";
18import { releaseDraftAttachments } from "./lib/attachmentUploadQueue";
19import {
20 carryDisplacedCustomAnswerIntoPrompt,
21 threadContextsReferencedByAnswer,
22} from "./pendingUserInput";
23
24/** Edited question answers, keyed by `userInputDraftKey`; suggested answers are not stored. */
25export const usePendingUserInputAnswers = create<
26 Record<string, Record<string, PendingUserInputDraftAnswer>>
27>(() => ({}));
28export const usePendingUserInputQuestionIndex = create<Record<string, number>>(() => ({}));
29
30export function questionAttachmentDraftPrefix(
31 environmentId: EnvironmentId,
32 threadId: ThreadId,
33): string {
34 return `${encodeURIComponent(JSON.stringify(environmentId))}:question-${encodeURIComponent(JSON.stringify(threadId))}-`;
35}
36
37export function questionAttachmentDraftId(
38 environmentId: EnvironmentId,
39 threadId: ThreadId,
40 requestId: string,
41 questionId: string,
42): DraftId {
43 return DraftId.make(
44 `${questionAttachmentDraftPrefix(environmentId, threadId)}${encodeURIComponent(JSON.stringify([requestId, questionId]))}`,
45 );
46}
47
48type OpenQuestionAttachmentDraft = {
49 readonly draftId: DraftId;
50 readonly requestKeys: ReadonlyArray<DraftId>;
51};
52
53const openQuestionDrafts = new Map<string, OpenQuestionAttachmentDraft>();
54
55export function openQuestionAttachmentDraft(
56 threadRef: ScopedThreadRef,
57): OpenQuestionAttachmentDraft | null {
58 return openQuestionDrafts.get(scopedThreadKey(threadRef)) ?? null;
59}
60
61export function countQuestionRequestAttachments(draftId: DraftId): number | null {
62 for (const draft of openQuestionDrafts.values()) {
63 if (draft.draftId === draftId) return countQuestionAttachments(draft.requestKeys);
64 }
65 return null;
66}
67
68export function trackOpenQuestionAttachmentDraft(
69 threadRef: ScopedThreadRef,
70 draft: OpenQuestionAttachmentDraft,
71): () => void {
72 const key = scopedThreadKey(threadRef);
73 openQuestionDrafts.set(key, draft);
74 return () => {
75 if (openQuestionDrafts.get(key) === draft) openQuestionDrafts.delete(key);
76 };
77}
78
79export const useQuestionAttachmentPreparation = create<{ counts: Record<string, number> }>(() => ({
80 counts: {},
81}));
82
83export function countQuestionAttachments(keys: ReadonlyArray<DraftId>): number {
84 const store = useComposerDraftStore.getState();
85 const { counts } = useQuestionAttachmentPreparation.getState();
86 return keys.reduce((total, key) => {
87 const draft = store.getComposerDraft(key);
88 return total + (draft?.images.length ?? 0) + (draft?.files.length ?? 0) + (counts[key] ?? 0);
89 }, 0);
90}
91
92export function changeQuestionAttachmentPreparation(key: DraftId, delta: number): void {
93 useQuestionAttachmentPreparation.setState((state) =>
94 delta < 0 && !(key in state.counts)
95 ? state
96 : {
97 counts: { ...state.counts, [key]: Math.max(0, (state.counts[key] ?? 0) + delta) },
98 },
99 );
100}
101
102export function clearQuestionAttachmentDraft(key: DraftId): void {
103 const store = useComposerDraftStore.getState();
104 const draft = store.getComposerDraft(key);
105 if (draft) {
106 releaseDraftAttachments([...draft.images, ...draft.files]);
107 for (const image of draft.images) {
108 if (image.previewUrl.startsWith("blob:")) URL.revokeObjectURL(image.previewUrl);
109 }
110 }
111 store.clearComposerContent(key);
112 useQuestionAttachmentPreparation.setState(({ counts }) => {
113 const next = { ...counts };
114 delete next[key];
115 return { counts: next };
116 });
117}
118
119/** Clears the thread's question attachment drafts for every request outside `pendingRequestIds`. */
120function pruneResolvedQuestionAttachmentDrafts(
121 environmentId: EnvironmentId,
122 threadId: ThreadId,
123 pendingRequestIds: ReadonlyArray<string>,
124): void {
125 const keys = new Set([
126 ...Object.keys(useComposerDraftStore.getState().draftsByThreadKey),
127 ...Object.keys(useQuestionAttachmentPreparation.getState().counts),
128 ]);
129 const prefix = questionAttachmentDraftPrefix(environmentId, threadId);
130 for (const key of resolvedQuestionDraftKeys(keys, prefix, pendingRequestIds)) {
131 clearQuestionAttachmentDraft(DraftId.make(key));
132 }
133}
134
135/**
136 * Stop cancels an open question and a provider timeout expires it; neither sent
137 * the typed answers, so they move to the thread's own draft (not a queued-message
138 * edit the composer may be showing) with the thread chips they reference.
139 */
140function returnUnsentAnswersToThreadDraft(
141 environmentId: EnvironmentId,
142 threadId: ThreadId,
143 unansweredRequestIds: ReadonlyArray<RuntimeRequestId>,
144): void {
145 const answers = usePendingUserInputAnswers.getState();
146 const store = useComposerDraftStore.getState();
147 const target = scopeThreadRef(environmentId, threadId);
148 const currentPrompt = store.getComposerDraft(target)?.prompt ?? "";
149 let prompt = currentPrompt;
150 const threadContexts: ThreadContextRecord[] = [];
151 for (const requestId of unansweredRequestIds) {
152 const drafts = answers[userInputDraftKey(environmentId, threadId, requestId)] ?? {};
153 for (const [questionId, { customAnswer }] of Object.entries(drafts)) {
154 const nextPrompt = carryDisplacedCustomAnswerIntoPrompt(prompt, customAnswer);
155 if (nextPrompt === prompt) continue;
156 prompt = nextPrompt;
157 const questionDraft = store.getComposerDraft(
158 questionAttachmentDraftId(environmentId, threadId, requestId, questionId),
159 );
160 threadContexts.push(
161 ...threadContextsReferencedByAnswer(
162 customAnswer ?? "",
163 questionDraft?.threadContexts ?? [],
164 ),
165 );
166 }
167 }
168 if (prompt === currentPrompt) return;
169 store.setPrompt(target, prompt);
170 store.addThreadContexts(target, threadContexts, { appendReference: false });
171}
172
173// Resolution can come from anywhere: another device, an interrupt, the turn
174// ending, or a reload over drafts for requests resolved meanwhile. One listener,
175// so unsent answers and their chips are read before their drafts are pruned.
176onThreadRequestsResolved(({ environmentId, threadId, pendingRequestIds, unansweredRequestIds }) => {
177 returnUnsentAnswersToThreadDraft(environmentId, threadId, unansweredRequestIds);
178 usePendingUserInputAnswers.setState(
179 (existing) =>
180 pruneResolvedUserInputDrafts(existing, environmentId, threadId, pendingRequestIds),
181 true,
182 );
183 usePendingUserInputQuestionIndex.setState(
184 (existing) =>
185 pruneResolvedUserInputDrafts(existing, environmentId, threadId, pendingRequestIds),
186 true,
187 );
188 pruneResolvedQuestionAttachmentDrafts(environmentId, threadId, pendingRequestIds);
189});