From 851fc61b43ebf5814f5edfe75e44aa2b34af5f33 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 20:52:33 -0500 Subject: [PATCH 1/6] feat(event): add reactive runtime event plumbing --- internal/assistant/client.go | 1 + internal/assistant/lifecycle.go | 51 +++++++++++++++++++- internal/assistant/runtime_lifecycle_test.go | 3 ++ internal/assistant/usage_events.go | 4 +- 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/internal/assistant/client.go b/internal/assistant/client.go index 9174373..68758e9 100644 --- a/internal/assistant/client.go +++ b/internal/assistant/client.go @@ -42,6 +42,7 @@ const ( jsonTextKey = "text" jsonThinkingKey = "thinking" jsonDisplayKey = "display" + jsonUsageKey = "usage" jsonUserRole = "user" functionToolType = "function" functionCallType = "function_call" diff --git a/internal/assistant/lifecycle.go b/internal/assistant/lifecycle.go index 58aecaa..4367d01 100644 --- a/internal/assistant/lifecycle.go +++ b/internal/assistant/lifecycle.go @@ -126,6 +126,18 @@ func (runtime *Runtime) dispatchTurnErrorLifecycle( runtime.dispatchLifecycle(ctx, extension.LifecycleAgentEnd, payload) } +func (runtime *Runtime) dispatchContextBuild( + ctx context.Context, + sessionID string, + cwd string, + systemPrompt string, + messages []database.MessageEntity, + usage model.TokenUsage, +) { + payload := contextBuildLifecyclePayload(sessionID, cwd, systemPrompt, messages, usage) + runtime.dispatchLifecycle(ctx, extension.LifecycleContextBuild, payload) +} + func promptLifecyclePayload(request *PromptRequest) map[string]any { if request == nil { return map[string]any{} @@ -185,7 +197,7 @@ func turnEndLifecyclePayload( "cached": cached, lifecycleErrorKey: "", jsonSessionIDKey: sessionID, - "usage": tokenUsageLifecyclePayload(usage), + jsonUsageKey: tokenUsageLifecyclePayload(usage), lifecycleUserEntryIDKey: userEntryID, } if turnErr != nil { @@ -195,6 +207,43 @@ func turnEndLifecyclePayload( return payload } +func contextBuildLifecyclePayload( + sessionID string, + cwd string, + systemPrompt string, + messages []database.MessageEntity, + usage model.TokenUsage, +) map[string]any { + return map[string]any{ + lifecycleCWDKey: cwd, + jsonSessionIDKey: sessionID, + "message_count": len(messages), + "system_tokens": estimateTokens(systemPrompt), + "message_tokens": estimateMessageTokens(messages), + jsonUsageKey: tokenUsageLifecyclePayload(usage), + "model_facing_roles": modelFacingRoleCounts(messages), + } +} + +func estimateMessageTokens(messages []database.MessageEntity) int { + tokens := 0 + for index := range messages { + tokens += estimateTokens(messages[index].Content) + } + + return tokens +} + +func modelFacingRoleCounts(messages []database.MessageEntity) map[string]int { + counts := map[string]int{} + for index := range messages { + role := string(messages[index].Role) + counts[role]++ + } + + return counts +} + func entryLifecyclePayload(entry *database.EntryEntity) map[string]any { return map[string]any{ lifecycleCreatedAtKey: entry.CreatedAt.Format(timeFormatRFC3339Nano), diff --git a/internal/assistant/runtime_lifecycle_test.go b/internal/assistant/runtime_lifecycle_test.go index 0fb3388..a099c35 100644 --- a/internal/assistant/runtime_lifecycle_test.go +++ b/internal/assistant/runtime_lifecycle_test.go @@ -25,6 +25,7 @@ local names = { "before_agent_start", "agent_start", "turn_start", + "context_build", "message_append", "turn_end", "agent_end", @@ -75,6 +76,7 @@ func TestRuntime_PromptEmitsOrderedSessionTurnLifecycleEvents(t *testing.T) { "before_agent_start|||session", "agent_start|||session", "turn_start|||session", + "context_build|||session", "message_append|assistant||session", "turn_end|||session", "agent_end|||session", @@ -92,6 +94,7 @@ func TestRuntime_PromptEmitsOrderedSessionTurnLifecycleEvents(t *testing.T) { "before_agent_start|||session", "agent_start|||session", "turn_start|||session", + "context_build|||session", "turn_end||error|session", "agent_end||error|session", }}, diff --git a/internal/assistant/usage_events.go b/internal/assistant/usage_events.go index 803b310..c2947cc 100644 --- a/internal/assistant/usage_events.go +++ b/internal/assistant/usage_events.go @@ -22,9 +22,9 @@ func (runtime *Runtime) emitUsage(ctx context.Context, onEvent func(StreamEvent) jsonInputTokensKey: usage.InputTokens, jsonOutputTokensKey: usage.OutputTokens, } - runtime.emit(ctx, "usage", payload) + runtime.emit(ctx, jsonUsageKey, payload) if runtime.extensions != nil { - if err := runtime.extensions.Emit(ctx, "usage", payload); err != nil { + if err := runtime.extensions.Emit(ctx, jsonUsageKey, payload); err != nil { runtime.logger.Debug("emit usage extension event failed", "error", err) } } From 38c4c1dc5fd996b12240854200f7fece32698098 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 20:52:41 -0500 Subject: [PATCH 2/6] feat(event): add context lifecycle event From 3882a7ac7294810ea745abdb2f2ca34e1934bb94 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 20:53:16 -0500 Subject: [PATCH 3/6] feat(assistant): publish context build events From 70ed8c8f61a17a6e71413a496bf9e505d84d11f0 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 22:59:39 -0500 Subject: [PATCH 4/6] fix(runtime): wire context build event From 263d6f12cf78d8a44bc73574790050c874e1864f Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 23:02:46 -0500 Subject: [PATCH 5/6] fix(runtime): wire context build event From e853469ca6b2521a6c77674c282b6c044d511bc8 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Mon, 18 May 2026 23:22:45 -0500 Subject: [PATCH 6/6] fix(runtime): wire context build event --- internal/assistant/runtime.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/assistant/runtime.go b/internal/assistant/runtime.go index ffba1b2..9dd91ac 100644 --- a/internal/assistant/runtime.go +++ b/internal/assistant/runtime.go @@ -650,6 +650,7 @@ func (runtime *Runtime) modelResponse( } estimatedUsage := estimateTokenUsage(systemPrompt, messages, &selectedModel) + runtime.dispatchContextBuild(ctx, sessionID, cwd, systemPrompt, messages, estimatedUsage) runtime.emitUsage(ctx, onEvent, estimatedUsage) request := &CompletionRequest{ Model: selectedModel,