462 lines
18 KiB
TypeScript
462 lines
18 KiB
TypeScript
import { create } from 'zustand'
|
||
import type { Room, Message, AgentInfo, SkillMeta, WsEvent, TodoItem } from './types'
|
||
|
||
interface TokenUsage {
|
||
promptTokens: number
|
||
completionTokens: number
|
||
totalTokens: number
|
||
}
|
||
|
||
interface AppState {
|
||
theme: 'light' | 'dark'
|
||
user: {
|
||
name: string
|
||
avatar_color: string
|
||
} | null
|
||
onboardingCompleted: boolean
|
||
rooms: Room[]
|
||
activeRoomId: string | null
|
||
messages: Record<string, Message[]>
|
||
tasks: Record<string, string>
|
||
workspace: Record<string, string[]>
|
||
agents: AgentInfo[]
|
||
skills: SkillMeta[]
|
||
page: 'chat' | 'agents' | 'skills' | 'market' | 'settings'
|
||
ws: Record<string, WebSocket>
|
||
tokenUsage: Record<string, TokenUsage>
|
||
todoItems: Record<string, TodoItem[]>
|
||
fileReaders: Record<string, Record<string, string[]>>
|
||
workingFiles: Record<string, { agent: string; filename: string; title?: string }[]>
|
||
workStartedAt: Record<string, number> // roomId -> timestamp (ms)
|
||
|
||
setTheme: (theme: 'light' | 'dark') => void
|
||
toggleTheme: () => void
|
||
setOnboardingCompleted: (completed: boolean) => void
|
||
fetchUser: () => Promise<void>
|
||
setPage: (p: AppState['page']) => void
|
||
setActiveRoom: (id: string) => void
|
||
fetchRooms: () => Promise<void>
|
||
fetchAgents: () => Promise<void>
|
||
fetchSkills: () => Promise<void>
|
||
sendMessage: (roomId: string, content: string) => void
|
||
connectRoom: (roomId: string) => void
|
||
markFileRead: (roomId: string, filename: string) => void
|
||
setRoomTeam: (roomId: string, team: string) => Promise<void>
|
||
}
|
||
|
||
const API = '/api'
|
||
|
||
export const useStore = create<AppState>((set, get) => {
|
||
const getInitialTheme = (): 'light' | 'dark' => {
|
||
if (typeof window !== 'undefined') {
|
||
const saved = localStorage.getItem('theme')
|
||
if (saved === 'light' || saved === 'dark') return saved
|
||
return window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light'
|
||
}
|
||
return 'dark'
|
||
}
|
||
|
||
const applyTheme = (theme: 'light' | 'dark') => {
|
||
if (typeof window !== 'undefined') {
|
||
document.documentElement.classList.remove('light', 'dark')
|
||
document.documentElement.classList.add(theme)
|
||
localStorage.setItem('theme', theme)
|
||
}
|
||
}
|
||
|
||
const initialTheme = getInitialTheme()
|
||
applyTheme(initialTheme)
|
||
|
||
return {
|
||
theme: initialTheme,
|
||
user: null,
|
||
onboardingCompleted: typeof window !== 'undefined' ? localStorage.getItem('onboarding_completed') === 'true' : false,
|
||
rooms: [],
|
||
activeRoomId: typeof window !== 'undefined' ? localStorage.getItem('activeRoomId') : null,
|
||
messages: {},
|
||
tasks: {},
|
||
workspace: {},
|
||
agents: [],
|
||
skills: [],
|
||
page: (typeof window !== 'undefined' ? localStorage.getItem('page') as AppState['page'] : null) || 'chat',
|
||
ws: {},
|
||
tokenUsage: {},
|
||
todoItems: {},
|
||
fileReaders: {},
|
||
workingFiles: {},
|
||
workStartedAt: {},
|
||
|
||
setTheme: (theme) => {
|
||
applyTheme(theme)
|
||
set({ theme })
|
||
},
|
||
toggleTheme: () => {
|
||
const newTheme = get().theme === 'dark' ? 'light' : 'dark'
|
||
applyTheme(newTheme)
|
||
set({ theme: newTheme })
|
||
},
|
||
setOnboardingCompleted: (completed) => {
|
||
localStorage.setItem('onboarding_completed', completed ? 'true' : 'false')
|
||
set({ onboardingCompleted: completed })
|
||
},
|
||
fetchUser: async () => {
|
||
try {
|
||
const res = await fetch(`${API}/user`)
|
||
const data = await res.json()
|
||
set({
|
||
user: {
|
||
name: data.name || '用户',
|
||
avatar_color: data.avatar_color || '#5865F2'
|
||
}
|
||
})
|
||
} catch {
|
||
set({
|
||
user: {
|
||
name: '用户',
|
||
avatar_color: '#5865F2'
|
||
}
|
||
})
|
||
}
|
||
},
|
||
setPage: (page) => {
|
||
localStorage.setItem('page', page)
|
||
set({ page })
|
||
},
|
||
setActiveRoom: (id) => {
|
||
localStorage.setItem('activeRoomId', id)
|
||
set({ activeRoomId: id })
|
||
get().connectRoom(id)
|
||
fetch(`${API}/rooms/${id}/tasks`).then(r => r.json()).then(d => {
|
||
set(s => ({ tasks: { ...s.tasks, [id]: d.content } }))
|
||
})
|
||
fetch(`${API}/rooms/${id}/workspace`).then(r => r.json()).then(files => {
|
||
set(s => ({ workspace: { ...s.workspace, [id]: files || [] } }))
|
||
})
|
||
fetch(`${API}/rooms/${id}/messages`).then(r => r.json()).then(msgs => {
|
||
if (msgs?.length) set(s => ({ messages: { ...s.messages, [id]: msgs } }))
|
||
}).catch(() => {})
|
||
fetch(`${API}/rooms/${id}/token-usage`).then(r => r.json()).then(data => {
|
||
if (data) set(s => ({
|
||
tokenUsage: { ...s.tokenUsage, [id]: {
|
||
promptTokens: data.prompt_tokens || 0,
|
||
completionTokens: data.completion_tokens || 0,
|
||
totalTokens: data.total_tokens || 0,
|
||
}}
|
||
}))
|
||
}).catch(() => {})
|
||
},
|
||
|
||
fetchRooms: async () => {
|
||
const rooms = await fetch(`${API}/rooms`).then(r => r.json())
|
||
set({ rooms: rooms || [] })
|
||
},
|
||
|
||
fetchAgents: async () => {
|
||
const agents = await fetch(`${API}/agents`).then(r => r.json())
|
||
set({ agents: agents || [] })
|
||
},
|
||
|
||
fetchSkills: async () => {
|
||
const skills = await fetch(`${API}/skills`).then(r => r.json())
|
||
set({ skills: skills || [] })
|
||
},
|
||
|
||
connectRoom: (roomId) => {
|
||
const { ws } = get()
|
||
// 已有连接且状态正常则跳过
|
||
if (ws[roomId] && ws[roomId].readyState <= 1) return
|
||
const socket = new WebSocket(`ws://${location.host}/ws/${roomId}`)
|
||
|
||
// 断线自动重连
|
||
socket.onclose = () => {
|
||
set(s => {
|
||
const newWs = { ...s.ws }
|
||
delete newWs[roomId]
|
||
return { ws: newWs }
|
||
})
|
||
// 3 秒后重连
|
||
setTimeout(() => get().connectRoom(roomId), 3000)
|
||
}
|
||
|
||
// 心跳保活：每 30 秒发一次 ping
|
||
const heartbeat = setInterval(() => {
|
||
if (socket.readyState === WebSocket.OPEN) {
|
||
socket.send(JSON.stringify({ type: 'ping' }))
|
||
}
|
||
}, 30000)
|
||
socket.addEventListener('close', () => clearInterval(heartbeat))
|
||
|
||
socket.onmessage = (e) => {
|
||
const ev: WsEvent = JSON.parse(e.data)
|
||
if (ev.type === 'agent_message') {
|
||
set(s => {
|
||
const msgs = [...(s.messages[roomId] || [])]
|
||
|
||
// tool_use 消息：同一个 agent 只保留一条，持续更新状态
|
||
if (ev.role === 'tool_use') {
|
||
let idx = -1
|
||
for (let i = msgs.length - 1; i >= 0; i--) {
|
||
if (msgs[i].role === 'tool_use' && msgs[i].agent === ev.agent) { idx = i; break }
|
||
}
|
||
if (idx !== -1) {
|
||
msgs[idx] = { ...msgs[idx], content: ev.content || msgs[idx].content, streaming: ev.streaming }
|
||
} else if (ev.content) {
|
||
msgs.push({ id: Date.now().toString(), agent: ev.agent, role: 'tool_use', content: ev.content, streaming: ev.streaming })
|
||
}
|
||
// 同步更新 todoItem 的 action 描述，同时将 assigned 转为 working
|
||
const actionText = ev.streaming
|
||
? (ev.content === 'thinking' ? '正在整理结果...' : '正在搜索...')
|
||
: undefined
|
||
const todos = (s.todoItems[roomId] || []).map(t =>
|
||
t.agent === ev.agent && (t.status === 'working' || t.status === 'assigned')
|
||
? { ...t, status: 'working' as const, action: actionText || t.action }
|
||
: t
|
||
)
|
||
return {
|
||
messages: { ...s.messages, [roomId]: msgs },
|
||
todoItems: { ...s.todoItems, [roomId]: todos },
|
||
}
|
||
}
|
||
|
||
// 普通消息：streaming 合并逻辑
|
||
let idx = -1
|
||
for (let i = msgs.length - 1; i >= 0; i--) {
|
||
if (msgs[i].streaming && msgs[i].agent === ev.agent) { idx = i; break }
|
||
}
|
||
if (idx !== -1) {
|
||
if (ev.action === 'replace') {
|
||
// action=replace: 用新内容替换（文件驱动模式，文档已保存到 workspace）
|
||
msgs[idx] = { ...msgs[idx], content: ev.content || '', streaming: false }
|
||
} else {
|
||
msgs[idx] = { ...msgs[idx], content: msgs[idx].content + (ev.content || ''), streaming: ev.streaming }
|
||
}
|
||
} else if (ev.action === 'replace') {
|
||
// 找最后一条同 agent 消息替换，找不到则新增
|
||
let found = false
|
||
for (let i = msgs.length - 1; i >= 0; i--) {
|
||
if (msgs[i].agent === ev.agent && msgs[i].role === ev.role) {
|
||
msgs[i] = { ...msgs[i], content: ev.content || '', streaming: false }
|
||
found = true
|
||
break
|
||
}
|
||
}
|
||
if (!found && ev.content) {
|
||
msgs.push({ id: Date.now().toString(), agent: ev.agent, role: ev.role, content: ev.content, streaming: false })
|
||
}
|
||
} else if (ev.content) {
|
||
msgs.push({ id: Date.now().toString(), agent: ev.agent, role: ev.role, content: ev.content, streaming: ev.streaming, created_at: new Date().toISOString() })
|
||
}
|
||
return { messages: { ...s.messages, [roomId]: msgs } }
|
||
})
|
||
} else if (ev.type === 'room_status') {
|
||
set(s => {
|
||
let todos = [...(s.todoItems[roomId] || [])]
|
||
if (ev.active_agent && ev.status === 'working') {
|
||
// 更新活跃 agent 状态为 working，附带 action 描述
|
||
todos = todos.map(t =>
|
||
t.agent === ev.active_agent
|
||
? { ...t, status: 'working' as const, action: ev.action || t.action }
|
||
: t
|
||
)
|
||
}
|
||
// 记录工作开始时间（thinking/working 时设置，pending 时清除）
|
||
const workStartedAt = { ...s.workStartedAt }
|
||
if (ev.status === 'pending') {
|
||
delete workStartedAt[roomId]
|
||
} else if (!workStartedAt[roomId]) {
|
||
workStartedAt[roomId] = Date.now()
|
||
}
|
||
// pending 时：清理所有残留 streaming 消息，只在所有任务都 done 后清空 todos
|
||
if (ev.status === 'pending') {
|
||
const cleanMsgs = (s.messages[roomId] || []).map(m =>
|
||
m.streaming ? { ...m, streaming: false } : m
|
||
)
|
||
if (todos.length > 0) {
|
||
const allDone = todos.every(t => t.status === 'done')
|
||
if (allDone) todos = []
|
||
}
|
||
return {
|
||
rooms: s.rooms.map(r => r.id === roomId
|
||
? { ...r, status: ev.status, activeAgent: ev.active_agent, action: ev.action }
|
||
: r),
|
||
messages: { ...s.messages, [roomId]: cleanMsgs },
|
||
todoItems: { ...s.todoItems, [roomId]: todos },
|
||
workingFiles: { ...s.workingFiles, [roomId]: [] },
|
||
workStartedAt,
|
||
}
|
||
}
|
||
return {
|
||
rooms: s.rooms.map(r => r.id === roomId
|
||
? { ...r, status: ev.status, activeAgent: ev.active_agent, action: ev.action }
|
||
: r),
|
||
todoItems: { ...s.todoItems, [roomId]: todos },
|
||
workStartedAt,
|
||
}
|
||
})
|
||
} else if (ev.type === 'tasks_update') {
|
||
set(s => ({ tasks: { ...s.tasks, [roomId]: ev.content } }))
|
||
} else if (ev.type === 'workspace_file') {
|
||
set(s => {
|
||
const files = [...(s.workspace[roomId] || [])]
|
||
if (!files.includes(ev.filename)) {
|
||
files.push(ev.filename)
|
||
}
|
||
return { workspace: { ...s.workspace, [roomId]: files } }
|
||
})
|
||
} else if (ev.type === 'task_assign') {
|
||
set(s => {
|
||
const msgs = [...(s.messages[roomId] || [])]
|
||
msgs.push({
|
||
id: Date.now().toString(),
|
||
agent: ev.from,
|
||
role: 'task_assign' as const,
|
||
content: ev.task,
|
||
title: ev.to,
|
||
})
|
||
// 去重：同一 agent 只保留最新任务，清除已完成的
|
||
const existing = (s.todoItems[roomId] || []).filter(t => t.status !== 'done' && t.agent !== ev.to)
|
||
existing.push({
|
||
id: `${Date.now()}-${ev.to}`,
|
||
agent: ev.to,
|
||
task: ev.task,
|
||
status: 'assigned',
|
||
})
|
||
return {
|
||
messages: { ...s.messages, [roomId]: msgs },
|
||
todoItems: { ...s.todoItems, [roomId]: existing },
|
||
}
|
||
})
|
||
} else if (ev.type === 'task_done') {
|
||
set(s => {
|
||
// 清理该 agent 的残留 streaming 消息
|
||
const msgs = (s.messages[roomId] || []).map(m =>
|
||
m.agent === ev.agent && m.streaming ? { ...m, streaming: false } : m
|
||
)
|
||
const todos = (s.todoItems[roomId] || []).map(t =>
|
||
t.agent === ev.agent && t.status !== 'done'
|
||
? { ...t, status: 'done' as const }
|
||
: t
|
||
)
|
||
return {
|
||
messages: { ...s.messages, [roomId]: msgs },
|
||
todoItems: { ...s.todoItems, [roomId]: todos },
|
||
}
|
||
})
|
||
} else if (ev.type === 'mode_change') {
|
||
set(s => ({
|
||
rooms: s.rooms.map(r => r.id === roomId ? { ...r, mode: ev.mode } : r)
|
||
}))
|
||
} else if (ev.type === 'schedule_run') {
|
||
set(s => {
|
||
const msgs = [...(s.messages[roomId] || [])]
|
||
msgs.push({
|
||
id: Date.now().toString(),
|
||
agent: ev.agent || 'scheduler',
|
||
role: 'user' as const,
|
||
content: `[定时任务] ${ev.content}`,
|
||
})
|
||
return { messages: { ...s.messages, [roomId]: msgs } }
|
||
})
|
||
} else if (ev.type === 'token_usage') {
|
||
set(s => {
|
||
const prev = s.tokenUsage[roomId] || { promptTokens: 0, completionTokens: 0, totalTokens: 0 }
|
||
return {
|
||
tokenUsage: {
|
||
...s.tokenUsage,
|
||
[roomId]: {
|
||
promptTokens: prev.promptTokens + ev.prompt_tokens,
|
||
completionTokens: prev.completionTokens + ev.completion_tokens,
|
||
totalTokens: prev.totalTokens + ev.total_tokens,
|
||
}
|
||
}
|
||
}
|
||
})
|
||
} else if (ev.type === 'artifact') {
|
||
// 添加 artifact 消息卡片到聊天
|
||
set(s => {
|
||
const msgs = [...(s.messages[roomId] || [])]
|
||
msgs.push({
|
||
id: Date.now().toString(),
|
||
agent: ev.agent,
|
||
role: 'artifact' as const,
|
||
content: ev.title || ev.filename,
|
||
filename: ev.filename,
|
||
title: ev.title,
|
||
})
|
||
// 同时更新 workspace 文件列表
|
||
const files = [...(s.workspace[roomId] || [])]
|
||
if (!files.includes(ev.filename)) {
|
||
files.push(ev.filename)
|
||
}
|
||
return {
|
||
messages: { ...s.messages, [roomId]: msgs },
|
||
workspace: { ...s.workspace, [roomId]: files },
|
||
}
|
||
})
|
||
} else if (ev.type === 'file_working') {
|
||
set(s => {
|
||
const files = [...(s.workingFiles[roomId] || [])]
|
||
// 避免重复
|
||
if (!files.some(f => f.agent === ev.agent && f.filename === ev.filename)) {
|
||
files.push({ agent: ev.agent, filename: ev.filename, title: ev.title })
|
||
}
|
||
return { workingFiles: { ...s.workingFiles, [roomId]: files } }
|
||
})
|
||
} else if (ev.type === 'file_done') {
|
||
set(s => {
|
||
const files = (s.workingFiles[roomId] || []).filter(
|
||
f => !(f.agent === ev.agent && f.filename === ev.filename)
|
||
)
|
||
return { workingFiles: { ...s.workingFiles, [roomId]: files } }
|
||
})
|
||
} else if (ev.type === 'file_read') {
|
||
set(s => {
|
||
const roomReaders = { ...(s.fileReaders[roomId] || {}) }
|
||
const readers = [...(roomReaders[ev.filename] || [])]
|
||
if (!readers.includes(ev.agent)) {
|
||
readers.push(ev.agent)
|
||
}
|
||
roomReaders[ev.filename] = readers
|
||
return { fileReaders: { ...s.fileReaders, [roomId]: roomReaders } }
|
||
})
|
||
}
|
||
}
|
||
set(s => ({ ws: { ...s.ws, [roomId]: socket } }))
|
||
},
|
||
|
||
sendMessage: (roomId, content) => {
|
||
const { ws, user } = get()
|
||
const userName = user?.name || '用户'
|
||
const userMsg = { id: Date.now().toString(), agent: userName, role: 'user' as const, content, created_at: new Date().toISOString() }
|
||
set(s => ({ messages: { ...s.messages, [roomId]: [...(s.messages[roomId] || []), userMsg] } }))
|
||
ws[roomId]?.send(JSON.stringify({ type: 'user_message', content, user_name: userName }))
|
||
},
|
||
|
||
markFileRead: (roomId: string, filename: string) => {
|
||
const userName = get().user?.name || '用户'
|
||
set(s => {
|
||
const roomReaders = { ...(s.fileReaders[roomId] || {}) }
|
||
const readers = [...(roomReaders[filename] || [])]
|
||
if (!readers.includes(userName)) {
|
||
readers.push(userName)
|
||
}
|
||
roomReaders[filename] = readers
|
||
return { fileReaders: { ...s.fileReaders, [roomId]: roomReaders } }
|
||
})
|
||
},
|
||
|
||
setRoomTeam: async (roomId, team) => {
|
||
const res = await fetch(`${API}/rooms/${roomId}/team`, {
|
||
method: 'PUT', headers: { 'Content-Type': 'application/json' },
|
||
body: JSON.stringify({ team })
|
||
})
|
||
if (res.ok) {
|
||
const data = await res.json()
|
||
set(s => ({
|
||
rooms: s.rooms.map(r => r.id === roomId ? { ...r, team, master: data.master, members: data.members } : r)
|
||
}))
|
||
}
|
||
},
|
||
}
|
||
})
|