Files
tkmind_go/ui/goose2/src/features/chat/hooks/useChat.ts
T
2026-05-08 20:10:55 +00:00

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,
};
}