improvement(mothership): chat history and stability

This commit is contained in:
Emir Karabeg
2026-03-09 15:57:13 -07:00
parent 917af6d141
commit 8fd8b1a248
2 changed files with 78 additions and 18 deletions

View File

@@ -26,6 +26,7 @@ import type {
import { SUBAGENT_LABELS } from '../types'
import {
extractFileResource,
extractResourcesFromHistory,
extractTableResource,
extractWorkflowResource,
RESOURCE_TOOL_NAMES,
@@ -184,6 +185,12 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
if (!chatHistory || appliedChatIdRef.current === chatHistory.id) return
appliedChatIdRef.current = chatHistory.id
setMessages(chatHistory.messages.map(mapStoredMessage))
const restored = extractResourcesFromHistory(chatHistory.messages)
if (restored.length > 0) {
setResources(restored)
setActiveResourceId(restored[restored.length - 1].id)
}
}, [chatHistory])
const processSSEStream = useCallback(
@@ -194,6 +201,9 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
const toolMap = new Map<string, number>()
let lastTableId: string | null = null
let lastWorkflowId: string | null = null
let runningText = ''
toolArgsMapRef.current.clear()
const ensureTextBlock = (): ContentBlock => {
const last = blocks[blocks.length - 1]
@@ -204,13 +214,9 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
}
const flush = () => {
const text = blocks
.filter((b) => b.type === 'text')
.map((b) => b.content ?? '')
.join('')
setMessages((prev) =>
prev.map((m) =>
m.id === assistantId ? { ...m, content: text, contentBlocks: [...blocks] } : m
m.id === assistantId ? { ...m, content: runningText, contentBlocks: [...blocks] } : m
)
)
}
@@ -267,6 +273,7 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
if (chunk) {
const tb = ensureTextBlock()
tb.content = (tb.content ?? '') + chunk
runningText += chunk
flush()
}
break
@@ -411,12 +418,7 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
[workspaceId, queryClient, addResource]
)
const finalize = useCallback(() => {
sendingRef.current = false
setIsSending(false)
abortControllerRef.current = null
chatBottomRef.current?.scrollIntoView({ behavior: 'smooth' })
const invalidateChatQueries = useCallback(() => {
const activeChatId = chatIdRef.current
if (activeChatId) {
queryClient.invalidateQueries({ queryKey: taskKeys.detail(activeChatId) })
@@ -424,6 +426,14 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
queryClient.invalidateQueries({ queryKey: taskKeys.list(workspaceId) })
}, [workspaceId, queryClient])
const finalize = useCallback(() => {
sendingRef.current = false
setIsSending(false)
abortControllerRef.current = null
chatBottomRef.current?.scrollIntoView({ behavior: 'smooth' })
invalidateChatQueries()
}, [invalidateChatQueries])
useEffect(() => {
const activeStreamId = chatHistory?.activeStreamId
if (!activeStreamId || !appliedChatIdRef.current || sendingRef.current) return
@@ -556,13 +566,8 @@ export function useChat(workspaceId: string, initialChatId?: string): UseChatRet
abortControllerRef.current = null
sendingRef.current = false
setIsSending(false)
const activeChatId = chatIdRef.current
if (activeChatId) {
queryClient.invalidateQueries({ queryKey: taskKeys.detail(activeChatId) })
}
queryClient.invalidateQueries({ queryKey: taskKeys.list(workspaceId) })
}, [workspaceId, queryClient])
invalidateChatQueries()
}, [invalidateChatQueries])
useEffect(() => {
return () => {

View File

@@ -1,3 +1,4 @@
import type { TaskStoredMessage } from '@/hooks/queries/tasks'
import type { MothershipResource, SSEPayload } from './types'
export const RESOURCE_TOOL_NAMES = new Set([
@@ -78,3 +79,57 @@ export function extractWorkflowResource(
return null
}
const GENERIC_TITLES = new Set(['Table', 'File', 'Workflow'])
/**
* Reconstructs the MothershipResource list from persisted tool calls.
* Adapts each stored tool call into an SSEPayload so the existing
* extract*Resource functions are reused with zero duplication.
* Deduplicates by type+id while preserving insertion order.
*/
export function extractResourcesFromHistory(messages: TaskStoredMessage[]): MothershipResource[] {
const resourceMap = new Map<string, MothershipResource>()
let lastTableId: string | null = null
let lastWorkflowId: string | null = null
for (const msg of messages) {
if (!msg.toolCalls) continue
for (const tc of msg.toolCalls) {
if (tc.status !== 'success' || !RESOURCE_TOOL_NAMES.has(tc.name)) continue
const payload: SSEPayload = {
type: 'tool_result',
result: tc.result as Record<string, unknown>,
success: true,
toolName: tc.name,
}
const args = tc.params as Record<string, unknown> | undefined
let resource: MothershipResource | null = null
if (tc.name === 'user_table') {
resource = extractTableResource(payload, args, lastTableId)
if (resource) lastTableId = resource.id
} else if (tc.name === 'workspace_file') {
resource = extractFileResource(payload, args)
} else if (tc.name === 'create_workflow' || tc.name === 'edit_workflow') {
resource = extractWorkflowResource(payload, lastWorkflowId)
if (resource) lastWorkflowId = resource.id
}
if (resource) {
const key = `${resource.type}:${resource.id}`
const existing = resourceMap.get(key)
if (
!existing ||
(GENERIC_TITLES.has(existing.title) && !GENERIC_TITLES.has(resource.title))
) {
resourceMap.set(key, resource)
}
}
}
}
return Array.from(resourceMap.values())
}