Skip to content

Commit 4bf190b

Browse files
mudlermudler
committed
fix events and tool stream in agent chat
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 040ad64 commit 4bf190b

7 files changed

Lines changed: 228 additions & 58 deletions

File tree

‎core/application/distributed.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*distribut
150150
// Initialize agent event bridge
151151
agentBridge := agents.NewEventBridge(natsClient, agentStore, cfg.Distributed.InstanceID)
152152

153+
// Start observable persister — captures observable_update events from workers
154+
// (which have no DB access) and persists them to PostgreSQL.
155+
if err := agentBridge.StartObservablePersister(); err != nil {
156+
xlog.Warn("Failed to start observable persister", "error", err)
157+
} else {
158+
xlog.Info("Observable persister started")
159+
}
160+
153161
// Initialize Phase 4 stores (MCP, Gallery, FineTune, Skills)
154162
distStores, err := distributed.InitStores(authDB)
155163
if err != nil {

‎core/http/react-ui/src/components/Sidebar.jsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ const sections = [
4848
{ path: '/app/users', icon: 'fas fa-users', label: 'Users', adminOnly: true, authOnly: true },
4949
{ path: '/app/backends', icon: 'fas fa-server', label: 'Backends', adminOnly: true },
5050
{ path: '/app/traces', icon: 'fas fa-chart-line', label: 'Traces', adminOnly: true },
51-
{ path: '/app/nodes', icon: 'fas fa-network-wired', label: 'Nodes', adminOnly: true },
51+
{ path: '/app/nodes', icon: 'fas fa-network-wired', label: 'Nodes', adminOnly: true, feature: 'distributed' },
5252
{ path: '/app/p2p', icon: 'fas fa-circle-nodes', label: 'Swarm', adminOnly: true },
5353
{ path: '/app/manage', icon: 'fas fa-desktop', label: 'System', adminOnly: true },
5454
{ path: '/app/settings', icon: 'fas fa-cog', label: 'Settings', adminOnly: true },

‎core/http/react-ui/src/pages/AgentChat.jsx‎

Lines changed: 66 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,8 @@ export default function AgentChat() {
125125
// Tracks which conversation initiated the current request — SSE responses
126126
// are pinned to this ID so switching tabs doesn't misdirect them.
127127
const processingChatIdRef = useRef(null)
128+
// Maps backend messageID → conversationId for robust SSE routing across navigations.
129+
const pendingRequestsRef = useRef(new Map())
128130

129131
const processing = processingChatId === activeId
130132

@@ -142,20 +144,29 @@ export default function AgentChat() {
142144
es.addEventListener('json_message', (e) => {
143145
try {
144146
const data = JSON.parse(e.data)
147+
const sender = data.sender || (data.role === 'user' ? 'user' : 'agent')
148+
// Skip user message echoes — already added locally in handleSend
149+
if (sender === 'user') return
145150
const msg = {
146151
id: nextId(),
147-
sender: data.sender || (data.role === 'user' ? 'user' : 'agent'),
152+
sender,
148153
content: data.content || data.message || '',
149154
timestamp: data.timestamp ? Math.floor(data.timestamp / 1e6) : Date.now(),
150155
}
151156
if (data.metadata && Object.keys(data.metadata).length > 0) {
152157
msg.metadata = data.metadata
153158
}
154-
// Pin message to the conversation that initiated the request
155-
const targetId = processingChatIdRef.current || activeIdRef.current
159+
// Route to conversation: try messageID mapping first, then processingChatIdRef, then active
160+
const msgId = data.message_id || ''
161+
const baseId = msgId.replace(/-agent$/, '')
162+
const targetId = pendingRequestsRef.current.get(baseId)
163+
|| processingChatIdRef.current
164+
|| activeIdRef.current
156165
addMessageToConvRef.current(targetId, msg)
157166
// Clear streaming + processing state when the final agent message arrives
158-
if (msg.sender === 'agent') {
167+
if (sender === 'agent') {
168+
pendingRequestsRef.current.delete(baseId)
169+
processingChatIdRef.current = null
159170
setProcessingChatId(null)
160171
setStreamContent('')
161172
setStreamReasoning('')
@@ -170,17 +181,20 @@ export default function AgentChat() {
170181
try {
171182
const data = JSON.parse(e.data)
172183
if (data.status === 'processing') {
173-
// Track which conversation is processing so responses go to the right place
174-
processingChatIdRef.current = activeIdRef.current
175-
setProcessingChatId(activeIdRef.current)
184+
// Track which conversation is processing so responses go to the right place.
185+
// Only set if not already pinned by handleSend (avoids race when user switches conversations).
186+
if (!processingChatIdRef.current) {
187+
processingChatIdRef.current = activeIdRef.current
188+
setProcessingChatId(activeIdRef.current)
189+
}
176190
setStreamContent('')
177191
setStreamReasoning('')
178192
setStreamToolCalls([])
179193
} else if (data.status === 'completed') {
180-
processingChatIdRef.current = null
181-
// Don't clear processingChatId or streaming state here —
194+
// Don't clear processingChatIdRef, processingChatId, or streaming state here —
182195
// they'll be cleared when the agent's json_message arrives,
183-
// so reasoning and tool calls remain visible until the response replaces them.
196+
// so reasoning and tool calls remain visible until the response replaces them
197+
// and late-arriving messages still route to the correct conversation.
184198
}
185199
} catch (_err) {
186200
// ignore
@@ -206,6 +220,16 @@ export default function AgentChat() {
206220
updated[updated.length - 1] = { ...updated[updated.length - 1], args: updated[updated.length - 1].args + args }
207221
return updated
208222
})
223+
} else if (data.type === 'tool_result') {
224+
const tname = data.tool_name || ''
225+
setStreamToolCalls(prev => {
226+
const updated = [...prev]
227+
const idx = updated.findLastIndex(tc => tc.name === tname && !tc.result)
228+
if (idx >= 0) {
229+
updated[idx] = { ...updated[idx], result: data.tool_result || 'done' }
230+
}
231+
return updated
232+
})
209233
} else if (data.type === 'done') {
210234
// Content will be finalized by json_message event
211235
}
@@ -244,6 +268,8 @@ export default function AgentChat() {
244268
return () => {
245269
es.close()
246270
eventSourceRef.current = null
271+
processingChatIdRef.current = null
272+
pendingRequestsRef.current.clear()
247273
}
248274
}, [name, userId, addToast, nextId])
249275

@@ -325,14 +351,22 @@ export default function AgentChat() {
325351
if (!msg || processing) return
326352
setInput('')
327353
if (textareaRef.current) textareaRef.current.style.height = 'auto'
354+
// Add user message locally immediately (like standard chat)
355+
addMessage({ id: nextId(), sender: 'user', content: msg, timestamp: Date.now() })
328356
setProcessingChatId(activeId)
357+
processingChatIdRef.current = activeId
329358
try {
330-
await agentsApi.chat(name, msg, userId)
359+
const resp = await agentsApi.chat(name, msg, userId)
360+
// Map backend messageID → conversation so SSE events route correctly
361+
if (resp && resp.message_id) {
362+
pendingRequestsRef.current.set(resp.message_id, activeId)
363+
}
331364
} catch (err) {
332365
addToast(`Failed to send message: ${err.message}`, 'error')
366+
processingChatIdRef.current = null
333367
setProcessingChatId(null)
334368
}
335-
}, [input, processing, name, activeId, addToast, userId])
369+
}, [input, processing, name, activeId, addToast, userId, addMessage, nextId])
336370

337371
const handleKeyDown = (e) => {
338372
if (e.key === 'Enter' && !e.shiftKey) {
@@ -622,16 +656,28 @@ export default function AgentChat() {
622656
</div>
623657
</details>
624658
)}
625-
{streamToolCalls.length > 0 && !streamContent && (
659+
{streamToolCalls.length > 0 && (
626660
<div className="chat-activity-group" style={{ marginBottom: 'var(--spacing-sm)' }}>
627661
{streamToolCalls.map((tc, idx) => (
628-
<div key={idx} className="chat-activity-item chat-activity-tool-call" style={{ padding: 'var(--spacing-xs) var(--spacing-sm)' }}>
629-
<span className="chat-activity-item-label">
630-
<i className="fas fa-bolt" style={{ marginRight: 'var(--spacing-xs)' }} />
631-
{tc.name}
632-
</span>
633-
<span style={{ opacity: 0.5, fontSize: '0.85em', marginLeft: 'var(--spacing-xs)' }}>calling...</span>
634-
</div>
662+
<details key={idx} className="chat-activity-item chat-activity-tool-call" style={{ padding: 'var(--spacing-xs) var(--spacing-sm)' }} open={!tc.result}>
663+
<summary className="chat-activity-item-label" style={{ cursor: 'pointer', display: 'flex', alignItems: 'center', gap: 'var(--spacing-xs)' }}>
664+
<i className={`fas ${tc.result ? 'fa-check' : 'fa-bolt'}`} />
665+
<strong>{tc.name}</strong>
666+
<span style={{ opacity: 0.5, fontSize: '0.85em' }}>
667+
{tc.result ? 'done' : 'calling...'}
668+
</span>
669+
</summary>
670+
{tc.args && (
671+
<pre style={{ margin: '4px 0', fontSize: '0.75rem', opacity: 0.8, whiteSpace: 'pre-wrap', wordBreak: 'break-word' }}>
672+
{(() => { try { return JSON.stringify(JSON.parse(tc.args), null, 2) } catch { return tc.args } })()}
673+
</pre>
674+
)}
675+
{tc.result && (
676+
<pre style={{ margin: '4px 0', fontSize: '0.75rem', opacity: 0.7, whiteSpace: 'pre-wrap', wordBreak: 'break-word', maxHeight: '200px', overflow: 'auto' }}>
677+
{tc.result}
678+
</pre>
679+
)}
680+
</details>
635681
))}
636682
</div>
637683
)}

‎core/http/routes/localai.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ func RegisterLocalAIRoutes(router *echo.Echo,
140140
"mcp": !appConfig.DisableMCP,
141141
"fine_tuning": true,
142142
"quantization": true,
143+
"distributed": appConfig.Distributed.Enabled,
143144
})
144145
})
145146

‎core/services/agents/dispatcher.go‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"encoding/json"
66
"fmt"
7+
"math/rand"
78
"strings"
89
"sync"
910
"sync/atomic"
@@ -133,6 +134,9 @@ func (d *LocalDispatcher) buildLocalCallbacks(writer SSEWriter, messageID string
133134
data["type"] = "content"
134135
data["content"] = ev.Content
135136
case cogito.StreamEventToolCall:
137+
if isInternalCogitoTool(ev.ToolName) {
138+
return
139+
}
136140
data["type"] = "tool_call"
137141
data["tool_name"] = ev.ToolName
138142
data["tool_args"] = ev.ToolArgs
@@ -152,6 +156,16 @@ func (d *LocalDispatcher) buildLocalCallbacks(writer SSEWriter, messageID string
152156
OnToolCall: func(name, args string) {
153157
// Already forwarded via OnStream
154158
},
159+
OnToolResult: func(name, result string) {
160+
if writer != nil {
161+
writer.SendEvent("stream_event", map[string]any{
162+
"type": "tool_result",
163+
"tool_name": name,
164+
"tool_result": result,
165+
"timestamp": time.Now().Format(time.RFC3339),
166+
})
167+
}
168+
},
155169
OnStatus: func(status string) {
156170
if writer != nil {
157171
writer.SendEvent("json_message_status", map[string]string{
@@ -297,7 +311,8 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent) {
297311
// Build execution options: skills come from the enriched NATS payload
298312
// (workers have no database access).
299313
opts := ExecuteChatOpts{
300-
UserID: evt.UserID,
314+
UserID: evt.UserID,
315+
MessageID: evt.MessageID,
301316
}
302317
if len(evt.Skills) > 0 {
303318
opts.SkillProvider = &staticSkillProvider{skills: evt.Skills}
@@ -333,13 +348,18 @@ func (p *staticSkillProvider) ListSkills() ([]SkillInfo, error) {
333348
func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks {
334349
// Observable tracking: build LocalAGI-compatible observable records
335350
// from cogito callbacks so the UI can render them properly.
351+
//
352+
// IDs must be globally unique (not just per-job) because the UI's buildTree
353+
// uses them as map keys. We use a random base + counter so IDs are unique
354+
// across jobs while parent–child relationships still work within a job.
355+
idBase := rand.Int31n(1<<30) + 1 // random base, avoids collisions across jobs
336356
var obsIDCounter atomic.Int32
337357
var mu sync.Mutex
338358
var currentToolObs *coreTypes.Observable
339359
var reasoningBuf strings.Builder
340360

341361
nextID := func() int32 {
342-
return obsIDCounter.Add(1)
362+
return idBase + obsIDCounter.Add(1)
343363
}
344364

345365
// Root observable for this chat job
@@ -374,6 +394,9 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks {
374394
data["type"] = "content"
375395
data["content"] = ev.Content
376396
case cogito.StreamEventToolCall:
397+
if isInternalCogitoTool(ev.ToolName) {
398+
return
399+
}
377400
data["type"] = "tool_call"
378401
data["tool_name"] = ev.ToolName
379402
data["tool_args"] = ev.ToolArgs
@@ -407,6 +430,15 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks {
407430
// Tool calls tracked via OnStream
408431
},
409432
OnToolResult: func(name, result string) {
433+
// Emit tool_result stream event for real-time UI display
434+
if d.eventBridge != nil {
435+
d.eventBridge.PublishStreamEvent(evt.AgentName, evt.UserID, map[string]any{
436+
"type": "tool_result",
437+
"tool_name": name,
438+
"tool_result": result,
439+
"timestamp": time.Now().Format(time.RFC3339),
440+
})
441+
}
410442
// Persist tool result: complete the current tool observable
411443
mu.Lock()
412444
obs := currentToolObs

0 commit comments

Comments
 (0)