56 lines
2.0 KiB
TypeScript
56 lines
2.0 KiB
TypeScript
import { useExternalStoreRuntime, type AppendMessage } from "@assistant-ui/react";
|
|
import { useMemo, useSyncExternalStore } from "react";
|
|
import type { MessageInput } from "@agent-studio/shared";
|
|
import { toThreadMessageLike } from "./message-mapper.js";
|
|
import type { HarnessClientState } from "../store/reducer.js";
|
|
|
|
export type HarnessRuntimeStore = {
|
|
getSnapshot(): HarnessClientState;
|
|
subscribe(subscriber: () => void): () => void;
|
|
sendMessage(input: MessageInput): Promise<unknown>;
|
|
cancel(): Promise<unknown>;
|
|
};
|
|
|
|
export function useHarnessRuntime(store: HarnessRuntimeStore, options: { isSendDisabled?: boolean } = {}) {
|
|
const snapshot = useSyncExternalStore(
|
|
(subscriber) => store.subscribe(subscriber),
|
|
() => store.getSnapshot(),
|
|
() => store.getSnapshot(),
|
|
);
|
|
const messages = useMemo(
|
|
() => snapshot.messageOrder.flatMap((id) => {
|
|
const message = snapshot.messages[id];
|
|
if (message === undefined) return [];
|
|
const tools = Object.values(snapshot.tools).filter((tool) => tool.runId === message.runId);
|
|
return [toThreadMessageLike(message, tools)];
|
|
}),
|
|
[snapshot.messageOrder, snapshot.messages, snapshot.tools],
|
|
);
|
|
const isRunning = Object.values(snapshot.runs).some((run) => run.status === "running");
|
|
|
|
return useExternalStoreRuntime({
|
|
messages,
|
|
convertMessage: (message) => message,
|
|
isRunning,
|
|
isSendDisabled: options.isSendDisabled,
|
|
onNew: async (message) => {
|
|
await store.sendMessage(extractMessageInput(message));
|
|
},
|
|
onCancel: async () => {
|
|
await store.cancel();
|
|
},
|
|
});
|
|
}
|
|
|
|
export function extractMessageInput(message: AppendMessage): MessageInput {
|
|
const text = message.content
|
|
.filter((part) => part.type === "text")
|
|
.map((part) => part.text)
|
|
.join("\n")
|
|
.trim();
|
|
const attachmentIds = (message.attachments ?? [])
|
|
.map((attachment) => attachment.id)
|
|
.filter((id) => /^att_[a-zA-Z0-9_-]+$/.test(id));
|
|
return { text, attachmentIds: [...new Set(attachmentIds)] };
|
|
}
|