feat: add keyboard shortcuts and hints (#1138, #1222)

apps/backend/src/queries/chat.queries.ts

@@ -5,7 +5,7 @@ import type {
 	GroupedChatListResponse,
 	LlmProvider,
 } from '@nao/shared/types';
-import { and, asc, desc, eq, gte, inArray, isNotNull, isNull, like, ne, or, sql } from 'drizzle-orm';
+import { and, asc, count, desc, eq, gte, inArray, isNotNull, isNull, like, ne, or, sql } from 'drizzle-orm';
 
 import s, {
 	DBChat,
@@ -262,17 +262,11 @@ const aggregateChatMessagParts = (
 			if (acc[row.chat_message.id]) {
 				acc[row.chat_message.id].parts.push(uiPart);
 			} else {
-				acc[row.chat_message.id] = {
-					id: row.chat_message.id,
-					role: row.chat_message.role,
-					parts: [uiPart],
-					feedback: row.message_feedback ?? undefined,
-					source: row.chat_message.source ?? undefined,
-					isForked: row.chat_message.isForked ?? undefined,
-					citation: row.chat_message.citation ?? undefined,
-					stopReason: row.chat_message.stopReason ?? undefined,
-					createdAt: row.chat_message.createdAt?.getTime(),
-				};
+				acc[row.chat_message.id] = convertDBMessageToUIMessage(
+					row.chat_message,
+					[uiPart],
+					row.message_feedback,
+				);
 			}
 			return acc;
 		},
@@ -282,6 +276,22 @@ const aggregateChatMessagParts = (
 	return Object.values(messagesMap);
 };
 
+const convertDBMessageToUIMessage = (
+	message: DBChatMessage,
+	parts: UIMessagePart[],
+	feedback?: MessageFeedback | null,
+): UIMessage => ({
+	id: message.id,
+	role: message.role,
+	parts,
+	feedback: feedback ?? undefined,
+	source: message.source ?? undefined,
+	isForked: message.isForked ?? undefined,
+	citation: message.citation ?? undefined,
+	stopReason: message.stopReason ?? undefined,
+	createdAt: message.createdAt?.getTime(),
+});
+
 export const getChatMessages = async (chatId: string): Promise<UIMessage[]> => {
 	const result = await db
 		.select()
@@ -294,6 +304,87 @@ export const getChatMessages = async (chatId: string): Promise<UIMessage[]> => {
 	return aggregateChatMessagParts(result);
 };
 
+export const getLastTurn = async (
+	chatId: string,
+): Promise<
+	| {
+			userMessageId: string;
+			versionGroupId: string | null;
+			versionGroupSize: number;
+			assistantMessage?: UIMessage;
+	  }
+	| undefined
+> => {
+	const activeMessages = await db
+		.select()
+		.from(s.chatMessage)
+		.where(and(eq(s.chatMessage.chatId, chatId), isNull(s.chatMessage.supersededAt)))
+		.orderBy(desc(s.chatMessage.createdAt))
+		.execute();
+	const lastMessage = activeMessages.at(0);
+	if (!lastMessage || lastMessage.role === 'system') {
+		return undefined;
+	}
+
+	const userMessage =
+		lastMessage.role === 'user' ? lastMessage : activeMessages.find((message) => message.role === 'user');
+	if (!userMessage) {
+		return undefined;
+	}
+
+	const versionGroupSize = userMessage.versionGroupId
+		? await countMessagesInVersionGroup(chatId, userMessage.versionGroupId)
+		: 1;
+	const assistantMessage =
+		lastMessage.role === 'assistant' ? await getMessageWithParts(chatId, lastMessage) : undefined;
+
+	return {
+		userMessageId: userMessage.id,
+		versionGroupId: userMessage.versionGroupId,
+		versionGroupSize,
+		assistantMessage,
+	};
+};
+
+const countMessagesInVersionGroup = async (chatId: string, versionGroupId: string): Promise<number> => {
+	const [result] = await db
+		.select({ value: count() })
+		.from(s.chatMessage)
+		.where(and(eq(s.chatMessage.chatId, chatId), eq(s.chatMessage.versionGroupId, versionGroupId)))
+		.execute();
+	return result?.value ?? 0;
+};
+
+const getMessageWithParts = async (chatId: string, message: DBChatMessage): Promise<UIMessage> => {
+	const result = await db
+		.select()
+		.from(s.chatMessage)
+		.innerJoin(s.messagePart, eq(s.messagePart.messageId, s.chatMessage.id))
+		.where(and(eq(s.chatMessage.chatId, chatId), eq(s.chatMessage.id, message.id)))
+		.orderBy(asc(s.messagePart.order))
+		.execute();
+	return aggregateChatMessagParts(result).at(0) ?? convertDBMessageToUIMessage(message, []);
+};
+
+export const deleteMessagesByIds = async (chatId: string, ids: string[]): Promise<void> => {
+	if (ids.length === 0) {
+		return;
+	}
+	await db
+		.delete(s.chatMessage)
+		.where(and(eq(s.chatMessage.chatId, chatId), inArray(s.chatMessage.id, ids)))
+		.execute();
+};
+
+export const countActiveMessages = async (chatId: string): Promise<number> => {
+	const [result] = await db
+		.select({ value: count() })
+		.from(s.chatMessage)
+		.where(and(eq(s.chatMessage.chatId, chatId), isNull(s.chatMessage.supersededAt)))
+		.execute();
+	return result?.value ?? 0;
+};
+
 export const getChatOwnerId = async (chatId: string): Promise<string | undefined> => {
 	const [result] = await db
 		.select({

apps/backend/src/services/agent.ts

@@ -344,6 +344,8 @@ export const MAX_OUTPUT_TOKENS = 16_000;
 
 class AgentManager {
 	private readonly _agent: ToolLoopAgent<never, AgentTools, never>;
+	private readonly _finished: Promise<void>;
+	private _resolveFinished: (() => void) | undefined;
 	private _streamWriter?: UIMessageStreamWriter<UIMessage>;
 
 	constructor(
@@ -357,6 +359,9 @@ class AgentManager {
 		stopWhen: StopCondition<AgentTools>[] = [hasToolCall('suggest_follow_ups'), hasToolCall('clarification')],
 		private readonly _systemPromptOverride?: string,
 	) {
+		this._finished = new Promise((resolve) => {
+			this._resolveFinished = resolve;
+		});
 		const callSettings = this._modelConfig.callSettings ?? {};
 		const provider = this._modelSelection.provider;
 		const providerOptions = fitThinkingBudget(this._modelConfig.providerOptions, this._maxOutputTokens);
@@ -518,7 +523,7 @@ class AgentManager {
 						llmModelId: this._modelSelection.modelId,
 					});
 				} finally {
-					this._onDispose();
+					this._finish();
 				}
 			},
 		});
@@ -802,7 +807,7 @@ class AgentManager {
 				responseParts: [],
 			};
 		} finally {
-			this._onDispose();
+			this._finish();
 		}
 	}
 
@@ -814,6 +819,23 @@ class AgentManager {
 		this._abortController.abort();
 	}
 
+	waitUntilFinished(): Promise<void> {
+		return this._finished;
+	}
+
+	private _markFinished(): void {
+		this._resolveFinished?.();
+		this._resolveFinished = undefined;
+	}
+
+	private _finish(): void {
+		try {
+			this._onDispose();
+		} finally {
+			this._markFinished();
+		}
+	}
+
 	private _addCitationContext(messages: UIMessage[]): UIMessage[] {
 		const [lastUserMessage] = findLastUserMessage(messages);
 		if (!lastUserMessage?.citation) {

apps/backend/src/trpc/chat.routes.ts

@@ -80,6 +80,45 @@ export const chatRoutes = {
 		posthog.capture(ctx.user.id, PostHogEvent.AgentStopped, { project_id: projectId, chat_id: input.chatId });
 	}),
 
+	cancel: chatOwnerProcedure.input(z.object({ chatId: z.string(), hadContent: z.boolean() })).mutation(
+		async ({
+			input,
+			ctx,
+		}): Promise<{
+			outcome: 'deleted' | 'kept';
+			chatDeleted: boolean;
+		}> => {
+			const agent = agentService.get(input.chatId);
+			if (agent) {
+				agent.stop();
+				posthog.capture(ctx.user.id, PostHogEvent.AgentStopped, {
+					project_id: agent.chat.projectId,
+					chat_id: input.chatId,
+				});
+				await Promise.race([agent.waitUntilFinished(), delay(10_000)]);
+			}
+
+			const turn = await chatQueries.getLastTurn(input.chatId);
+			if (!turn) {
+				return { outcome: 'kept', chatDeleted: false };
+			}
+
+			if (input.hadContent || turn.versionGroupSize > 1) {
+				return { outcome: 'kept', chatDeleted: false };
+			}
+
+			const idsToDelete = [turn.userMessageId, ...(turn.assistantMessage ? [turn.assistantMessage.id] : [])];
+			await chatQueries.deleteMessagesByIds(input.chatId, idsToDelete);
+
+			const remaining = await chatQueries.countActiveMessages(input.chatId);
+			if (remaining === 0) {
+				await chatQueries.deleteChat(input.chatId);
+				return { outcome: 'deleted', chatDeleted: true };
+			}
+			return { outcome: 'deleted', chatDeleted: false };
+		},
+	),
+
 	rename: chatOwnerProcedure
 		.input(z.object({ chatId: z.string(), title: z.string().min(1).max(255) }))
 		.mutation(async ({ input, ctx }): Promise<void> => {
@@ -135,3 +174,8 @@ export const chatRoutes = {
 			return usage;
 		}),
 };
+
+const delay = (milliseconds: number): Promise<void> =>
+	new Promise((resolve) => {
+		setTimeout(resolve, milliseconds);
+	});

apps/backend/src/utils/ai.ts

@@ -97,6 +97,16 @@ export const createChatTitle = ({ text }: { text: string }) => {
 	return text.slice(0, 64);
 };
 
+export const checkAssistantMessageHasContent = (message: UIMessage): boolean =>
+	message.parts.some(
+		(part) =>
+			part.type !== 'step-start' &&
+			part.type !== 'tool-suggest_follow_ups' &&
+			part.type !== 'reasoning' &&
+			part.type !== 'data-newChat' &&
+			part.type !== 'data-newUserMessage',
+	);
+
 export const joinAllTextParts = (message: UIMessage, separator: string = '\n'): string => {
 	return message.parts
 		.filter((part) => part.type === 'text')

apps/frontend/src/components/chat-input.tsx

@@ -33,6 +33,7 @@ import { cn } from '@/lib/utils';
 import { useChatId } from '@/hooks/use-chat-id';
 import { usePermissions } from '@/hooks/use-permissions';
 import { messageQueueStore } from '@/stores/chat-message-queue';
+import { chatInputRestoreStore, useChatInputRestore } from '@/stores/chat-input-restore';
 import { chatPendingCitationStore } from '@/stores/chat-pending-citation';
 import { useChatPendingCitation } from '@/hooks/use-chat-pending-citation';
 import { SelectionCitationBanner } from '@/components/selection-citation-banner';
@@ -96,7 +97,7 @@ function ChatInputBase({
 	const [inputText, setInputText] = useState('');
 	const {
 		isRunning,
-		stopAgent,
+		cancelAgent,
 		isLoadingMessages,
 		adminMode,
 		setAdminMode,
@@ -110,6 +111,7 @@ function ChatInputBase({
 
 	const isAdminMode = isAdmin && adminMode;
 	const imageUpload = useImageUpload();
+	const chatInputRestore = useChatInputRestore(!!allowQueueing);
 	const effectivePlaceholder = isRunning && allowQueueing ? 'Add a follow-up...' : placeholder;
 
 	const agentSettings = useQuery(trpc.project.getAgentSettings.queryOptions());
@@ -132,6 +134,33 @@ function ChatInputBase({
 
 	useEffect(() => promptRef.current?.focus(), [chatId, promptRef]);
 
+	useEffect(() => {
+		if (!allowQueueing || !chatInputRestore) {
+			return;
+		}
+
+		chatInputRestoreStore.clear();
+		promptRef.current?.clear();
+		promptRef.current?.insertText(chatInputRestore.text);
+		setInputText(chatInputRestore.text);
+		imageUpload.clearImages();
+
+		if (chatInputRestore.citation && chatId) {
+			chatPendingCitationStore.set({ ...chatInputRestore.citation, chatId });
+		}
+
+		const restoreImages = async () => {
+			const files = await Promise.all(
+				chatInputRestore.images.map(({ url, mediaType }, index) =>
+					dataUrlToFile(url, mediaType, `image-${index + 1}`),
+				),
+			);
+			await imageUpload.addFiles(files);
+			requestAnimationFrame(() => promptRef.current?.focus());
+		};
+		restoreImages();
+	}, [allowQueueing, chatInputRestore]); // eslint-disable-line react-hooks/exhaustive-deps
+
 	useEffect(() => {
 		const el = dropZoneRef.current;
 		if (!el) {
@@ -382,14 +411,14 @@ function ChatInputBase({
 								<ChatButton
 									showStop={isInputEmpty}
 									disabled={!isInputEmpty && isBudgetExceeded}
-									onClick={isInputEmpty ? stopAgent : handleSubmitMessage}
+									onClick={isInputEmpty ? cancelAgent : handleSubmitMessage}
 									type='button'
 								/>
 							) : (
 								<ChatButton
 									showStop={isRunning}
 									disabled={isLoadingMessages || isInputEmpty || (!isRunning && isBudgetExceeded)}
-									onClick={isRunning ? stopAgent : handleSubmitMessage}
+									onClick={isRunning ? cancelAgent : handleSubmitMessage}
 									type='button'
 								/>
 							)}
@@ -401,6 +430,12 @@ function ChatInputBase({
 	);
 }
 
+async function dataUrlToFile(url: string, mediaType: string, name: string): Promise<File> {
+	const response = await fetch(url);
+	const blob = await response.blob();
+	return new File([blob], name, { type: mediaType });
+}
+
 const CHAT_INPUT_BORDER_RADIUS = 18;
 const CHAT_INPUT_BORDER_STROKE = 1;
 

apps/frontend/src/components/chat-messages/assistant-message.tsx

@@ -11,6 +11,7 @@ import { TextShimmer } from '@/components/ui/text-shimmer';
 import { AssistantMessageActions } from '@/components/chat-messages/assistant-message-actions';
 import { cn, isLast } from '@/lib/utils';
 import { useChatId } from '@/hooks/use-chat-id';
+import { useIsCancellingMessage } from '@/hooks/use-is-cancelling-message-store';
 import { useToolCallDensity } from '@/hooks/use-tool-call-density';
 import { AssistantMessageProvider, useAssistantMessage } from '@/contexts/assistant-message';
 
@@ -37,6 +38,7 @@ export const AssistantMessage = memo(
 			[message.parts, toolCallDensity],
 		);
 		const hasContent = useMemo(() => checkAssistantMessageHasContent(message), [message]);
+		const isCancelling = useIsCancellingMessage(message.id);
 		const isCompacting = message.parts.at(-1)?.type === 'data-compactionSummaryStarted';
 		const showActions = message.id !== storyIntroMessageId;
 		const hasFeedback = message.feedback != null;
@@ -45,6 +47,10 @@ export const AssistantMessage = memo(
 			return null;
 		}
 
+		if (isCancelling && isSettled && !hasContent) {
+			return null;
+		}
+
 		return (
 			<AssistantMessageProvider isSettled={isSettled}>
 				<div className={cn('group px-3 flex flex-col gap-2 bg-transparent')}>