129 lines
3.5 KiB
TypeScript
129 lines
3.5 KiB
TypeScript
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
|
|
import { useMemo } from "react";
|
|
|
|
import { useThread } from "@/components/workspace/messages/context";
|
|
|
|
import { loadArtifactContent, loadArtifactContentFromToolCall } from "./loader";
|
|
import { updateArtifactContent } from "./api";
|
|
|
|
function hasToolResult(thread: ReturnType<typeof useThread>["thread"], toolCallId?: string) {
|
|
if (!toolCallId) {
|
|
return false;
|
|
}
|
|
return thread.messages.some(
|
|
(message) => message.type === "tool" && message.tool_call_id === toolCallId,
|
|
);
|
|
}
|
|
|
|
function applyPendingTextEdit({
|
|
content,
|
|
filepath,
|
|
thread,
|
|
}: {
|
|
content: string | undefined;
|
|
filepath: string;
|
|
thread: ReturnType<typeof useThread>["thread"];
|
|
}) {
|
|
if (!content || !thread.isLoading) {
|
|
return null;
|
|
}
|
|
|
|
for (const message of [...thread.messages].reverse()) {
|
|
if (message.type !== "ai" || !message.tool_calls) {
|
|
continue;
|
|
}
|
|
for (const toolCall of [...message.tool_calls].reverse()) {
|
|
if (toolCall.name !== "str_replace" || hasToolResult(thread, toolCall.id)) {
|
|
continue;
|
|
}
|
|
const args = toolCall.args as {
|
|
path?: unknown;
|
|
old_str?: unknown;
|
|
new_str?: unknown;
|
|
replace_all?: unknown;
|
|
};
|
|
if (args.path !== filepath) {
|
|
continue;
|
|
}
|
|
if (typeof args.old_str !== "string" || typeof args.new_str !== "string") {
|
|
continue;
|
|
}
|
|
if (!content.includes(args.old_str)) {
|
|
return { content, kind: "editing" as const };
|
|
}
|
|
const nextContent =
|
|
args.replace_all === true
|
|
? content.split(args.old_str).join(args.new_str)
|
|
: content.replace(args.old_str, args.new_str);
|
|
return { content: nextContent, kind: "editing" as const };
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
export function useArtifactContent({
|
|
filepath,
|
|
threadId,
|
|
enabled,
|
|
}: {
|
|
filepath: string;
|
|
threadId: string;
|
|
enabled?: boolean;
|
|
}) {
|
|
const isWriteFile = useMemo(() => {
|
|
return filepath.startsWith("write-file:");
|
|
}, [filepath]);
|
|
const { thread, isMock } = useThread();
|
|
const content = useMemo(() => {
|
|
if (isWriteFile) {
|
|
return loadArtifactContentFromToolCall({ url: filepath, thread });
|
|
}
|
|
return null;
|
|
}, [filepath, isWriteFile, thread]);
|
|
|
|
const { data, isLoading, error } = useQuery({
|
|
queryKey: ["artifact", filepath, threadId, isMock],
|
|
queryFn: () => {
|
|
return loadArtifactContent({ filepath, threadId, isMock });
|
|
},
|
|
enabled,
|
|
refetchInterval: thread.isLoading ? 1000 : false,
|
|
// Cache artifact content for 5 minutes to avoid repeated fetches (especially for .skill ZIP extraction)
|
|
staleTime: 5 * 60 * 1000,
|
|
});
|
|
const pendingEdit = useMemo(() => {
|
|
if (isWriteFile) {
|
|
return null;
|
|
}
|
|
return applyPendingTextEdit({
|
|
content: data?.content,
|
|
filepath,
|
|
thread,
|
|
});
|
|
}, [data?.content, filepath, isWriteFile, thread]);
|
|
|
|
return {
|
|
content: isWriteFile ? content : (pendingEdit?.content ?? data?.content),
|
|
url: isWriteFile ? undefined : data?.url,
|
|
isLoading,
|
|
isApplyingEdit: pendingEdit?.kind === "editing",
|
|
error,
|
|
};
|
|
}
|
|
|
|
export function useUpdateArtifactContent() {
|
|
const queryClient = useQueryClient();
|
|
return useMutation({
|
|
mutationFn: updateArtifactContent,
|
|
onSuccess: (content, variables) => {
|
|
queryClient.setQueryData(
|
|
["artifact", variables.filepath, variables.threadId, variables.isMock],
|
|
{
|
|
content,
|
|
url: undefined,
|
|
},
|
|
);
|
|
},
|
|
});
|
|
}
|