| 1 |
import { buildToolMessages } from '@agent/lib/tool-messages'; |
| 2 |
import { makeId } from '@agent/lib/util'; |
| 3 |
import { useStatusStore } from '@agent/state/status'; |
| 4 |
import apiFetch from '@wordpress/api-fetch'; |
| 5 |
import { create } from 'zustand'; |
| 6 |
import { createJSONStorage, devtools, persist } from 'zustand/middleware'; |
| 7 |
|
| 8 |
const { chatHistory } = window.extAgentData; |
| 9 |
|
| 10 |
const state = (set, get) => ({ |
| 11 |
messages: chatHistory?.length ? chatHistory.toReversed() : [], |
| 12 |
// API messages, back to the last finished workflow. |
| 13 |
getCurrentMessages: ({ includeTools = true } = {}) => { |
| 14 |
const messages = []; |
| 15 |
let foundUserMessage = false; |
| 16 |
for (const { type, details } of get().messages.toReversed()) { |
| 17 |
const finished = ['completed', 'canceled'].includes(details.status); |
| 18 |
if (type === 'workflow' && finished) break; |
| 19 |
if (type === 'workflow-component' && finished) break; |
| 20 |
if (type === 'message' && details.error) break; |
| 21 |
// A call with no result leaves the model answering a phantom. |
| 22 |
if (type === 'tool' && !('result' in details)) continue; |
| 23 |
if (type === 'tool' && includeTools) { |
| 24 |
// buildToolMessages returns [call, result]; push reversed so the |
| 25 |
// final toReversed() restores call-before-result order. |
| 26 |
for (const m of buildToolMessages(details).toReversed()) |
| 27 |
messages.push(m); |
| 28 |
} |
| 29 |
// This prevents a loop of assistant messages from being at the end |
| 30 |
if (type === 'message' && details.role === 'user') { |
| 31 |
foundUserMessage = true; |
| 32 |
} |
| 33 |
if (type === 'message' && !foundUserMessage) continue; |
| 34 |
if (type === 'message') messages.push(details); |
| 35 |
} |
| 36 |
return messages.toReversed(); |
| 37 |
}, |
| 38 |
// Most recent consecutive runs — another workflow in between resets the |
| 39 |
// memory. |
| 40 |
getMessagesFor: (workflowId) => { |
| 41 |
if (!workflowId) return []; |
| 42 |
const segments = []; |
| 43 |
let segment = []; |
| 44 |
for (const { type, details } of get().messages) { |
| 45 |
const finished = ['completed', 'canceled'].includes(details.status); |
| 46 |
if (['workflow', 'workflow-component'].includes(type) && finished) { |
| 47 |
segments.push({ workflowId: details.workflowId, segment }); |
| 48 |
segment = []; |
| 49 |
continue; |
| 50 |
} |
| 51 |
// An error ended a run without a marker; what precedes it is dead. |
| 52 |
if (type === 'message' && details.error) { |
| 53 |
segment = []; |
| 54 |
continue; |
| 55 |
} |
| 56 |
if (type === 'tool' && 'result' in details) { |
| 57 |
segment.push(...buildToolMessages(details, { summarize: true })); |
| 58 |
continue; |
| 59 |
} |
| 60 |
if (type === 'message') segment.push(details); |
| 61 |
} |
| 62 |
const broken = segments.findLastIndex( |
| 63 |
(finished) => finished.workflowId !== workflowId, |
| 64 |
); |
| 65 |
return segments.slice(broken + 1).flatMap(({ segment }) => segment); |
| 66 |
}, |
| 67 |
getLastAssistantMessage: () => |
| 68 |
get()?.messages?.findLast( |
| 69 |
(message) => |
| 70 |
message.type === 'message' && message.details?.role === 'assistant', |
| 71 |
), |
| 72 |
hasMessages: () => get().messages.length > 0, |
| 73 |
addMessage: (type, details) => { |
| 74 |
const id = makeId(); |
| 75 |
set((state) => { |
| 76 |
// max 250 messages |
| 77 |
const max = Math.max(0, state.messages.length - 249); |
| 78 |
const next = { id, type, details }; |
| 79 |
return { |
| 80 |
// { id: 1, type: message, details: { role: 'user', content: 'Hello' } } |
| 81 |
// { id: 2, type: message, details: { role: 'assistant', content: 'Hi there!' } } |
| 82 |
// { id: 3, type: workflow, details: { name: 'Workflow 1' } } |
| 83 |
messages: [...state.messages.toSpliced(0, max), next], |
| 84 |
}; |
| 85 |
}); |
| 86 |
// A real message supersedes any in-flight progress status. |
| 87 |
useStatusStore.getState().clearStatuses(); |
| 88 |
return id; |
| 89 |
}, |
| 90 |
updateMessage: (id, details) => |
| 91 |
set((state) => ({ |
| 92 |
messages: state.messages.map((message) => |
| 93 |
message.id === id |
| 94 |
? { ...message, details: { ...message.details, ...details } } |
| 95 |
: message, |
| 96 |
), |
| 97 |
})), |
| 98 |
// pop messages all the way back to the last agent message |
| 99 |
popMessage: () => { |
| 100 |
set((state) => ({ |
| 101 |
messages: state.messages?.slice(0, -1) || [], |
| 102 |
})); |
| 103 |
}, |
| 104 |
clearMessages: () => set({ messages: [] }), |
| 105 |
}); |
| 106 |
|
| 107 |
const path = '/extendify/v1/agent/chat-events'; |
| 108 |
let lastSave = Promise.resolve(); |
| 109 |
const storage = { |
| 110 |
getItem: async () => await apiFetch({ path }), |
| 111 |
setItem: (_name, state) => { |
| 112 |
lastSave = apiFetch({ path, method: 'POST', data: { state } }); |
| 113 |
return lastSave; |
| 114 |
}, |
| 115 |
}; |
| 116 |
|
| 117 |
// Await before navigating — an aborted in-flight save loses the newest messages. |
| 118 |
export const flushChatStorage = (timeout = 5000) => |
| 119 |
Promise.race([ |
| 120 |
lastSave.catch(() => null), |
| 121 |
new Promise((resolve) => setTimeout(resolve, timeout)), |
| 122 |
]); |
| 123 |
|
| 124 |
export const useChatStore = create()( |
| 125 |
persist(devtools(state, { name: 'Extendify Agent Chat' }), { |
| 126 |
name: `extendify-agent-chat-${window.extSharedData.siteId}`, |
| 127 |
storage: createJSONStorage(() => storage), |
| 128 |
skipHydration: true, |
| 129 |
}), |
| 130 |
); |
| 131 |
|