From f348b968a249720bad802fb324df6b7845fae0ad Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 4 Aug 2026 14:21:45 +0200 Subject: [PATCH] fix: wire up event handler so chat stream receives events After the fs/session refactor, NewEventHandler was defined but never called. Agents created via session restore or CreateAgent had nil output handlers, so emit() was a no-op and the chat stream stayed empty. - Add Agent.SetOutput() to allow setting the event handler post-creation - Wire up SetOutput(NewEventHandler(al)) in fs/newroot.go for restored sessions - Wire up SetOutput(NewEventHandler(al)) in fs/handlers.go for new agents - Also includes expanded UsageStats fields and TaskState persistence --- agent/agent.go | 38 ++++++++++++++++++++++++++++++-------- fs/handlers.go | 3 ++- fs/newroot.go | 4 ++++ session/persist.go | 30 ++++++++++++++++++++---------- 4 files changed, 56 insertions(+), 19 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index d0acb56..1580dd6 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -79,12 +79,24 @@ func (ag *Agent) Messages() []backend.Message { return ag.history.history() } +// TaskState returns the current task state, if any. +func (ag *Agent) TaskState() *TaskState { + if ag.history == nil { + return nil + } + return ag.history.TaskState +} + // UsageStats holds token usage and cost information. type UsageStats struct { - TotalInputTokens int - TotalOutputTokens int - TotalRequests int - SessionCostUSD float64 + TotalInputTokens int + TotalCachedInputTokens int + TotalCacheCreationTokens int + TotalOutputTokens int + TotalRequests int + Estimated bool + LastTurnCostUSD float64 + SessionCostUSD float64 } // Usage returns the agent's usage statistics. @@ -93,10 +105,14 @@ func (ag *Agent) Usage() *UsageStats { return nil } return &UsageStats{ - TotalInputTokens: ag.history.TotalInputTokens, - TotalOutputTokens: ag.history.TotalOutputTokens, - TotalRequests: ag.history.TotalRequests, - SessionCostUSD: ag.history.SessionCostUSD, + TotalInputTokens: ag.history.TotalInputTokens, + TotalCachedInputTokens: ag.history.TotalCachedInputTokens, + TotalCacheCreationTokens: ag.history.TotalCacheCreationTokens, + TotalOutputTokens: ag.history.TotalOutputTokens, + TotalRequests: ag.history.TotalRequests, + Estimated: ag.history.Estimated, + LastTurnCostUSD: ag.history.LastTurnCostUSD, + SessionCostUSD: ag.history.SessionCostUSD, } } @@ -221,6 +237,12 @@ func (ag *Agent) emit(ev Event) { } } +// SetOutput sets the agent's event output handler. +// This allows the fs layer to wire up the event handler after agent creation. +func (ag *Agent) SetOutput(h EventHandler) { + ag.output = h +} + // CWD returns the agent's working directory. func (ag *Agent) Cwd() string { ag.stateMu.RLock() diff --git a/fs/handlers.go b/fs/handlers.go index 793f0a6..df614c5 100644 --- a/fs/handlers.go +++ b/fs/handlers.go @@ -393,9 +393,10 @@ func requestAgentNew(ctx HandlerCtx, data []byte) ([]byte, error) { if err != nil { return nil, fmt.Errorf("create agent: %w", err) } - // Create AgentLog for this agent + // Create AgentLog for this agent and wire up the event handler al := NewAgentLog("") ctx.Session.SetAgentLog(ag.ID(), al) + ag.SetOutput(NewEventHandler(al)) return []byte(ag.ID() + "\n"), nil } diff --git a/fs/newroot.go b/fs/newroot.go index 6c7396d..14d8933 100644 --- a/fs/newroot.go +++ b/fs/newroot.go @@ -83,6 +83,10 @@ func NewRoot(cfg Config) *Tree { al := NewAgentLog(ra.Remote) node.SetAgentLog(ra.ID, al) replayMessagesToLog(al, ra.Messages) + // Wire up the event handler so agent events flow to the AgentLog + if ag := r.Session.FindAgent(ra.ID); ag != nil { + ag.SetOutput(NewEventHandler(al)) + } } } diff --git a/session/persist.go b/session/persist.go index c181abb..1b6991c 100644 --- a/session/persist.go +++ b/session/persist.go @@ -31,11 +31,16 @@ type PersistedAgent struct { Backend string `json:"backend"` Model string `json:"model"` Messages []backend.Message `json:"messages"` + TaskState *agent.TaskState `json:"taskState,omitempty"` // Usage tracking - TotalInputTokens int `json:"totalInputTokens,omitempty"` - TotalOutputTokens int `json:"totalOutputTokens,omitempty"` - TotalRequests int `json:"totalRequests,omitempty"` - SessionCostUSD float64 `json:"sessionCostUSD,omitempty"` + TotalInputTokens int `json:"totalInputTokens,omitempty"` + TotalCachedInputTokens int `json:"totalCachedInputTokens,omitempty"` + TotalCacheCreationTokens int `json:"totalCacheCreationTokens,omitempty"` + TotalOutputTokens int `json:"totalOutputTokens,omitempty"` + TotalRequests int `json:"totalRequests,omitempty"` + Estimated bool `json:"estimated,omitempty"` + LastTurnCostUSD float64 `json:"lastTurnCostUSD,omitempty"` + SessionCostUSD float64 `json:"sessionCostUSD,omitempty"` } // ActiveSessionsDir returns the directory for persisted sessions. @@ -63,17 +68,22 @@ func PersistSession(name string) error { for _, ag := range agents { pa := PersistedAgent{ - ID: ag.ID(), - Name: ag.Name(), - Profile: ag.Profile(), - Backend: ag.BackendName(), - Model: ag.ModelName(), - Messages: ag.Messages(), + ID: ag.ID(), + Name: ag.Name(), + Profile: ag.Profile(), + Backend: ag.BackendName(), + Model: ag.ModelName(), + Messages: ag.Messages(), + TaskState: ag.TaskState(), } if usage := ag.Usage(); usage != nil { pa.TotalInputTokens = usage.TotalInputTokens + pa.TotalCachedInputTokens = usage.TotalCachedInputTokens + pa.TotalCacheCreationTokens = usage.TotalCacheCreationTokens pa.TotalOutputTokens = usage.TotalOutputTokens pa.TotalRequests = usage.TotalRequests + pa.Estimated = usage.Estimated + pa.LastTurnCostUSD = usage.LastTurnCostUSD pa.SessionCostUSD = usage.SessionCostUSD } ps.Agents = append(ps.Agents, pa)