mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-19 13:33:35 +08:00
refactor: chatStore supports per-session isolated state
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
程序员阿江(Relakkes)
parent
c056b3c4fd
commit
9dc2294954
@ -11,7 +11,7 @@ import type { AttachmentRef, ChatState, UIAttachment, UIMessage, ServerMessage,
|
||||
|
||||
type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'reconnecting'
|
||||
|
||||
type ChatStore = {
|
||||
export type PerSessionState = {
|
||||
messages: UIMessage[]
|
||||
chatState: ChatState
|
||||
connectionState: ConnectionState
|
||||
@ -29,31 +29,11 @@ type ChatStore = {
|
||||
tokenUsage: TokenUsage
|
||||
elapsedSeconds: number
|
||||
statusVerb: string
|
||||
connectedSessionId: string | null
|
||||
slashCommands: Array<{ name: string; description: string }>
|
||||
|
||||
// Actions
|
||||
connectToSession: (sessionId: string) => void
|
||||
disconnectSession: () => void
|
||||
sendMessage: (content: string, attachments?: AttachmentRef[]) => void
|
||||
respondToPermission: (requestId: string, allowed: boolean, rule?: string) => void
|
||||
setSessionPermissionMode: (mode: PermissionMode) => void
|
||||
stopGeneration: () => void
|
||||
loadHistory: (sessionId: string) => Promise<void>
|
||||
clearMessages: () => void
|
||||
handleServerMessage: (msg: ServerMessage) => void
|
||||
elapsedTimer: ReturnType<typeof setInterval> | null
|
||||
}
|
||||
|
||||
const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite'])
|
||||
|
||||
/** Track tool_use IDs for task-related tools, so we can refresh on tool_result */
|
||||
const pendingTaskToolUseIds = new Set<string>()
|
||||
|
||||
let msgCounter = 0
|
||||
const nextId = () => `msg-${++msgCounter}-${Date.now()}`
|
||||
let elapsedTimer: ReturnType<typeof setInterval> | null = null
|
||||
|
||||
export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
const DEFAULT_SESSION_STATE: PerSessionState = {
|
||||
messages: [],
|
||||
chatState: 'idle',
|
||||
connectionState: 'disconnected',
|
||||
@ -66,97 +46,127 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
tokenUsage: { input_tokens: 0, output_tokens: 0 },
|
||||
elapsedSeconds: 0,
|
||||
statusVerb: '',
|
||||
connectedSessionId: null,
|
||||
slashCommands: [],
|
||||
elapsedTimer: null,
|
||||
}
|
||||
|
||||
connectToSession: (sessionId: string) => {
|
||||
const current = get().connectedSessionId
|
||||
if (current === sessionId) return
|
||||
function createDefaultSessionState(): PerSessionState {
|
||||
return { ...DEFAULT_SESSION_STATE, messages: [], tokenUsage: { input_tokens: 0, output_tokens: 0 } }
|
||||
}
|
||||
|
||||
// Disconnect previous
|
||||
if (current) wsManager.disconnect()
|
||||
type ChatStore = {
|
||||
sessions: Record<string, PerSessionState>
|
||||
|
||||
set({
|
||||
connectedSessionId: sessionId,
|
||||
connectionState: 'connecting',
|
||||
messages: [],
|
||||
chatState: 'idle',
|
||||
streamingText: '',
|
||||
streamingToolInput: '',
|
||||
activeToolUseId: null,
|
||||
activeToolName: null,
|
||||
activeThinkingId: null,
|
||||
pendingPermission: null,
|
||||
elapsedSeconds: 0,
|
||||
slashCommands: [],
|
||||
})
|
||||
getSession: (sessionId: string) => PerSessionState
|
||||
connectToSession: (sessionId: string) => void
|
||||
disconnectSession: (sessionId: string) => void
|
||||
sendMessage: (sessionId: string, content: string, attachments?: AttachmentRef[]) => void
|
||||
respondToPermission: (sessionId: string, requestId: string, allowed: boolean, rule?: string) => void
|
||||
setSessionPermissionMode: (sessionId: string, mode: PermissionMode) => void
|
||||
stopGeneration: (sessionId: string) => void
|
||||
loadHistory: (sessionId: string) => Promise<void>
|
||||
clearMessages: (sessionId: string) => void
|
||||
handleServerMessage: (sessionId: string, msg: ServerMessage) => void
|
||||
}
|
||||
|
||||
// Clear all previous handlers before registering new ones
|
||||
// This prevents handler accumulation causing message duplication
|
||||
wsManager.clearHandlers()
|
||||
const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite'])
|
||||
const pendingTaskToolUseIds = new Set<string>()
|
||||
|
||||
let msgCounter = 0
|
||||
const nextId = () => `msg-${++msgCounter}-${Date.now()}`
|
||||
|
||||
/** Helper: immutably update a specific session within the sessions record */
|
||||
function updateSessionIn(
|
||||
sessions: Record<string, PerSessionState>,
|
||||
sessionId: string,
|
||||
updater: (s: PerSessionState) => Partial<PerSessionState>,
|
||||
): Record<string, PerSessionState> {
|
||||
const session = sessions[sessionId]
|
||||
if (!session) return sessions
|
||||
return { ...sessions, [sessionId]: { ...session, ...updater(session) } }
|
||||
}
|
||||
|
||||
export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
sessions: {},
|
||||
|
||||
getSession: (sessionId) => get().sessions[sessionId] ?? createDefaultSessionState(),
|
||||
|
||||
connectToSession: (sessionId) => {
|
||||
const existing = get().sessions[sessionId]
|
||||
if (existing && existing.connectionState !== 'disconnected') return
|
||||
|
||||
set((s) => ({
|
||||
sessions: {
|
||||
...s.sessions,
|
||||
[sessionId]: {
|
||||
...createDefaultSessionState(),
|
||||
connectionState: 'connecting',
|
||||
messages: existing?.messages ?? [],
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
wsManager.clearHandlers(sessionId)
|
||||
wsManager.connect(sessionId)
|
||||
wsManager.onMessage((msg) => {
|
||||
wsManager.onMessage(sessionId, (msg) => {
|
||||
if (msg.type === 'connected') {
|
||||
set({ connectionState: 'connected' })
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ connectionState: 'connected' })) }))
|
||||
}
|
||||
get().handleServerMessage(msg)
|
||||
get().handleServerMessage(sessionId, msg)
|
||||
})
|
||||
|
||||
// Load history and tasks
|
||||
get().loadHistory(sessionId)
|
||||
useCLITaskStore.getState().fetchSessionTasks(sessionId)
|
||||
sessionsApi.getSlashCommands(sessionId)
|
||||
.then(({ commands }) => {
|
||||
if (get().connectedSessionId === sessionId) {
|
||||
set({ slashCommands: commands })
|
||||
if (get().sessions[sessionId]) {
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ slashCommands: commands })) }))
|
||||
}
|
||||
})
|
||||
.catch(() => {
|
||||
if (get().connectedSessionId === sessionId) {
|
||||
set({ slashCommands: [] })
|
||||
if (get().sessions[sessionId]) {
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ slashCommands: [] })) }))
|
||||
}
|
||||
})
|
||||
},
|
||||
|
||||
disconnectSession: () => {
|
||||
wsManager.disconnect()
|
||||
if (elapsedTimer) clearInterval(elapsedTimer)
|
||||
useCLITaskStore.getState().clearTasks()
|
||||
set({ connectedSessionId: null, chatState: 'idle' })
|
||||
disconnectSession: (sessionId) => {
|
||||
const session = get().sessions[sessionId]
|
||||
if (session?.elapsedTimer) clearInterval(session.elapsedTimer)
|
||||
wsManager.disconnect(sessionId)
|
||||
set((s) => {
|
||||
const { [sessionId]: _, ...rest } = s.sessions
|
||||
return { sessions: rest }
|
||||
})
|
||||
},
|
||||
|
||||
sendMessage: (content: string, attachments?: AttachmentRef[]) => {
|
||||
sendMessage: (sessionId, content, attachments?) => {
|
||||
const userFacingContent = content.trim()
|
||||
const uiAttachments: UIAttachment[] | undefined =
|
||||
attachments && attachments.length > 0
|
||||
? attachments.map((attachment) => ({
|
||||
type: attachment.type,
|
||||
name: attachment.name || attachment.path || attachment.mimeType || attachment.type,
|
||||
data: attachment.data,
|
||||
mimeType: attachment.mimeType,
|
||||
? attachments.map((a) => ({
|
||||
type: a.type,
|
||||
name: a.name || a.path || a.mimeType || a.type,
|
||||
data: a.data,
|
||||
mimeType: a.mimeType,
|
||||
}))
|
||||
: undefined
|
||||
|
||||
// If all tasks are completed, inline the task summary before the new user message
|
||||
const taskStore = useCLITaskStore.getState()
|
||||
const allTasksDone = taskStore.tasks.length > 0 && taskStore.tasks.every((t) => t.status === 'completed')
|
||||
|
||||
// Add user message to UI (with optional task summary before it)
|
||||
set((s) => {
|
||||
const newMessages = [...s.messages]
|
||||
const session = s.sessions[sessionId]
|
||||
if (!session) return s
|
||||
|
||||
const newMessages = [...session.messages]
|
||||
if (allTasksDone) {
|
||||
newMessages.push({
|
||||
id: nextId(),
|
||||
type: 'task_summary',
|
||||
tasks: taskStore.tasks.map((t) => ({
|
||||
id: t.id,
|
||||
subject: t.subject,
|
||||
status: t.status,
|
||||
activeForm: t.activeForm,
|
||||
})),
|
||||
tasks: taskStore.tasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm })),
|
||||
timestamp: Date.now(),
|
||||
})
|
||||
// Clear sticky task bar since we inlined the summary
|
||||
taskStore.clearTasks()
|
||||
}
|
||||
newMessages.push({
|
||||
@ -166,64 +176,66 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
attachments: uiAttachments,
|
||||
timestamp: Date.now(),
|
||||
})
|
||||
|
||||
if (session.elapsedTimer) clearInterval(session.elapsedTimer)
|
||||
|
||||
const timer = setInterval(() => {
|
||||
set((st) => ({ sessions: updateSessionIn(st.sessions, sessionId, (sess) => ({ elapsedSeconds: sess.elapsedSeconds + 1 })) }))
|
||||
}, 1000)
|
||||
|
||||
return {
|
||||
messages: newMessages,
|
||||
chatState: 'thinking',
|
||||
elapsedSeconds: 0,
|
||||
streamingText: '',
|
||||
statusVerb: randomSpinnerVerb(),
|
||||
sessions: {
|
||||
...s.sessions,
|
||||
[sessionId]: {
|
||||
...session,
|
||||
messages: newMessages,
|
||||
chatState: 'thinking',
|
||||
elapsedSeconds: 0,
|
||||
streamingText: '',
|
||||
statusVerb: randomSpinnerVerb(),
|
||||
elapsedTimer: timer,
|
||||
},
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
// Start elapsed timer
|
||||
if (elapsedTimer) clearInterval(elapsedTimer)
|
||||
elapsedTimer = setInterval(() => {
|
||||
set((s) => ({ elapsedSeconds: s.elapsedSeconds + 1 }))
|
||||
}, 1000)
|
||||
|
||||
wsManager.send({ type: 'user_message', content, attachments })
|
||||
wsManager.send(sessionId, { type: 'user_message', content, attachments })
|
||||
},
|
||||
|
||||
respondToPermission: (requestId: string, allowed: boolean, rule?: string) => {
|
||||
wsManager.send({ type: 'permission_response', requestId, allowed, ...(rule ? { rule } : {}) })
|
||||
set({ pendingPermission: null, chatState: allowed ? 'tool_executing' : 'idle' })
|
||||
respondToPermission: (sessionId, requestId, allowed, rule?) => {
|
||||
wsManager.send(sessionId, { type: 'permission_response', requestId, allowed, ...(rule ? { rule } : {}) })
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ pendingPermission: null, chatState: allowed ? 'tool_executing' : 'idle' })) }))
|
||||
},
|
||||
|
||||
setSessionPermissionMode: (mode) => {
|
||||
if (!get().connectedSessionId) return
|
||||
wsManager.send({ type: 'set_permission_mode', mode })
|
||||
setSessionPermissionMode: (sessionId, mode) => {
|
||||
if (!get().sessions[sessionId]) return
|
||||
wsManager.send(sessionId, { type: 'set_permission_mode', mode })
|
||||
},
|
||||
|
||||
stopGeneration: () => {
|
||||
wsManager.send({ type: 'stop_generation' })
|
||||
if (elapsedTimer) clearInterval(elapsedTimer)
|
||||
set({ chatState: 'idle' })
|
||||
stopGeneration: (sessionId) => {
|
||||
wsManager.send(sessionId, { type: 'stop_generation' })
|
||||
set((s) => {
|
||||
const session = s.sessions[sessionId]
|
||||
if (!session) return s
|
||||
if (session.elapsedTimer) clearInterval(session.elapsedTimer)
|
||||
return { sessions: { ...s.sessions, [sessionId]: { ...session, chatState: 'idle', elapsedTimer: null } } }
|
||||
})
|
||||
},
|
||||
|
||||
loadHistory: async (sessionId: string) => {
|
||||
loadHistory: async (sessionId) => {
|
||||
try {
|
||||
const { messages } = await sessionsApi.getMessages(sessionId)
|
||||
const uiMessages = mapHistoryMessagesToUiMessages(messages)
|
||||
|
||||
set((state) => {
|
||||
if (state.connectedSessionId !== sessionId || state.messages.length > 0) {
|
||||
return state
|
||||
}
|
||||
return { ...state, messages: uiMessages }
|
||||
const session = state.sessions[sessionId]
|
||||
if (!session || session.messages.length > 0) return state
|
||||
return { sessions: updateSessionIn(state.sessions, sessionId, () => ({ messages: uiMessages })) }
|
||||
})
|
||||
|
||||
// Extract the last TodoWrite input from history so TaskBar shows for V1 sessions
|
||||
const lastTodos = extractLastTodoWriteFromHistory(messages)
|
||||
if (lastTodos && lastTodos.length > 0) {
|
||||
const taskStore = useCLITaskStore.getState()
|
||||
// Only set if V2 task fetch didn't already populate tasks
|
||||
if (taskStore.tasks.length === 0) {
|
||||
taskStore.setTasksFromTodos(lastTodos)
|
||||
}
|
||||
if (taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos)
|
||||
}
|
||||
|
||||
// For both V1 and V2: if all tasks completed and the user already
|
||||
// continued chatting, suppress the sticky TaskBar on page refresh.
|
||||
if (hasUserMessagesAfterTaskCompletion(messages)) {
|
||||
useCLITaskStore.getState().markCompletedAndDismissed()
|
||||
}
|
||||
@ -232,89 +244,77 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
}
|
||||
},
|
||||
|
||||
clearMessages: () => set({ messages: [], streamingText: '', chatState: 'idle' }),
|
||||
clearMessages: (sessionId) => {
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ messages: [], streamingText: '', chatState: 'idle' })) }))
|
||||
},
|
||||
|
||||
handleServerMessage: (sessionId, msg) => {
|
||||
const update = (updater: (session: PerSessionState) => Partial<PerSessionState>) => {
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, updater) }))
|
||||
}
|
||||
|
||||
handleServerMessage: (msg: ServerMessage) => {
|
||||
switch (msg.type) {
|
||||
case 'connected':
|
||||
break
|
||||
|
||||
case 'status':
|
||||
set({
|
||||
update((session) => ({
|
||||
chatState: msg.state,
|
||||
// Only override statusVerb if the server sends something other than
|
||||
// the generic 'Thinking' — otherwise keep the random verb we picked
|
||||
// in sendMessage so the user sees fun loading text.
|
||||
...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}),
|
||||
...(msg.tokens ? { tokenUsage: { ...get().tokenUsage, output_tokens: msg.tokens } } : {}),
|
||||
...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}),
|
||||
...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}),
|
||||
})
|
||||
if (msg.state === 'idle' && elapsedTimer) {
|
||||
clearInterval(elapsedTimer)
|
||||
elapsedTimer = null
|
||||
}))
|
||||
if (msg.state === 'idle') {
|
||||
const session = get().sessions[sessionId]
|
||||
if (session?.elapsedTimer) {
|
||||
clearInterval(session.elapsedTimer)
|
||||
update(() => ({ elapsedTimer: null }))
|
||||
}
|
||||
}
|
||||
break
|
||||
|
||||
case 'content_start': {
|
||||
// Flush any accumulated streamingText as assistant_text BEFORE
|
||||
// switching to the next block. This preserves intermediate text
|
||||
// segments between tool calls (e.g. "项目已经有 node_modules，直接启动").
|
||||
const pendingText = get().streamingText.trim()
|
||||
const session = get().sessions[sessionId]
|
||||
if (!session) break
|
||||
const pendingText = session.streamingText.trim()
|
||||
if (pendingText) {
|
||||
set((s) => ({
|
||||
messages: [...s.messages, {
|
||||
id: nextId(),
|
||||
type: 'assistant_text',
|
||||
content: pendingText,
|
||||
timestamp: Date.now(),
|
||||
}],
|
||||
update((s) => ({
|
||||
messages: [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }],
|
||||
streamingText: '',
|
||||
}))
|
||||
}
|
||||
|
||||
if (msg.blockType === 'text') {
|
||||
set({ streamingText: '', chatState: 'streaming', activeThinkingId: null })
|
||||
update(() => ({ streamingText: '', chatState: 'streaming', activeThinkingId: null }))
|
||||
} else if (msg.blockType === 'tool_use') {
|
||||
set({
|
||||
update(() => ({
|
||||
activeToolUseId: msg.toolUseId ?? null,
|
||||
activeToolName: msg.toolName ?? null,
|
||||
streamingToolInput: '',
|
||||
chatState: 'tool_executing',
|
||||
activeThinkingId: null,
|
||||
})
|
||||
}))
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
case 'content_delta':
|
||||
if (msg.text !== undefined) {
|
||||
set((s) => ({ streamingText: s.streamingText + msg.text }))
|
||||
}
|
||||
if (msg.toolInput !== undefined) {
|
||||
set((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput }))
|
||||
}
|
||||
if (msg.text !== undefined) update((s) => ({ streamingText: s.streamingText + msg.text }))
|
||||
if (msg.toolInput !== undefined) update((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput }))
|
||||
break
|
||||
|
||||
case 'thinking':
|
||||
// Merge consecutive thinking deltas into one message.
|
||||
// Also flush any pending streamingText first — otherwise the text
|
||||
// becomes invisible because MessageList only renders streamingText
|
||||
// when chatState === 'streaming', and we're about to set it to 'thinking'.
|
||||
set((s) => {
|
||||
update((s) => {
|
||||
const pendingText = s.streamingText.trim()
|
||||
const base = pendingText
|
||||
? [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }]
|
||||
: s.messages
|
||||
|
||||
const last = base[base.length - 1]
|
||||
if (last && last.type === 'thinking') {
|
||||
// Append to existing thinking message
|
||||
const updated = [...base]
|
||||
updated[updated.length - 1] = { ...last, content: last.content + msg.text }
|
||||
return { messages: updated, chatState: 'thinking', activeThinkingId: last.id, streamingText: '' }
|
||||
}
|
||||
const id = nextId()
|
||||
// Create new thinking message
|
||||
return {
|
||||
messages: [...base, { id, type: 'thinking', content: msg.text, timestamp: Date.now() }],
|
||||
chatState: 'thinking',
|
||||
@ -325,49 +325,33 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
break
|
||||
|
||||
case 'tool_use_complete': {
|
||||
const toolName = msg.toolName || get().activeToolName || 'unknown'
|
||||
set((s) => ({
|
||||
const session = get().sessions[sessionId]
|
||||
const toolName = msg.toolName || session?.activeToolName || 'unknown'
|
||||
update((s) => ({
|
||||
messages: [...s.messages, {
|
||||
id: nextId(),
|
||||
type: 'tool_use',
|
||||
toolName,
|
||||
id: nextId(), type: 'tool_use', toolName,
|
||||
toolUseId: msg.toolUseId || s.activeToolUseId || '',
|
||||
input: msg.input,
|
||||
timestamp: Date.now(),
|
||||
parentToolUseId: msg.parentToolUseId,
|
||||
input: msg.input, timestamp: Date.now(), parentToolUseId: msg.parentToolUseId,
|
||||
}],
|
||||
activeToolUseId: null,
|
||||
activeToolName: null,
|
||||
activeThinkingId: null,
|
||||
streamingToolInput: '',
|
||||
activeToolUseId: null, activeToolName: null, activeThinkingId: null, streamingToolInput: '',
|
||||
}))
|
||||
// TodoWrite: input contains the full todo list — update tasks immediately
|
||||
if (toolName === 'TodoWrite' && Array.isArray((msg.input as any)?.todos)) {
|
||||
useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos)
|
||||
} else if (TASK_TOOL_NAMES.has(toolName)) {
|
||||
// V2 task tools — refresh will happen on tool_result
|
||||
// when the tool has actually finished executing and written to disk
|
||||
const useId = msg.toolUseId || get().activeToolUseId
|
||||
const useId = msg.toolUseId || session?.activeToolUseId
|
||||
if (useId) pendingTaskToolUseIds.add(useId)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
case 'tool_result':
|
||||
set((s) => ({
|
||||
update((s) => ({
|
||||
messages: [...s.messages, {
|
||||
id: nextId(),
|
||||
type: 'tool_result',
|
||||
toolUseId: msg.toolUseId,
|
||||
content: msg.content,
|
||||
isError: msg.isError,
|
||||
timestamp: Date.now(),
|
||||
parentToolUseId: msg.parentToolUseId,
|
||||
id: nextId(), type: 'tool_result', toolUseId: msg.toolUseId,
|
||||
content: msg.content, isError: msg.isError, timestamp: Date.now(), parentToolUseId: msg.parentToolUseId,
|
||||
}],
|
||||
chatState: 'thinking',
|
||||
activeThinkingId: null,
|
||||
chatState: 'thinking', activeThinkingId: null,
|
||||
}))
|
||||
// Refresh tasks after a task tool has finished executing
|
||||
if (pendingTaskToolUseIds.has(msg.toolUseId)) {
|
||||
pendingTaskToolUseIds.delete(msg.toolUseId)
|
||||
useCLITaskStore.getState().refreshTasks()
|
||||
@ -375,230 +359,116 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
break
|
||||
|
||||
case 'permission_request':
|
||||
set({
|
||||
pendingPermission: {
|
||||
requestId: msg.requestId,
|
||||
toolName: msg.toolName,
|
||||
input: msg.input,
|
||||
description: msg.description,
|
||||
},
|
||||
update((s) => ({
|
||||
pendingPermission: { requestId: msg.requestId, toolName: msg.toolName, input: msg.input, description: msg.description },
|
||||
chatState: 'permission_pending',
|
||||
activeThinkingId: null,
|
||||
})
|
||||
set((s) => ({
|
||||
messages: [...s.messages, {
|
||||
id: nextId(),
|
||||
type: 'permission_request',
|
||||
requestId: msg.requestId,
|
||||
toolName: msg.toolName,
|
||||
input: msg.input,
|
||||
description: msg.description,
|
||||
timestamp: Date.now(),
|
||||
id: nextId(), type: 'permission_request', requestId: msg.requestId,
|
||||
toolName: msg.toolName, input: msg.input, description: msg.description, timestamp: Date.now(),
|
||||
}],
|
||||
}))
|
||||
break
|
||||
|
||||
case 'message_complete': {
|
||||
// Flush streaming text as a message
|
||||
const text = get().streamingText
|
||||
const session = get().sessions[sessionId]
|
||||
if (!session) break
|
||||
const text = session.streamingText
|
||||
if (text) {
|
||||
set((s) => ({
|
||||
update((s) => ({
|
||||
messages: [...s.messages, { id: nextId(), type: 'assistant_text', content: text, timestamp: Date.now() }],
|
||||
streamingText: '',
|
||||
}))
|
||||
}
|
||||
set({ tokenUsage: msg.usage, chatState: 'idle', activeThinkingId: null })
|
||||
if (elapsedTimer) {
|
||||
clearInterval(elapsedTimer)
|
||||
elapsedTimer = null
|
||||
}
|
||||
if (session.elapsedTimer) clearInterval(session.elapsedTimer)
|
||||
update(() => ({ tokenUsage: msg.usage, chatState: 'idle', activeThinkingId: null, elapsedTimer: null }))
|
||||
break
|
||||
}
|
||||
|
||||
case 'error':
|
||||
set((s) => ({
|
||||
update((s) => ({
|
||||
messages: [...s.messages, { id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }],
|
||||
chatState: 'idle',
|
||||
activeThinkingId: null,
|
||||
chatState: 'idle', activeThinkingId: null,
|
||||
}))
|
||||
if (elapsedTimer) {
|
||||
clearInterval(elapsedTimer)
|
||||
elapsedTimer = null
|
||||
{
|
||||
const session = get().sessions[sessionId]
|
||||
if (session?.elapsedTimer) {
|
||||
clearInterval(session.elapsedTimer)
|
||||
update(() => ({ elapsedTimer: null }))
|
||||
}
|
||||
}
|
||||
break
|
||||
|
||||
case 'team_created':
|
||||
useTeamStore.getState().handleTeamCreated(msg.teamName)
|
||||
break
|
||||
|
||||
case 'team_update':
|
||||
useTeamStore.getState().handleTeamUpdate(msg.teamName, msg.members)
|
||||
break
|
||||
|
||||
case 'team_deleted':
|
||||
useTeamStore.getState().handleTeamDeleted(msg.teamName)
|
||||
break
|
||||
|
||||
case 'task_update':
|
||||
break
|
||||
|
||||
case 'session_title_updated':
|
||||
useSessionStore.getState().updateSessionTitle(msg.sessionId, msg.title)
|
||||
break
|
||||
|
||||
case 'system_notification':
|
||||
// Cache slash commands from CLI init
|
||||
if (msg.subtype === 'slash_commands' && Array.isArray(msg.data)) {
|
||||
set({ slashCommands: msg.data as Array<{ name: string; description: string }> })
|
||||
update(() => ({ slashCommands: msg.data as Array<{ name: string; description: string }> }))
|
||||
}
|
||||
break
|
||||
|
||||
case 'pong':
|
||||
break
|
||||
}
|
||||
},
|
||||
}))
|
||||
|
||||
type AssistantHistoryBlock = {
|
||||
type: string
|
||||
text?: string
|
||||
thinking?: string
|
||||
name?: string
|
||||
id?: string
|
||||
input?: unknown
|
||||
}
|
||||
// ─── History mapping helpers (unchanged from original) ─────────
|
||||
|
||||
type UserHistoryBlock = {
|
||||
type: string
|
||||
text?: string
|
||||
tool_use_id?: string
|
||||
content?: unknown
|
||||
is_error?: boolean
|
||||
source?: { data?: string }
|
||||
mimeType?: string
|
||||
media_type?: string
|
||||
name?: string
|
||||
}
|
||||
type AssistantHistoryBlock = { type: string; text?: string; thinking?: string; name?: string; id?: string; input?: unknown }
|
||||
type UserHistoryBlock = { type: string; text?: string; tool_use_id?: string; content?: unknown; is_error?: boolean; source?: { data?: string }; mimeType?: string; media_type?: string; name?: string }
|
||||
|
||||
export function mapHistoryMessagesToUiMessages(messages: MessageEntry[]): UIMessage[] {
|
||||
const uiMessages: UIMessage[] = []
|
||||
|
||||
for (const msg of messages) {
|
||||
const timestamp = new Date(msg.timestamp).getTime()
|
||||
|
||||
if (msg.type === 'user' && typeof msg.content === 'string') {
|
||||
uiMessages.push({
|
||||
id: msg.id || nextId(),
|
||||
type: 'user_text',
|
||||
content: msg.content,
|
||||
timestamp,
|
||||
})
|
||||
uiMessages.push({ id: msg.id || nextId(), type: 'user_text', content: msg.content, timestamp })
|
||||
continue
|
||||
}
|
||||
|
||||
if (msg.type === 'assistant' && typeof msg.content === 'string') {
|
||||
uiMessages.push({
|
||||
id: msg.id || nextId(),
|
||||
type: 'assistant_text',
|
||||
content: msg.content,
|
||||
timestamp,
|
||||
model: msg.model,
|
||||
})
|
||||
uiMessages.push({ id: msg.id || nextId(), type: 'assistant_text', content: msg.content, timestamp, model: msg.model })
|
||||
continue
|
||||
}
|
||||
|
||||
// Server marks assistant messages containing tool_use blocks as type 'tool_use',
|
||||
// but the content array structure is the same as 'assistant' — handle both.
|
||||
if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
|
||||
for (const block of msg.content as AssistantHistoryBlock[]) {
|
||||
if (block.type === 'thinking' && block.thinking) {
|
||||
uiMessages.push({
|
||||
id: nextId(),
|
||||
type: 'thinking',
|
||||
content: block.thinking,
|
||||
timestamp,
|
||||
})
|
||||
} else if (block.type === 'text' && block.text) {
|
||||
uiMessages.push({
|
||||
id: nextId(),
|
||||
type: 'assistant_text',
|
||||
content: block.text,
|
||||
timestamp,
|
||||
model: msg.model,
|
||||
})
|
||||
} else if (block.type === 'tool_use') {
|
||||
uiMessages.push({
|
||||
id: nextId(),
|
||||
type: 'tool_use',
|
||||
toolName: block.name ?? 'unknown',
|
||||
toolUseId: block.id ?? '',
|
||||
input: block.input,
|
||||
timestamp,
|
||||
parentToolUseId: msg.parentToolUseId,
|
||||
})
|
||||
}
|
||||
if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp })
|
||||
else if (block.type === 'text' && block.text) uiMessages.push({ id: nextId(), type: 'assistant_text', content: block.text, timestamp, model: msg.model })
|
||||
else if (block.type === 'tool_use') uiMessages.push({ id: nextId(), type: 'tool_use', toolName: block.name ?? 'unknown', toolUseId: block.id ?? '', input: block.input, timestamp, parentToolUseId: msg.parentToolUseId })
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Server marks user messages containing tool_result blocks as type 'tool_result',
|
||||
// but the content array structure is the same as 'user' — handle both.
|
||||
if ((msg.type === 'user' || msg.type === 'tool_result') && Array.isArray(msg.content)) {
|
||||
const textParts: string[] = []
|
||||
const attachments: UIAttachment[] = []
|
||||
|
||||
for (const block of msg.content as UserHistoryBlock[]) {
|
||||
if (block.type === 'text' && block.text) {
|
||||
textParts.push(block.text)
|
||||
} else if (block.type === 'image') {
|
||||
attachments.push({
|
||||
type: 'image',
|
||||
name: block.name || 'image',
|
||||
data: block.source?.data,
|
||||
mimeType: block.mimeType || block.media_type,
|
||||
})
|
||||
} else if (block.type === 'file') {
|
||||
attachments.push({
|
||||
type: 'file',
|
||||
name: block.name || 'file',
|
||||
})
|
||||
} else if (block.type === 'tool_result') {
|
||||
uiMessages.push({
|
||||
id: nextId(),
|
||||
type: 'tool_result',
|
||||
toolUseId: block.tool_use_id ?? '',
|
||||
content: block.content,
|
||||
isError: !!block.is_error,
|
||||
timestamp,
|
||||
parentToolUseId: msg.parentToolUseId,
|
||||
})
|
||||
}
|
||||
if (block.type === 'text' && block.text) textParts.push(block.text)
|
||||
else if (block.type === 'image') attachments.push({ type: 'image', name: block.name || 'image', data: block.source?.data, mimeType: block.mimeType || block.media_type })
|
||||
else if (block.type === 'file') attachments.push({ type: 'file', name: block.name || 'file' })
|
||||
else if (block.type === 'tool_result') uiMessages.push({ id: nextId(), type: 'tool_result', toolUseId: block.tool_use_id ?? '', content: block.content, isError: !!block.is_error, timestamp, parentToolUseId: msg.parentToolUseId })
|
||||
}
|
||||
|
||||
if (textParts.length > 0 || attachments.length > 0) {
|
||||
uiMessages.push({
|
||||
id: nextId(),
|
||||
type: 'user_text',
|
||||
content: textParts.join('\n'),
|
||||
attachments: attachments.length > 0 ? attachments : undefined,
|
||||
timestamp,
|
||||
})
|
||||
uiMessages.push({ id: nextId(), type: 'user_text', content: textParts.join('\n'), attachments: attachments.length > 0 ? attachments : undefined, timestamp })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return uiMessages
|
||||
}
|
||||
|
||||
/** Scan history messages for the last TodoWrite tool_use and return the todos array.
|
||||
* Returns null if all todos were completed and the user continued chatting after. */
|
||||
function extractLastTodoWriteFromHistory(
|
||||
messages: MessageEntry[],
|
||||
): Array<{ content: string; status: string; activeForm?: string }> | null {
|
||||
function extractLastTodoWriteFromHistory(messages: MessageEntry[]): Array<{ content: string; status: string; activeForm?: string }> | null {
|
||||
let foundIndex = -1
|
||||
let todos: Array<{ content: string; status: string; activeForm?: string }> | null = null
|
||||
|
||||
// Walk backwards to find the most recent TodoWrite
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const msg = messages[i]!
|
||||
if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
|
||||
@ -617,45 +487,28 @@ function extractLastTodoWriteFromHistory(
|
||||
if (todos) break
|
||||
}
|
||||
}
|
||||
|
||||
if (!todos) return null
|
||||
|
||||
// If all todos are completed and there are user messages after, the tasks
|
||||
// were already inlined — don't re-show the sticky bar.
|
||||
const allDone = todos.every((t) => t.status === 'completed')
|
||||
if (allDone) {
|
||||
for (let i = foundIndex + 1; i < messages.length; i++) {
|
||||
const msg = messages[i]!
|
||||
if (msg.type === 'user' && msg.content) return null
|
||||
if (messages[i]!.type === 'user' && messages[i]!.content) return null
|
||||
}
|
||||
}
|
||||
|
||||
return todos
|
||||
}
|
||||
|
||||
const TASK_RELATED_TOOL_NAMES = new Set(['TodoWrite', 'TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList'])
|
||||
|
||||
/** Check if there are user messages after the last task-related tool call.
|
||||
* Used to determine if completed tasks should be shown as sticky or inline. */
|
||||
function hasUserMessagesAfterTaskCompletion(messages: MessageEntry[]): boolean {
|
||||
// Find the index of the last task-related tool call
|
||||
let lastTaskIndex = -1
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const msg = messages[i]!
|
||||
if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
|
||||
const blocks = msg.content as AssistantHistoryBlock[]
|
||||
if (blocks.some((b) => b.type === 'tool_use' && TASK_RELATED_TOOL_NAMES.has(b.name ?? ''))) {
|
||||
lastTaskIndex = i
|
||||
break
|
||||
}
|
||||
if (blocks.some((b) => b.type === 'tool_use' && TASK_RELATED_TOOL_NAMES.has(b.name ?? ''))) { lastTaskIndex = i; break }
|
||||
}
|
||||
}
|
||||
|
||||
if (lastTaskIndex < 0) return false
|
||||
|
||||
// Check for user messages after the last task tool
|
||||
for (let i = lastTaskIndex + 1; i < messages.length; i++) {
|
||||
if (messages[i]!.type === 'user') return true
|
||||
}
|
||||
for (let i = lastTaskIndex + 1; i < messages.length; i++) { if (messages[i]!.type === 'user') return true }
|
||||
return false
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user