527 lines
20 KiB
TypeScript
527 lines
20 KiB
TypeScript
"use client";
|
|
|
|
import type { ClipboardEvent, PropsWithChildren } from "react";
|
|
import { createContext, useContext, useEffect, useId, useLayoutEffect, useMemo, useRef } from "react";
|
|
import type { AssistantRuntime, AttachmentAdapter, ThreadListRuntime, Toolkit } from "@assistant-ui/react";
|
|
import { AssistantRuntimeProvider, INTERNAL, Tools, useAui, useAuiEvent } from "@assistant-ui/react";
|
|
|
|
import { showErrorNotification } from "@homarr/notifications";
|
|
import { useScopedI18n } from "@homarr/translation/client";
|
|
|
|
interface AssistantRuntimeProviderWithToolsProps extends PropsWithChildren {
|
|
runtime: AssistantRuntime;
|
|
toolkit: Toolkit;
|
|
}
|
|
|
|
export const AssistantRuntimeProviderWithTools = ({
|
|
children,
|
|
runtime,
|
|
toolkit,
|
|
}: AssistantRuntimeProviderWithToolsProps) => {
|
|
const toolsAui = useAui({ tools: Tools({ toolkit }) });
|
|
return (
|
|
<AssistantRuntimeProvider aui={toolsAui} runtime={runtime}>
|
|
{children}
|
|
</AssistantRuntimeProvider>
|
|
);
|
|
};
|
|
|
|
type ThreadCore = INTERNAL.ThreadRuntimeCore;
|
|
type ThreadCoreWithAdapters = ConstructorParameters<typeof INTERNAL.DefaultThreadComposerRuntimeCore>[0];
|
|
type ComposerConnectionDisposer = () => void;
|
|
|
|
export const assistantSurfaceComposerCacheLimit = 12;
|
|
const assistantSurfaceAttachmentLimit = 5;
|
|
const composerConnectionDisposers = new WeakMap<object, ComposerConnectionDisposer>();
|
|
|
|
class AssistantSurfaceComposerRuntimeCore extends INTERNAL.DefaultThreadComposerRuntimeCore {
|
|
public override connect() {
|
|
const dispose = super.connect();
|
|
composerConnectionDisposers.set(this, dispose);
|
|
return dispose;
|
|
}
|
|
|
|
public dispose() {
|
|
composerConnectionDisposers.get(this)?.();
|
|
composerConnectionDisposers.delete(this);
|
|
}
|
|
}
|
|
|
|
interface AssistantSurfaceRuntime extends AssistantRuntime {
|
|
dispose(): void;
|
|
}
|
|
|
|
interface AssistantSurfaceComposerEntry {
|
|
composer: AssistantSurfaceComposerRuntimeCore;
|
|
source: ThreadCore;
|
|
}
|
|
|
|
const getThreadAdapters = (threadCore: ThreadCore) => (threadCore as ThreadCoreWithAdapters).adapters;
|
|
|
|
const bindRuntimeValue = (target: object, property: PropertyKey, cache: WeakMap<object, Map<PropertyKey, unknown>>) => {
|
|
const value = Reflect.get(target, property, target) as unknown;
|
|
if (typeof value !== "function") return value;
|
|
|
|
let targetCache = cache.get(target);
|
|
if (!targetCache) {
|
|
targetCache = new Map();
|
|
cache.set(target, targetCache);
|
|
}
|
|
const cached = targetCache.get(property);
|
|
if (cached) return cached;
|
|
const bound = value.bind(target) as unknown;
|
|
targetCache.set(property, bound);
|
|
return bound;
|
|
};
|
|
|
|
const createSurfaceAttachmentAdapter = (threadCore: ThreadCore): AttachmentAdapter => {
|
|
const pendingAttachmentIds = new Set<string>();
|
|
let pendingAdds = 0;
|
|
const getAdapter = () => {
|
|
const adapter = getThreadAdapters(threadCore)?.attachments;
|
|
if (!adapter) throw new Error("Attachments are not supported");
|
|
return adapter;
|
|
};
|
|
|
|
return {
|
|
get accept() {
|
|
return getThreadAdapters(threadCore)?.attachments?.accept ?? "";
|
|
},
|
|
add({ file }) {
|
|
if (pendingAttachmentIds.size + pendingAdds >= assistantSurfaceAttachmentLimit) {
|
|
throw new Error(`A message can include up to ${assistantSurfaceAttachmentLimit} attachments.`);
|
|
}
|
|
|
|
pendingAdds += 1;
|
|
let result: ReturnType<AttachmentAdapter["add"]>;
|
|
try {
|
|
result = getAdapter().add({ file });
|
|
} catch (error) {
|
|
pendingAdds -= 1;
|
|
throw error;
|
|
}
|
|
|
|
if (Symbol.asyncIterator in result) {
|
|
return (async function* () {
|
|
let reservationPending = true;
|
|
try {
|
|
for await (const attachment of result) {
|
|
if (reservationPending) {
|
|
reservationPending = false;
|
|
pendingAdds -= 1;
|
|
}
|
|
pendingAttachmentIds.add(attachment.id);
|
|
yield attachment;
|
|
}
|
|
} finally {
|
|
if (reservationPending) pendingAdds -= 1;
|
|
}
|
|
})();
|
|
}
|
|
|
|
return result.then(
|
|
(attachment) => {
|
|
pendingAdds -= 1;
|
|
pendingAttachmentIds.add(attachment.id);
|
|
return attachment;
|
|
},
|
|
(error: unknown) => {
|
|
pendingAdds -= 1;
|
|
throw error;
|
|
},
|
|
);
|
|
},
|
|
async remove(attachment) {
|
|
try {
|
|
await getAdapter().remove(attachment);
|
|
} finally {
|
|
pendingAttachmentIds.delete(attachment.id);
|
|
}
|
|
},
|
|
async send(attachment) {
|
|
try {
|
|
return await getAdapter().send(attachment);
|
|
} finally {
|
|
pendingAttachmentIds.delete(attachment.id);
|
|
}
|
|
},
|
|
};
|
|
};
|
|
|
|
const createSurfaceComposerHost = (
|
|
threadCore: ThreadCore,
|
|
boundRuntimeValues: WeakMap<object, Map<PropertyKey, unknown>>,
|
|
) => {
|
|
const attachments = createSurfaceAttachmentAdapter(threadCore);
|
|
const adapters = {
|
|
get attachments() {
|
|
return getThreadAdapters(threadCore)?.attachments ? attachments : undefined;
|
|
},
|
|
get dictation() {
|
|
return getThreadAdapters(threadCore)?.dictation;
|
|
},
|
|
};
|
|
|
|
return new Proxy(threadCore as ThreadCoreWithAdapters, {
|
|
get(target, property) {
|
|
if (property === "adapters") return adapters;
|
|
return bindRuntimeValue(target, property, boundRuntimeValues);
|
|
},
|
|
});
|
|
};
|
|
|
|
/**
|
|
* assistant-ui intentionally stores the thread composer on the thread runtime. That is the right
|
|
* default for a single chat surface, but Homarr can render the same conversation in the floating
|
|
* panel and in one or more board widgets at the same time. A second runtime-bound Lexical input
|
|
* would otherwise mirror the first input's draft, attachments and selection.
|
|
*
|
|
* This adapter keeps the thread runtime itself shared and replaces only its composer core. The
|
|
* local composer still delegates send/cancel to the shared thread, so messages, runs, tools,
|
|
* branches and history remain one conversation. Composers are cached per conversation id because
|
|
* some assistant-ui runtimes reuse the same underlying core while switching conversations.
|
|
*/
|
|
export const createAssistantSurfaceRuntime = (runtime: AssistantRuntime): AssistantSurfaceRuntime => {
|
|
const sharedThread = runtime.thread as InstanceType<typeof INTERNAL.ThreadRuntimeImpl>;
|
|
// oxlint-disable-next-line no-underscore-dangle -- assistant-ui exposes the binding only through its internal adapter API.
|
|
const sharedBinding = sharedThread.__internal_threadBinding;
|
|
const composerByConversation = new Map<string, AssistantSurfaceComposerEntry>();
|
|
const surfaceCoreByConversation = new Map<string, { source: ThreadCore; surface: ThreadCore }>();
|
|
const activeConversationCounts = new Map<string, number>();
|
|
const boundRuntimeValues = new WeakMap<object, Map<PropertyKey, unknown>>();
|
|
|
|
const disposeComposer = (conversationId: string, entry: AssistantSurfaceComposerEntry) => {
|
|
const reset = entry.composer.reset();
|
|
entry.composer.dispose();
|
|
composerByConversation.delete(conversationId);
|
|
surfaceCoreByConversation.delete(conversationId);
|
|
void reset.catch(() => {});
|
|
};
|
|
|
|
const evictInactiveComposers = () => {
|
|
while (composerByConversation.size > assistantSurfaceComposerCacheLimit) {
|
|
let candidate: [string, AssistantSurfaceComposerEntry] | undefined;
|
|
for (const entry of composerByConversation) {
|
|
if (!activeConversationCounts.has(entry[0])) {
|
|
candidate = entry;
|
|
break;
|
|
}
|
|
}
|
|
if (!candidate) return;
|
|
disposeComposer(candidate[0], candidate[1]);
|
|
}
|
|
};
|
|
|
|
const retainConversation = (conversationId: string) => {
|
|
activeConversationCounts.set(conversationId, (activeConversationCounts.get(conversationId) ?? 0) + 1);
|
|
};
|
|
|
|
const releaseConversation = (conversationId: string) => {
|
|
const count = activeConversationCounts.get(conversationId);
|
|
if (count === undefined) return;
|
|
if (count <= 1) activeConversationCounts.delete(conversationId);
|
|
else activeConversationCounts.set(conversationId, count - 1);
|
|
evictInactiveComposers();
|
|
};
|
|
|
|
const getComposer = (conversationId: string, threadCore: ThreadCore) => {
|
|
const existing = composerByConversation.get(conversationId);
|
|
if (existing?.source === threadCore) {
|
|
composerByConversation.delete(conversationId);
|
|
composerByConversation.set(conversationId, existing);
|
|
return existing.composer;
|
|
}
|
|
const draft = existing
|
|
? {
|
|
attachments: existing.composer.attachments,
|
|
quote: existing.composer.quote,
|
|
role: existing.composer.role,
|
|
runConfig: existing.composer.runConfig,
|
|
text: existing.composer.text,
|
|
}
|
|
: undefined;
|
|
if (existing) disposeComposer(conversationId, existing);
|
|
|
|
const host = createSurfaceComposerHost(threadCore, boundRuntimeValues);
|
|
const composer = new AssistantSurfaceComposerRuntimeCore(host);
|
|
composerByConversation.set(conversationId, { composer, source: threadCore });
|
|
if (draft) {
|
|
composer.setText(draft.text);
|
|
composer.setRole(draft.role);
|
|
composer.setRunConfig(draft.runConfig);
|
|
composer.setQuote(draft.quote);
|
|
for (const attachment of draft.attachments) {
|
|
if (attachment.status.type === "complete" && attachment.content) {
|
|
void composer
|
|
.addAttachment({
|
|
id: attachment.id,
|
|
type: attachment.type,
|
|
name: attachment.name,
|
|
contentType: attachment.contentType,
|
|
content: attachment.content,
|
|
})
|
|
.catch(() => {});
|
|
} else if (attachment.file) {
|
|
void composer.addAttachment(attachment.file).catch(() => {});
|
|
}
|
|
}
|
|
}
|
|
evictInactiveComposers();
|
|
return composer;
|
|
};
|
|
|
|
const getSurfaceCore = (conversationId: string, threadCore: ThreadCore) => {
|
|
const existing = surfaceCoreByConversation.get(conversationId);
|
|
if (existing?.source === threadCore) return existing.surface;
|
|
|
|
const composer = getComposer(conversationId, threadCore);
|
|
const surfaceCore = new Proxy(threadCore, {
|
|
get(target, property) {
|
|
if (property === "composer") return composer;
|
|
return bindRuntimeValue(target, property, boundRuntimeValues);
|
|
},
|
|
});
|
|
surfaceCoreByConversation.set(conversationId, { source: threadCore, surface: surfaceCore });
|
|
return surfaceCore;
|
|
};
|
|
|
|
const getConversationId = () => runtime.threads.mainItem.getState().id;
|
|
|
|
const surfaceBinding: INTERNAL.ThreadRuntimeCoreBinding = {
|
|
path: sharedBinding.path,
|
|
getState: () => getSurfaceCore(getConversationId(), sharedBinding.getState()),
|
|
outerSubscribe: (callback) => sharedBinding.outerSubscribe(callback),
|
|
subscribe(callback) {
|
|
let activeCore = sharedBinding.getState();
|
|
let activeConversationId = getConversationId();
|
|
retainConversation(activeConversationId);
|
|
let unsubscribeComposer = getComposer(activeConversationId, activeCore).subscribe(callback);
|
|
const handleSharedUpdate = () => {
|
|
const nextCore = sharedBinding.getState();
|
|
const nextConversationId = getConversationId();
|
|
if (nextCore !== activeCore || nextConversationId !== activeConversationId) {
|
|
unsubscribeComposer();
|
|
releaseConversation(activeConversationId);
|
|
activeCore = nextCore;
|
|
activeConversationId = nextConversationId;
|
|
retainConversation(activeConversationId);
|
|
unsubscribeComposer = getComposer(activeConversationId, activeCore).subscribe(callback);
|
|
}
|
|
callback();
|
|
};
|
|
const unsubscribeThread = sharedBinding.subscribe(handleSharedUpdate);
|
|
const unsubscribeThreadItem = runtime.threads.mainItem.subscribe(handleSharedUpdate);
|
|
return () => {
|
|
unsubscribeComposer();
|
|
releaseConversation(activeConversationId);
|
|
unsubscribeThread();
|
|
unsubscribeThreadItem();
|
|
};
|
|
},
|
|
};
|
|
|
|
const sharedMainItem = runtime.threads.mainItem;
|
|
const mainItemBinding: INTERNAL.ThreadListItemRuntimeBinding = {
|
|
path: sharedMainItem.path,
|
|
getState: () => sharedMainItem.getState(),
|
|
subscribe: (callback) => sharedMainItem.subscribe(callback),
|
|
};
|
|
const surfaceThread = new INTERNAL.ThreadRuntimeImpl(surfaceBinding, mainItemBinding);
|
|
const sharedThreads = runtime.threads;
|
|
const getSurfaceThreadById = (threadId: string) => {
|
|
const requestedItem = sharedThreads.getItemById(threadId).getState();
|
|
if (requestedItem.id === sharedMainItem.getState().id) return surfaceThread;
|
|
// Homarr renders only the selected conversation. Preserve the upstream behavior for inactive
|
|
// thread inspection without manufacturing another composer outside this surface's LRU cache.
|
|
return sharedThreads.getById(threadId);
|
|
};
|
|
const surfaceThreads: ThreadListRuntime = {
|
|
main: surfaceThread,
|
|
mainItem: sharedMainItem,
|
|
getState: () => sharedThreads.getState(),
|
|
subscribe: (callback) => sharedThreads.subscribe(callback),
|
|
getById: getSurfaceThreadById,
|
|
getItemById: (threadId) => sharedThreads.getItemById(threadId),
|
|
getItemByIndex: (index) => sharedThreads.getItemByIndex(index),
|
|
getArchivedItemByIndex: (index) => sharedThreads.getArchivedItemByIndex(index),
|
|
switchToThread: (threadId, options) => sharedThreads.switchToThread(threadId, options),
|
|
switchToNewThread: () => sharedThreads.switchToNewThread(),
|
|
getLoadThreadsPromise: () => sharedThreads.getLoadThreadsPromise(),
|
|
reload: () => sharedThreads.reload(),
|
|
reloadMainThread: () => sharedThreads.reloadMainThread(),
|
|
loadMore: () => sharedThreads.loadMore(),
|
|
};
|
|
|
|
return {
|
|
threads: surfaceThreads,
|
|
thread: surfaceThread,
|
|
// The root runtime already owns the inherited model-context client. Re-registering that same
|
|
// provider for every surface would make one unmount remove it from the shared runtime's Set.
|
|
registerModelContextProvider: () => () => {},
|
|
dispose() {
|
|
activeConversationCounts.clear();
|
|
for (const [conversationId, entry] of composerByConversation) disposeComposer(conversationId, entry);
|
|
},
|
|
};
|
|
};
|
|
|
|
interface AssistantComposerSurfaceValue {
|
|
id: string;
|
|
}
|
|
|
|
const AssistantComposerSurfaceContext = createContext<AssistantComposerSurfaceValue | null>(null);
|
|
|
|
export const useAssistantComposerSurface = () => {
|
|
const value = useContext(AssistantComposerSurfaceContext);
|
|
if (!value) throw new Error("useAssistantComposerSurface must be used within AssistantComposerSurfaceProvider");
|
|
return value;
|
|
};
|
|
|
|
const surfaceElements = new Map<string, HTMLElement>();
|
|
|
|
const getSurfaceIdForElement = (element: Element | null) => {
|
|
if (!element) return null;
|
|
for (const [surfaceId, surfaceElement] of surfaceElements) {
|
|
if (surfaceElement.contains(element)) return surfaceId;
|
|
}
|
|
return null;
|
|
};
|
|
|
|
const getClipboardAttachmentName = (file: File, index: number) => {
|
|
if (file.name.trim().length > 0) return file.name;
|
|
const extension = file.type.split("/")[1]?.replace("jpeg", "jpg") ?? "bin";
|
|
return `pasted-file-${index + 1}.${extension}`;
|
|
};
|
|
|
|
const AssistantComposerSurfaceEvents = () => {
|
|
const t = useScopedI18n("common.assistant");
|
|
useAuiEvent("composer.attachmentAddError", ({ message }) => {
|
|
showErrorNotification({ title: t("attachments.errorTitle"), message });
|
|
});
|
|
return null;
|
|
};
|
|
|
|
const AssistantComposerSurfaceBoundary = ({ children, surfaceId }: PropsWithChildren<{ surfaceId: string }>) => {
|
|
const t = useScopedI18n("common.assistant");
|
|
const aui = useAui();
|
|
const elementRef = useRef<HTMLDivElement>(null);
|
|
const contextValue = useMemo(() => ({ id: surfaceId }), [surfaceId]);
|
|
|
|
useLayoutEffect(() => {
|
|
const element = elementRef.current;
|
|
if (!element) return undefined;
|
|
surfaceElements.set(surfaceId, element);
|
|
|
|
return () => {
|
|
if (surfaceElements.get(surfaceId) === element) surfaceElements.delete(surfaceId);
|
|
};
|
|
}, [surfaceId]);
|
|
|
|
const handlePasteCapture = (event: ClipboardEvent<HTMLDivElement>) => {
|
|
const target = event.target;
|
|
if (!(target instanceof HTMLElement) || !target.closest("[data-assistant-composer-input]")) return;
|
|
const files = Array.from(event.clipboardData.files);
|
|
if (files.length === 0) return;
|
|
|
|
event.preventDefault();
|
|
const composer = aui.composer();
|
|
const acceptedTypes = new Set(
|
|
composer
|
|
.getState()
|
|
.attachmentAccept.split(",")
|
|
.map((type) => type.trim())
|
|
.filter(Boolean),
|
|
);
|
|
|
|
files.forEach((file, index) => {
|
|
if (
|
|
file.type.startsWith("image/") &&
|
|
!acceptedTypes.has("*") &&
|
|
!acceptedTypes.has(file.type) &&
|
|
!acceptedTypes.has("image/*")
|
|
) {
|
|
showErrorNotification({
|
|
title: t("attachments.errorTitle"),
|
|
message: t("attachments.imageUnsupported"),
|
|
});
|
|
return;
|
|
}
|
|
const attachment =
|
|
file.name.trim().length > 0
|
|
? file
|
|
: new File([file], getClipboardAttachmentName(file, index), {
|
|
type: file.type,
|
|
lastModified: file.lastModified,
|
|
});
|
|
void composer.addAttachment(attachment).catch(() => {
|
|
// The scoped composer event above presents the adapter's actionable error message.
|
|
});
|
|
});
|
|
};
|
|
|
|
return (
|
|
<AssistantComposerSurfaceContext.Provider value={contextValue}>
|
|
<div ref={elementRef} style={{ display: "contents" }} onPasteCapture={handlePasteCapture}>
|
|
<AssistantComposerSurfaceEvents />
|
|
{children}
|
|
</div>
|
|
</AssistantComposerSurfaceContext.Provider>
|
|
);
|
|
};
|
|
|
|
/**
|
|
* The current react-lexical FocusPlugin focuses every mounted composer when a run starts. Mount
|
|
* this once above all surfaces so programmatic focus returns to the element/surface the user was
|
|
* actually interacting with after those listeners have run.
|
|
*/
|
|
export const AssistantRunFocusPreserver = () => {
|
|
useAuiEvent("thread.runStart", () => {
|
|
const focusedBeforeRun = document.activeElement instanceof HTMLElement ? document.activeElement : null;
|
|
const surfaceIdBeforeRun = getSurfaceIdForElement(focusedBeforeRun);
|
|
|
|
const restoreFocus = () => {
|
|
const surfaceIdAfterRun = getSurfaceIdForElement(document.activeElement);
|
|
if (!surfaceIdAfterRun || surfaceIdAfterRun === surfaceIdBeforeRun) return;
|
|
|
|
if (focusedBeforeRun?.isConnected) focusedBeforeRun.focus({ preventScroll: true });
|
|
if (
|
|
getSurfaceIdForElement(document.activeElement) !== surfaceIdBeforeRun &&
|
|
document.activeElement instanceof HTMLElement
|
|
) {
|
|
document.activeElement.blur();
|
|
}
|
|
};
|
|
|
|
queueMicrotask(() => queueMicrotask(restoreFocus));
|
|
requestAnimationFrame(restoreFocus);
|
|
});
|
|
return null;
|
|
};
|
|
|
|
interface AssistantComposerSurfaceProviderProps extends PropsWithChildren {
|
|
surfaceId: string;
|
|
}
|
|
|
|
/**
|
|
* Gives one rendered assistant surface a local composer while preserving the shared conversation.
|
|
* Every component below this boundary can keep using the normal assistant-ui primitives and hooks.
|
|
*/
|
|
export const AssistantComposerSurfaceProvider = ({ children, surfaceId }: AssistantComposerSurfaceProviderProps) => {
|
|
const parentAui = useAui();
|
|
const instanceId = useId();
|
|
const resolvedSurfaceId = `${surfaceId}:${instanceId}`;
|
|
// oxlint-disable-next-line no-underscore-dangle -- required to derive a scoped runtime without cloning the shared thread.
|
|
const sharedRuntime = parentAui.threads().__internal_getAssistantRuntime?.();
|
|
const surfaceRuntime = useMemo(() => {
|
|
if (!sharedRuntime) throw new Error("Assistant runtime is not available for this composer surface");
|
|
return createAssistantSurfaceRuntime(sharedRuntime);
|
|
}, [sharedRuntime]);
|
|
useEffect(() => () => surfaceRuntime.dispose(), [surfaceRuntime]);
|
|
|
|
return (
|
|
<AssistantRuntimeProvider aui={parentAui} runtime={surfaceRuntime}>
|
|
<AssistantComposerSurfaceBoundary surfaceId={resolvedSurfaceId}>{children}</AssistantComposerSurfaceBoundary>
|
|
</AssistantRuntimeProvider>
|
|
);
|
|
};
|