fix: harden assistant authoring and UI
This commit is contained in:
ajnart authored and 
Thomas Camlong committed
1 parent
fa0e55cc93
commit
edc7b4d55b
51 files changed
+4118
-916
No files matched your search
@@ -1,8 +1,12 @@
|
||||
"use client";
|
||||
|
||||
import type { PropsWithChildren } from "react";
|
||||
import type { AssistantRuntime, Toolkit } from "@assistant-ui/react";
|
||||
import { AssistantRuntimeProvider, Tools, useAui } from "@assistant-ui/react";
|
||||
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;
|
||||
@@ -21,3 +25,501 @@ export const AssistantRuntimeProviderWithTools = ({
|
||||
</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(),
|
||||
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>
|
||||
);
|
||||
};
|
||||
Reference in new issue
Block a user