ea5802c380
Co-authored-by: Douwe Osinga <douwe@squareup.com>
488 lines
15 KiB
TypeScript
488 lines
15 KiB
TypeScript
import { useCallback, useRef } from "react";
|
|
import { useChatStore } from "../stores/chatStore";
|
|
import { useChatSessionStore } from "../stores/chatSessionStore";
|
|
import { clearReplayBuffer, getAndDeleteReplayBuffer } from "./replayBuffer";
|
|
import {
|
|
type ChatAttachmentDraft,
|
|
type Message,
|
|
createSystemNotificationMessage,
|
|
createUserMessage,
|
|
} from "@/shared/types/messages";
|
|
import type { ChatState, TokenState } from "@/shared/types/chat";
|
|
import { INITIAL_SESSION_CHAT_RUNTIME } from "@/shared/types/chat";
|
|
import {
|
|
acpSendMessage,
|
|
acpCancelSession,
|
|
acpLoadSession,
|
|
} from "@/shared/api/acp";
|
|
import { useAgentStore } from "@/features/agents/stores/agentStore";
|
|
import {
|
|
getSessionTitleFromDraft,
|
|
isDefaultChatTitle,
|
|
} from "../lib/sessionTitle";
|
|
import { perfLog } from "@/shared/lib/perfLog";
|
|
import {
|
|
appendAttachmentPaths,
|
|
buildAcpImages,
|
|
buildMessageAttachments,
|
|
} from "../lib/attachments";
|
|
import { i18n } from "@/shared/i18n";
|
|
import type { ChatSendOptions } from "../types";
|
|
|
|
// TODO: Remove this fallback once goose2 has first-class /-commands.
|
|
const MANUAL_COMPACT_TRIGGER = "/compact";
|
|
const EMPTY_MESSAGES: Message[] = [];
|
|
type CompactConversationResult = "completed" | "failed" | "skipped";
|
|
type EnsurePrepared = (personaId?: string) => Promise<boolean | undefined>;
|
|
|
|
function createCompactionConfirmationMessage() {
|
|
return createSystemNotificationMessage(
|
|
i18n.t("chat:notifications.compactionComplete"),
|
|
"compaction",
|
|
);
|
|
}
|
|
|
|
function getErrorMessage(error: unknown): string {
|
|
// Tauri command rejections typically arrive as plain strings, so handle
|
|
// that shape first before falling back to standard Error objects.
|
|
if (typeof error === "string" && error.trim()) {
|
|
return error;
|
|
}
|
|
|
|
if (error instanceof Error && error.message.trim()) {
|
|
return error.message;
|
|
}
|
|
|
|
if (
|
|
typeof error === "object" &&
|
|
error !== null &&
|
|
"message" in error &&
|
|
typeof error.message === "string" &&
|
|
error.message.trim()
|
|
) {
|
|
return error.message;
|
|
}
|
|
|
|
return "Unknown error";
|
|
}
|
|
|
|
async function ensurePreparedForPrompt(
|
|
ensurePrepared: EnsurePrepared | undefined,
|
|
personaId?: string,
|
|
) {
|
|
const prepared = await ensurePrepared?.(personaId);
|
|
if (prepared === false) {
|
|
throw new Error(i18n.t("chat:errors.sessionPreparationSuperseded"));
|
|
}
|
|
}
|
|
|
|
function markMessageStopped(sessionId: string, messageId: string) {
|
|
useChatStore.getState().updateMessage(sessionId, messageId, (message) => {
|
|
if (
|
|
message.metadata?.completionStatus === "completed" ||
|
|
message.metadata?.completionStatus === "error" ||
|
|
message.metadata?.completionStatus === "stopped"
|
|
) {
|
|
return message;
|
|
}
|
|
|
|
return {
|
|
...message,
|
|
metadata: {
|
|
...message.metadata,
|
|
completionStatus: "stopped",
|
|
},
|
|
content: message.content.map((block) =>
|
|
block.type === "toolRequest" && block.status === "in_progress"
|
|
? { ...block, status: "stopped" }
|
|
: block,
|
|
),
|
|
};
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Hook for managing a chat session -- sending messages, handling streaming,
|
|
* and managing chat lifecycle.
|
|
*/
|
|
export function useChat(
|
|
sessionId: string,
|
|
providerOverride?: string,
|
|
systemPromptOverride?: string,
|
|
personaInfo?: { id: string; name: string },
|
|
options?: {
|
|
onMessageAccepted?: (sessionId: string) => void;
|
|
ensurePrepared?: EnsurePrepared;
|
|
},
|
|
) {
|
|
const abortRef = useRef<AbortController | null>(null);
|
|
|
|
const messages = useChatStore(
|
|
(s) => s.messagesBySession[sessionId] ?? EMPTY_MESSAGES,
|
|
);
|
|
const runtime = useChatStore(
|
|
(s) => s.sessionStateById[sessionId] ?? INITIAL_SESSION_CHAT_RUNTIME,
|
|
);
|
|
const setActiveSession = useChatStore((s) => s.setActiveSession);
|
|
const addMessage = useChatStore((s) => s.addMessage);
|
|
const setMessages = useChatStore((s) => s.setMessages);
|
|
const clearMessages = useChatStore((s) => s.clearMessages);
|
|
const setChatState = useChatStore((s) => s.setChatState);
|
|
const setError = useChatStore((s) => s.setError);
|
|
const setStreamingMessageId = useChatStore((s) => s.setStreamingMessageId);
|
|
const setPendingAssistantProvider = useChatStore(
|
|
(s) => s.setPendingAssistantProvider,
|
|
);
|
|
const clearDraft = useChatStore((s) => s.clearDraft);
|
|
const setSessionLoading = useChatStore((s) => s.setSessionLoading);
|
|
|
|
const { chatState, tokenState, error, streamingMessageId } = runtime;
|
|
const isStreaming = chatState === "streaming" || streamingMessageId !== null;
|
|
|
|
const resolvePersonaInfo = useCallback(
|
|
(overridePersonaId?: string, overridePersonaName?: string) => {
|
|
if (overridePersonaId) {
|
|
// Read the latest persona snapshot at call time so override lookups
|
|
// still work even if the agent store changed after this hook rendered.
|
|
const personaName =
|
|
overridePersonaName ??
|
|
useAgentStore.getState().getPersonaById(overridePersonaId)
|
|
?.displayName ??
|
|
overridePersonaId;
|
|
return { id: overridePersonaId, name: personaName };
|
|
}
|
|
|
|
return personaInfo;
|
|
},
|
|
[personaInfo],
|
|
);
|
|
|
|
const sendMessage = useCallback(
|
|
async (
|
|
text: string,
|
|
overridePersona?: { id: string; name?: string },
|
|
attachments?: ChatAttachmentDraft[],
|
|
sendOptions?: ChatSendOptions,
|
|
) => {
|
|
const sid = sessionId.slice(0, 8);
|
|
const tSendStart = performance.now();
|
|
const images = buildAcpImages(attachments);
|
|
const hasAttachments = (attachments?.length ?? 0) > 0;
|
|
const hasAssistantPrompt = Boolean(sendOptions?.assistantPrompt?.trim());
|
|
const currentChatState = useChatStore
|
|
.getState()
|
|
.getSessionRuntime(sessionId).chatState;
|
|
if (
|
|
(!text.trim() && !hasAttachments && !hasAssistantPrompt) ||
|
|
currentChatState === "streaming" ||
|
|
currentChatState === "thinking" ||
|
|
currentChatState === "compacting"
|
|
)
|
|
return;
|
|
perfLog(
|
|
`[perf:send] ${sid} useChat.sendMessage start (textLen=${text.length}, attachments=${attachments?.length ?? 0})`,
|
|
);
|
|
|
|
const effectivePersonaInfo = resolvePersonaInfo(
|
|
overridePersona?.id,
|
|
overridePersona?.name,
|
|
);
|
|
const agent = useAgentStore.getState().getActiveAgent();
|
|
const providerId = providerOverride ?? agent?.provider ?? "goose";
|
|
const systemPrompt =
|
|
systemPromptOverride ?? agent?.systemPrompt ?? undefined;
|
|
|
|
// Ensure active session
|
|
setActiveSession(sessionId);
|
|
setPendingAssistantProvider(sessionId, providerId);
|
|
|
|
// Create and add user message
|
|
const userMessage = createUserMessage(
|
|
sendOptions?.displayText ?? text,
|
|
buildMessageAttachments(attachments),
|
|
sendOptions?.chips,
|
|
);
|
|
if (effectivePersonaInfo) {
|
|
userMessage.metadata = {
|
|
...userMessage.metadata,
|
|
targetPersonaId: effectivePersonaInfo.id,
|
|
targetPersonaName: effectivePersonaInfo.name,
|
|
};
|
|
}
|
|
// Embed image content blocks into the user message for local display
|
|
if (images && images.length > 0) {
|
|
for (const img of images) {
|
|
userMessage.content.push({
|
|
type: "image",
|
|
data: img.base64,
|
|
mimeType: img.mimeType,
|
|
});
|
|
}
|
|
}
|
|
addMessage(sessionId, userMessage);
|
|
setChatState(sessionId, "thinking");
|
|
setError(sessionId, null);
|
|
|
|
const sessionStore = useChatSessionStore.getState();
|
|
const session = sessionStore.getSession(sessionId);
|
|
|
|
// Immediately set the session/sidebar title from the user's message when
|
|
// the session still has the default placeholder. This gives instant
|
|
// feedback instead of waiting for acp:done or acp:session_info.
|
|
// A better backend-generated title will overwrite this if it arrives
|
|
// via the acp:session_info event.
|
|
if (session && isDefaultChatTitle(session.title)) {
|
|
sessionStore.patchSession(sessionId, {
|
|
title: getSessionTitleFromDraft(text, attachments),
|
|
updatedAt: new Date().toISOString(),
|
|
});
|
|
} else {
|
|
sessionStore.patchSession(sessionId, {
|
|
updatedAt: new Date().toISOString(),
|
|
});
|
|
}
|
|
|
|
options?.onMessageAccepted?.(sessionId);
|
|
|
|
clearDraft(sessionId);
|
|
|
|
const abort = new AbortController();
|
|
abortRef.current = abort;
|
|
|
|
try {
|
|
await ensurePreparedForPrompt(
|
|
options?.ensurePrepared,
|
|
effectivePersonaInfo?.id,
|
|
);
|
|
|
|
setChatState(sessionId, "streaming");
|
|
const promptWithPaths = appendAttachmentPaths(text.trim(), attachments);
|
|
const acpPrompt =
|
|
promptWithPaths || (images?.length ? " " : promptWithPaths);
|
|
const tAcp = performance.now();
|
|
perfLog(
|
|
`[perf:send] ${sid} → acpSendMessage (setup took ${(tAcp - tSendStart).toFixed(1)}ms)`,
|
|
);
|
|
await acpSendMessage(sessionId, acpPrompt, {
|
|
systemPrompt,
|
|
...(sendOptions?.assistantPrompt
|
|
? { assistantPrompt: sendOptions.assistantPrompt }
|
|
: {}),
|
|
|
|
images: images?.map(
|
|
(img) => [img.base64, img.mimeType] as [string, string],
|
|
),
|
|
});
|
|
perfLog(
|
|
`[perf:send] ${sid} acpSendMessage returned after ${(performance.now() - tAcp).toFixed(1)}ms (total sendMessage ${(performance.now() - tSendStart).toFixed(1)}ms)`,
|
|
);
|
|
|
|
setChatState(sessionId, "idle");
|
|
setStreamingMessageId(sessionId, null);
|
|
} catch (err) {
|
|
if (err instanceof DOMException && err.name === "AbortError") {
|
|
setChatState(sessionId, "idle");
|
|
} else {
|
|
const errorMessage = getErrorMessage(err);
|
|
const liveStore = useChatStore.getState();
|
|
const { streamingMessageId } = liveStore.getSessionRuntime(sessionId);
|
|
if (streamingMessageId) {
|
|
liveStore.updateMessage(
|
|
sessionId,
|
|
streamingMessageId,
|
|
(message) => ({
|
|
...message,
|
|
metadata: {
|
|
...message.metadata,
|
|
completionStatus: "error",
|
|
},
|
|
}),
|
|
);
|
|
}
|
|
|
|
liveStore.addMessage(
|
|
sessionId,
|
|
createSystemNotificationMessage(errorMessage, "error"),
|
|
);
|
|
setError(sessionId, errorMessage);
|
|
setChatState(sessionId, "idle");
|
|
setStreamingMessageId(sessionId, null);
|
|
}
|
|
setPendingAssistantProvider(sessionId, null);
|
|
} finally {
|
|
abortRef.current = null;
|
|
}
|
|
},
|
|
[
|
|
sessionId,
|
|
setActiveSession,
|
|
setPendingAssistantProvider,
|
|
addMessage,
|
|
setChatState,
|
|
setError,
|
|
clearDraft,
|
|
setStreamingMessageId,
|
|
providerOverride,
|
|
systemPromptOverride,
|
|
resolvePersonaInfo,
|
|
options,
|
|
],
|
|
);
|
|
|
|
const stopGeneration = useCallback(() => {
|
|
abortRef.current?.abort();
|
|
const activeStreamingMessageId = useChatStore
|
|
.getState()
|
|
.getSessionRuntime(sessionId).streamingMessageId;
|
|
|
|
setChatState(sessionId, "idle");
|
|
setStreamingMessageId(sessionId, null);
|
|
setPendingAssistantProvider(sessionId, null);
|
|
// Cancel the backend ACP session to stop orphaned streaming events
|
|
acpCancelSession(sessionId)
|
|
.then((wasCancelled) => {
|
|
if (wasCancelled && activeStreamingMessageId) {
|
|
markMessageStopped(sessionId, activeStreamingMessageId);
|
|
}
|
|
})
|
|
.catch(() => {
|
|
// Best-effort cancellation — ignore errors
|
|
});
|
|
}, [
|
|
setChatState,
|
|
setPendingAssistantProvider,
|
|
setStreamingMessageId,
|
|
sessionId,
|
|
]);
|
|
|
|
const clearChat = useCallback(() => {
|
|
abortRef.current?.abort();
|
|
clearMessages(sessionId);
|
|
setChatState(sessionId, "idle");
|
|
setStreamingMessageId(sessionId, null);
|
|
setPendingAssistantProvider(sessionId, null);
|
|
}, [
|
|
sessionId,
|
|
clearMessages,
|
|
setChatState,
|
|
setStreamingMessageId,
|
|
setPendingAssistantProvider,
|
|
]);
|
|
|
|
const getWorkingDir = useCallback(
|
|
() =>
|
|
useChatSessionStore.getState().activeWorkspaceBySession[sessionId]?.path,
|
|
[sessionId],
|
|
);
|
|
|
|
const compactConversation = useCallback(
|
|
async (overridePersona?: { id: string; name?: string }) => {
|
|
const currentChatState = useChatStore
|
|
.getState()
|
|
.getSessionRuntime(sessionId).chatState;
|
|
if (currentChatState !== "idle") {
|
|
return "skipped" as CompactConversationResult;
|
|
}
|
|
|
|
const effectivePersonaInfo = resolvePersonaInfo(
|
|
overridePersona?.id,
|
|
overridePersona?.name,
|
|
);
|
|
|
|
setActiveSession(sessionId);
|
|
setChatState(sessionId, "compacting");
|
|
setStreamingMessageId(sessionId, null);
|
|
setError(sessionId, null);
|
|
|
|
try {
|
|
await ensurePreparedForPrompt(
|
|
options?.ensurePrepared,
|
|
effectivePersonaInfo?.id,
|
|
);
|
|
} catch (err) {
|
|
const errorMessage = getErrorMessage(err);
|
|
addMessage(
|
|
sessionId,
|
|
createSystemNotificationMessage(errorMessage, "error"),
|
|
);
|
|
setError(sessionId, errorMessage);
|
|
setChatState(sessionId, "idle");
|
|
return "failed" as CompactConversationResult;
|
|
}
|
|
|
|
setSessionLoading(sessionId, true);
|
|
clearReplayBuffer(sessionId);
|
|
|
|
try {
|
|
await acpSendMessage(sessionId, MANUAL_COMPACT_TRIGGER);
|
|
|
|
// Command responses are streamed via prompt notifications, but the ACP
|
|
// layer does not currently forward history replacement events. Drop those
|
|
// transient chunks and refresh the session from replay instead.
|
|
clearReplayBuffer(sessionId);
|
|
const workingDir = getWorkingDir();
|
|
await acpLoadSession(sessionId, workingDir);
|
|
|
|
setSessionLoading(sessionId, false);
|
|
|
|
const buffer = getAndDeleteReplayBuffer(sessionId);
|
|
if (buffer) {
|
|
setMessages(sessionId, [
|
|
...buffer,
|
|
createCompactionConfirmationMessage(),
|
|
]);
|
|
} else {
|
|
addMessage(sessionId, createCompactionConfirmationMessage());
|
|
}
|
|
return "completed" as CompactConversationResult;
|
|
} catch (err) {
|
|
clearReplayBuffer(sessionId);
|
|
setSessionLoading(sessionId, false);
|
|
|
|
const errorMessage = getErrorMessage(err);
|
|
addMessage(
|
|
sessionId,
|
|
createSystemNotificationMessage(errorMessage, "error"),
|
|
);
|
|
setError(sessionId, errorMessage);
|
|
return "failed" as CompactConversationResult;
|
|
} finally {
|
|
setChatState(sessionId, "idle");
|
|
setStreamingMessageId(sessionId, null);
|
|
setPendingAssistantProvider(sessionId, null);
|
|
setSessionLoading(sessionId, false);
|
|
}
|
|
},
|
|
[
|
|
getWorkingDir,
|
|
options,
|
|
resolvePersonaInfo,
|
|
sessionId,
|
|
setActiveSession,
|
|
setChatState,
|
|
setStreamingMessageId,
|
|
setError,
|
|
addMessage,
|
|
setSessionLoading,
|
|
setMessages,
|
|
setPendingAssistantProvider,
|
|
],
|
|
);
|
|
|
|
const stopStreaming = stopGeneration;
|
|
|
|
return {
|
|
messages,
|
|
chatState: chatState as ChatState,
|
|
tokenState: tokenState as TokenState,
|
|
error,
|
|
streamingMessageId,
|
|
sendMessage,
|
|
stopGeneration,
|
|
stopStreaming,
|
|
clearChat,
|
|
compactConversation,
|
|
isStreaming,
|
|
};
|
|
}
|