Repository navigation
Refactor streaming and add message edit/delete endpoints - #395
Conversation
… message recovery - Add missing PATCH/DELETE routes for page-agent messages that caused 404 HTML errors when editing/deleting messages in agent conversation mode - Wrap Global Assistant streaming in createUIMessageStream so server-side processing continues even if client disconnects (matches page chat route pattern) - Stabilize DefaultChatTransport creation using refs in all chat views (AiChatView, GlobalAssistantView, SidebarChatTab, GlobalChatContext) to prevent useChat from resetting state when transport instances are recreated - Add optimistic local updates for message editing so edits are preserved even when the server save fails https://claude.ai/code/session_019a3TCSNUwiz1SCRXePfUvG
📝 WalkthroughWalkthroughThis PR refactors AI chat streaming and message handling by extracting transport, streaming state, and stop logic into reusable hooks. It introduces a new API route for editing and deleting conversation messages with permissions validation, and adds optimistic UI updates for message actions to improve user experience. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
… stop patterns Deduplicates transport stability, streaming registration, and stop handler patterns that were copy-pasted across 4 chat components into 3 composable hooks. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In
`@apps/web/src/app/api/ai/page-agents/`[agentId]/conversations/[conversationId]/messages/[messageId]/route.ts:
- Around line 52-61: The message retrieval uses
chatMessageRepository.getMessageById which does not check the soft-delete flag,
so add an isActive guard after fetching the message: after const message = await
chatMessageRepository.getMessageById(messageId) and the not-found check, verify
message.isActive === true and return a 404 JSON (same shape as other errors) if
false; apply the same isActive check to the DELETE handler path (so deleting an
already soft-deleted message returns 404 instead of re-soft-deleting).
In `@apps/web/src/lib/ai/shared/hooks/useMessageActions.ts`:
- Around line 118-123: The current catch around the entire edit flow conflates
failures of patch() and the subsequent fetchWithAuth(), causing a misleading
"failed to sync" toast even when patch() succeeded; update useMessageActions.ts
to split the operations into two try/catch blocks: first await patch(...) and on
its failure show the existing "failed to edit" toast and revert/handle as
before, then in a separate try call fetchWithAuth(...) to refresh the cache and
on its failure log the error and show a distinct, accurate toast (e.g., "Edit
saved to server but failed to refresh local data") without undoing the
optimistic update; reference the existing patch() call and fetchWithAuth() call
in your changes and ensure console.error/processLogger entries reflect which
step failed.
🧹 Nitpick comments (7)
apps/web/src/app/api/ai/page-agents/[agentId]/conversations/[conversationId]/messages/[messageId]/route.ts (2)
3-4: Merge duplicate import sources.
canUserEditPageandloggersare both imported from@pagespace/lib/serveron separate lines.♻️ Suggested fix
-import { canUserEditPage } from '@pagespace/lib/server'; -import { loggers } from '@pagespace/lib/server'; +import { canUserEditPage, loggers } from '@pagespace/lib/server';
72-90: Activity logging is actually blocking despite the "non-blocking" comment.
await getActorInfo(userId)on line 74 is awaited, which delays the response. The same pattern appears in the DELETE handler (line 160). If latency matters, fire-and-forget the entire block:♻️ Suggested approach (for both PATCH and DELETE)
- // Log activity for audit trail (non-blocking) - try { - const actorInfo = await getActorInfo(userId); - logMessageActivity(userId, 'message_update', { ... }, actorInfo, { ... }); - } catch (loggingError) { - loggers.api.error('Failed to log agent message update activity', loggingError as Error, { ... }); - } + // Log activity for audit trail (non-blocking) + getActorInfo(userId) + .then((actorInfo) => { + logMessageActivity(userId, 'message_update', { ... }, actorInfo, { ... }); + }) + .catch((loggingError) => { + loggers.api.error('Failed to log agent message update activity', loggingError as Error, { ... }); + });Alternatively, if the await is intentional, update the comment to say "// Log activity for audit trail".
apps/web/src/lib/ai/shared/hooks/useMessageActions.ts (2)
67-126:messagesinuseCallbackdeps causes handler instability during streaming.All three handlers capture
messagesdirectly and include it in their dependency arrays, producing a new function identity on every message update. During streaming this means rapid re-creation, which can trigger unnecessary re-renders in any component that receives these handlers as props.A
useReffor the latest messages avoids the dependency while keeping reads fresh:♻️ Suggested refactor
+import { useCallback, useRef, useEffect } from 'react'; -import { useCallback } from 'react'; export function useMessageActions({ ... messages, setMessages, ... }: UseMessageActionsOptions): UseMessageActionsResult { + const messagesRef = useRef(messages); + useEffect(() => { messagesRef.current = messages; }, [messages]); const handleEdit = useCallback( async (messageId: string, newContent: string) => { if (!conversationId) return; - const updatedMessages = messages.map((m) => { + const updatedMessages = messagesRef.current.map((m) => { ... }); setMessages(updatedMessages); ... }, - [isAgentMode, agentId, conversationId, messages, setMessages, onEditVersionChange] + [isAgentMode, agentId, conversationId, setMessages, onEditVersionChange] );Apply the same pattern to
handleDeleteandhandleRetry.
160-203: Sequential deletes inhandleRetrycould be parallelized.Lines 173-185
awaiteach deletion sequentially. When multiple assistant messages need cleanup,Promise.allSettledwould be faster and still handle partial failures gracefully.♻️ Suggested refactor
- for (const msg of assistantMessagesToDelete) { - try { - if (isAgentMode) { - await del( - `/api/ai/page-agents/${agentId}/conversations/${conversationId}/messages/${msg.id}` - ); - } else { - await del(`/api/ai/global/${conversationId}/messages/${msg.id}`); - } - } catch (error) { - console.error('Failed to delete old assistant message:', error); - } - } + await Promise.allSettled( + assistantMessagesToDelete.map((msg) => { + const url = isAgentMode + ? `/api/ai/page-agents/${agentId}/conversations/${conversationId}/messages/${msg.id}` + : `/api/ai/global/${conversationId}/messages/${msg.id}`; + return del(url).catch((error) => + console.error('Failed to delete old assistant message:', error) + ); + }) + );apps/web/src/lib/ai/shared/hooks/useChatTransport.ts (1)
23-29:apiparameter changes are not tracked — transport won't update if onlyapichanges.The recreation condition at Line 23 only checks
conversationId. Ifapichanges whileconversationIdstays the same, the stale transport with the oldapiwill be returned. Currently this is safe because callers always coupleapiwithconversationId, but it's a subtle implicit contract.Consider including
apiin the tracking check to make the hook robust against future misuse:♻️ Suggested fix
- const trackingIdRef = useRef<string | null>(null); + const trackingIdRef = useRef<string | null>(null); + const apiRef = useRef<string>(api); if (!conversationId) { return null; } - if (trackingIdRef.current !== conversationId || !transportRef.current) { + if (trackingIdRef.current !== conversationId || apiRef.current !== api || !transportRef.current) { transportRef.current = new DefaultChatTransport({ api, fetch: createStreamTrackingFetch({ chatId: conversationId }), }); trackingIdRef.current = conversationId; + apiRef.current = api; }apps/web/src/app/api/ai/global/[id]/messages/route.ts (1)
740-801: Solid streaming refactor usingcreateUIMessageStreamfor server-side processing continuity.The architecture correctly ensures
onFinishcallbacks (DB save, usage tracking) run even if the client disconnects. TheusagePromiseclosure timing is correct —executecompletes beforeonFinishfires, so the promise is always assigned.One minor nit: the
.then((usage) => usage)on Line 786 is an identity transform that can be dropped:♻️ Simplify usagePromise
- usagePromise = aiResult.totalUsage - .then((usage) => usage) - .catch((error) => { + usagePromise = aiResult.totalUsage + .catch((error) => {apps/web/src/components/layout/middle-content/page-views/ai-page/AiChatView.tsx (1)
153-166: Non-null assertion ontransport!is fragile — prefer an explicit guard.
transportisDefaultChatTransport | nullperuseChatTransport's return type. It's currently safe becausestreamTrackingIdis always a string (currentConversationId || page.id), so the hook never returnsnull. However, the!assertion silently suppresses null checking and could mask a runtime error ifstreamTrackingIdbecomes nullable in the future.In contrast,
GlobalAssistantView.tsxuses the safer pattern: theagentChatConfigmemo returnsnullwhen transport is null (line 265), anduseChatreceivesagentChatConfig || {}.Consider matching that pattern here:
Proposed guard for null transport
const chatConfig = useMemo( - () => ({ + () => !transport ? null : ({ id: page.id, messages: initialMessages, - transport: transport!, + transport, experimental_throttle: 100, onError: (error: Error) => { console.error('AiChatView: Chat error:', error); }, }), [page.id, transport, initialMessages] ); const { messages, sendMessage, status, error, regenerate, setMessages, stop: chatStop } = - useChat(chatConfig); + useChat(chatConfig || {});
…picks - Add isActive check to PATCH/DELETE handlers to prevent editing/deleting soft-deleted messages - Split edit error handling in useMessageActions so patch() failure and refetch failure show accurate error toasts - Merge duplicate imports in route.ts - Track api parameter changes in useChatTransport for robustness - Remove transport! non-null assertion in AiChatView, use null guard - Remove identity .then() transform in global messages route - Fix "non-blocking" comment on blocking activity logging Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Review Feedback Addressed (193e497)Actionable Comments1. Missing 2. Misleading error toast in Nitpicks Addressed
Nitpicks Acknowledged (not changed)
|
2 participants
Summary
This PR refactors the AI chat streaming implementation to use the new
createUIMessageStreamAPI for better server-side processing continuity, adds message edit/delete endpoints for page agents, optimizes transport recreation in chat components to prevent unnecessary state resets, and extracts shared chat hooks to reduce duplication.Key Changes
Streaming Architecture
createUIMessageStreamandcreateUIMessageStreamResponseinstead of directtoUIMessageStreamResponse()callNew Endpoints
Page Agent Message Edit (
PATCH /api/ai/page-agents/[agentId]/conversations/[conversationId]/messages/[messageId])editedAttimestampisActiveguard, and activity loggingPage Agent Message Delete (
DELETE /api/ai/page-agents/[agentId]/conversations/[conversationId]/messages/[messageId])isActiveto falseShared Chat Hooks (new)
Extracted 3 composable hooks from duplicated patterns across 4 chat components:
useChatTransport— StableDefaultChatTransportinstance that only recreates when conversation ID or API endpoint changesuseStreamingRegistration— Registers/unregisters streaming state withuseEditingStoreuseChatStop— Memoized stop function composing server abort + client stop via try/finallyComponent Optimizations
AiChatView,GlobalAssistantView,SidebarChatTab,GlobalChatContextwithuseChatTransporthookuseEffectblocks withuseStreamingRegistrationhookuseChatStophookMessage Actions
useMessageActionshookCode Quality
.then()transformHow to Validate
Implementation Notes
usagePromise) for async token countingonFinishfiresuseChatto reset and lose message state