From ab3ce260c836fda0fc4d23eb7ae39321780a023d Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Thu, 26 Mar 2026 13:38:06 +0300 Subject: [PATCH] feat: add subagent monitoring extension with horizontal widget layout Add new extension API hooks for tracking spawned subagents: - OnSubagentStart, OnSubagentChunk, OnSubagentEnd events - Extensions bridge for forwarding child subagent events Create subagent-monitor.go extension: - Displays horizontally-stacked widgets above input box - Shows real-time scrolling output from each subagent - Yaegi-safe: no sync.Mutex, no goroutines, nil-guarded context calls - Race-free design with on-demand elapsed time calculation Add comprehensive tests: - SessionStart, SubagentLifecycle, MultipleSubagents, SessionShutdown Update symbols.go to export new event types for Yaegi interpreter. --- .kit/extensions/subagent-monitor.go | 304 +++++++++++++++++++ examples/extensions/subagent-monitor_test.go | 159 ++++++++++ internal/extensions/api.go | 82 ++++- internal/extensions/events.go | 13 + internal/extensions/events_test.go | 7 +- internal/extensions/loader.go | 18 ++ internal/extensions/symbols.go | 5 + internal/extensions/test_api.go | 18 ++ pkg/kit/extensions_bridge.go | 144 +++++++++ 9 files changed, 747 insertions(+), 3 deletions(-) create mode 100644 .kit/extensions/subagent-monitor.go create mode 100644 examples/extensions/subagent-monitor_test.go diff --git a/.kit/extensions/subagent-monitor.go b/.kit/extensions/subagent-monitor.go new file mode 100644 index 00000000..f7e86d0c --- /dev/null +++ b/.kit/extensions/subagent-monitor.go @@ -0,0 +1,304 @@ +//go:build ignore + +// subagent-monitor — live horizontal widget strip for spawned subagents +// +// Subscribes to subagents spawned by the main Kit agent and displays a +// single widget just above the input box. Each subagent occupies one column +// in a side-by-side horizontal layout. Columns show scrolling real-time +// output as the subagent works. When a subagent finishes its column is +// removed automatically. +// +// Yaegi-safe design notes: +// - No sync.Mutex (Yaegi has reflection issues with sync primitives) +// - No channels in maps (Yaegi panics on range over map[string]chan) +// - All ctx.* calls guarded with nil checks +// - Simple data structures only +package main + +import ( + "fmt" + "strings" + "time" + + "kit/ext" +) + +// --------------------------------------------------------------------------- +// Per-subagent state +// --------------------------------------------------------------------------- + +type submonEntry struct { + id int + callID string + task string + lines []string + started time.Time + elapsed time.Duration +} + +const ( + submonColWidth = 34 // visible character width per column + submonMaxLines = 5 // scrolling output lines per column + submonColGap = 2 // spaces between columns +) + +// --------------------------------------------------------------------------- +// Package-level state - all simple types +// --------------------------------------------------------------------------- + +var ( + submonCtx ext.Context + submonHasCtx bool + submonEntries []*submonEntry + submonNextID int +) + +func submonInit() { + submonEntries = nil + submonNextID = 1 +} + +// --------------------------------------------------------------------------- +// String helpers +// --------------------------------------------------------------------------- + +func submonPad(s string, w int) string { + r := []rune(s) + if len(r) >= w { + return string(r[:w]) + } + return s + strings.Repeat(" ", w-len(r)) +} + +func submonTrunc(s string, w int) string { + r := []rune(s) + if len(r) <= w { + return s + } + if w <= 1 { + return "…" + } + return string(r[:w-1]) + "…" +} + +// --------------------------------------------------------------------------- +// Widget rendering +// --------------------------------------------------------------------------- + +func submonRenderColumn(e *submonEntry) []string { + var rows []string + + // Calculate elapsed time on-demand to avoid race conditions with ticker + elapsed := e.elapsed + if elapsed == 0 && !e.started.IsZero() { + elapsed = time.Since(e.started) + } + secs := int(elapsed.Seconds()) + timeStr := fmt.Sprintf("%ds", secs) + taskMax := submonColWidth - len(timeStr) - 3 + taskPart := submonTrunc(e.task, taskMax) + header := fmt.Sprintf("#%d %s %s", e.id, taskPart, timeStr) + rows = append(rows, submonPad(header, submonColWidth)) + + display := e.lines + if len(display) > submonMaxLines { + display = display[len(display)-submonMaxLines:] + } + for _, l := range display { + rows = append(rows, submonPad(" "+submonTrunc(l, submonColWidth-2), submonColWidth)) + } + for len(rows) < submonMaxLines+1 { + if len(rows) == 1 && len(e.lines) == 0 { + rows = append(rows, submonPad(" waiting…", submonColWidth)) + } else { + rows = append(rows, strings.Repeat(" ", submonColWidth)) + } + } + return rows +} + +func submonBuildWidget() string { + if len(submonEntries) == 0 { + return "" + } + + numCols := len(submonEntries) + numRows := submonMaxLines + 1 + cols := make([][]string, numCols) + for i, e := range submonEntries { + rows := submonRenderColumn(e) + col := make([]string, numRows) + for j := 0; j < numRows; j++ { + if j < len(rows) { + col[j] = rows[j] + } else { + col[j] = strings.Repeat(" ", submonColWidth) + } + } + cols[i] = col + } + + gap := strings.Repeat(" ", submonColGap) + var sb strings.Builder + for row := 0; row < numRows; row++ { + for ci := range cols { + if ci > 0 { + sb.WriteString(gap) + } + sb.WriteString(cols[ci][row]) + } + if row < numRows-1 { + sb.WriteString("\n") + } + } + return sb.String() +} + +func submonPushWidget() { + if !submonHasCtx { + return + } + if submonCtx.SetWidget == nil { + return + } + + text := submonBuildWidget() + if len(submonEntries) == 0 { + if submonCtx.RemoveWidget != nil { + submonCtx.RemoveWidget("submon") + } + return + } + submonCtx.SetWidget(ext.WidgetConfig{ + ID: "submon", + Placement: ext.WidgetAbove, + Content: ext.WidgetContent{Text: text}, + Style: ext.WidgetStyle{BorderColor: "#89b4fa"}, + Priority: 0, + }) +} + +func submonAppendLine(e *submonEntry, line string) { + line = strings.TrimRight(line, "\r\n") + if strings.TrimSpace(line) == "" { + return + } + e.lines = append(e.lines, line) +} + +// --------------------------------------------------------------------------- +// Init +// --------------------------------------------------------------------------- + +func Init(api ext.API) { + submonInit() + + api.OnSessionStart(func(_ ext.SessionStartEvent, ctx ext.Context) { + submonCtx = ctx + submonHasCtx = true + submonInit() + if ctx.RemoveWidget != nil { + ctx.RemoveWidget("submon") + } + }) + + api.OnAgentEnd(func(_ ext.AgentEndEvent, ctx ext.Context) { + submonCtx = ctx + submonHasCtx = true + }) + + // ── SubagentStart ──────────────────────────────────────────────────────── + api.OnSubagentStart(func(e ext.SubagentStartEvent, ctx ext.Context) { + submonCtx = ctx + submonHasCtx = true + + id := submonNextID + submonNextID++ + entry := &submonEntry{ + id: id, + callID: e.ToolCallID, + task: e.Task, + started: time.Now(), + } + submonEntries = append(submonEntries, entry) + + submonPushWidget() + }) + + // ── SubagentChunk ──────────────────────────────────────────────────────── + api.OnSubagentChunk(func(e ext.SubagentChunkEvent, ctx ext.Context) { + submonCtx = ctx + submonHasCtx = true + + var entry *submonEntry + for _, en := range submonEntries { + if en.callID == e.ToolCallID { + entry = en + break + } + } + if entry == nil { + return + } + + switch e.ChunkType { + case "text": + for _, line := range strings.Split(e.Content, "\n") { + submonAppendLine(entry, line) + } + case "tool_call": + submonAppendLine(entry, "→ "+e.ToolName) + case "tool_execution_start": + submonAppendLine(entry, "⚙ "+e.ToolName) + case "tool_result": + if e.IsError { + submonAppendLine(entry, "✗ "+e.ToolName) + } else { + submonAppendLine(entry, "✓ "+e.ToolName) + } + } + + submonPushWidget() + }) + + // ── SubagentEnd ────────────────────────────────────────────────────────── + api.OnSubagentEnd(func(e ext.SubagentEndEvent, ctx ext.Context) { + submonCtx = ctx + submonHasCtx = true + + var entry *submonEntry + for _, en := range submonEntries { + if en.callID == e.ToolCallID { + entry = en + break + } + } + if entry != nil { + entry.elapsed = time.Since(entry.started) + if e.ErrorMsg != "" { + submonAppendLine(entry, "✗ "+submonTrunc(e.ErrorMsg, submonColWidth-2)) + } + } + + submonPushWidget() + + // Remove the entry immediately (no goroutine to avoid races) + newEntries := submonEntries[:0] + for _, en := range submonEntries { + if en.callID != e.ToolCallID { + newEntries = append(newEntries, en) + } + } + submonEntries = newEntries + submonPushWidget() + }) + + // ── SessionShutdown ────────────────────────────────────────────────────── + api.OnSessionShutdown(func(_ ext.SessionShutdownEvent, ctx ext.Context) { + submonInit() + // Guard ctx access - may be nil during shutdown + if ctx.RemoveWidget != nil { + ctx.RemoveWidget("submon") + } + }) +} diff --git a/examples/extensions/subagent-monitor_test.go b/examples/extensions/subagent-monitor_test.go new file mode 100644 index 00000000..0129d73f --- /dev/null +++ b/examples/extensions/subagent-monitor_test.go @@ -0,0 +1,159 @@ +package main + +import ( + "fmt" + "testing" + "time" + + "github.com/mark3labs/kit/internal/extensions" + "github.com/mark3labs/kit/pkg/extensions/test" +) + +// TestSubagentMonitor_SessionStart verifies OnSessionStart initializes state +// without panicking and properly guards nil ctx calls. +func TestSubagentMonitor_SessionStart(t *testing.T) { + harness := test.New(t) + harness.LoadFile("../../.kit/extensions/subagent-monitor.go") + + // Emit SessionStart - should not panic even with nil ctx functions + _, err := harness.Emit(extensions.SessionStartEvent{SessionID: "test-session"}) + if err != nil { + t.Fatalf("SessionStart should not error: %v", err) + } +} + +// TestSubagentMonitor_SubagentLifecycle verifies the full subagent lifecycle +// creates entries and emits widget updates. +func TestSubagentMonitor_SubagentLifecycle(t *testing.T) { + harness := test.New(t) + harness.LoadFile("../../.kit/extensions/subagent-monitor.go") + + // Start session + _, err := harness.Emit(extensions.SessionStartEvent{SessionID: "test-session"}) + if err != nil { + t.Fatalf("SessionStart should not error: %v", err) + } + + // Emit SubagentStart + _, err = harness.Emit(extensions.SubagentStartEvent{ + ToolCallID: "call-1", + Task: "test task", + }) + if err != nil { + t.Fatalf("SubagentStart should not error: %v", err) + } + + // Emit a few chunks + for i := range 3 { + _, err = harness.Emit(extensions.SubagentChunkEvent{ + ToolCallID: "call-1", + Task: "test task", + ChunkType: "text", + Content: fmt.Sprintf("line %d", i), + }) + if err != nil { + t.Fatalf("SubagentChunk %d should not error: %v", i, err) + } + } + + // Emit tool call chunk + _, err = harness.Emit(extensions.SubagentChunkEvent{ + ToolCallID: "call-1", + Task: "test task", + ChunkType: "tool_call", + ToolName: "bash", + }) + if err != nil { + t.Fatalf("SubagentChunk tool_call should not error: %v", err) + } + + // Emit SubagentEnd + _, err = harness.Emit(extensions.SubagentEndEvent{ + ToolCallID: "call-1", + Task: "test task", + Response: "done", + }) + if err != nil { + t.Fatalf("SubagentEnd should not error: %v", err) + } + + // Give time for cleanup goroutine + time.Sleep(100 * time.Millisecond) +} + +// TestSubagentMonitor_MultipleSubagents verifies multiple parallel subagents. +func TestSubagentMonitor_MultipleSubagents(t *testing.T) { + harness := test.New(t) + harness.LoadFile("../../.kit/extensions/subagent-monitor.go") + + _, err := harness.Emit(extensions.SessionStartEvent{SessionID: "test-session"}) + if err != nil { + t.Fatalf("SessionStart should not error: %v", err) + } + + // Start 3 subagents + for i := 1; i <= 3; i++ { + _, err := harness.Emit(extensions.SubagentStartEvent{ + ToolCallID: fmt.Sprintf("call-%d", i), + Task: fmt.Sprintf("task %d", i), + }) + if err != nil { + t.Fatalf("SubagentStart %d should not error: %v", i, err) + } + } + + // Emit chunks for each + for i := 1; i <= 3; i++ { + _, err := harness.Emit(extensions.SubagentChunkEvent{ + ToolCallID: fmt.Sprintf("call-%d", i), + Task: fmt.Sprintf("task %d", i), + ChunkType: "text", + Content: fmt.Sprintf("output from agent %d", i), + }) + if err != nil { + t.Fatalf("SubagentChunk %d should not error: %v", i, err) + } + } + + // End all subagents + for i := 1; i <= 3; i++ { + _, err := harness.Emit(extensions.SubagentEndEvent{ + ToolCallID: fmt.Sprintf("call-%d", i), + Task: fmt.Sprintf("task %d", i), + Response: "completed", + }) + if err != nil { + t.Fatalf("SubagentEnd %d should not error: %v", i, err) + } + } + + time.Sleep(100 * time.Millisecond) +} + +// TestSubagentMonitor_SessionShutdown verifies shutdown doesn't panic +// even with nil ctx functions. +func TestSubagentMonitor_SessionShutdown(t *testing.T) { + harness := test.New(t) + harness.LoadFile("../../.kit/extensions/subagent-monitor.go") + + // Start then shutdown + _, err := harness.Emit(extensions.SessionStartEvent{SessionID: "test-session"}) + if err != nil { + t.Fatalf("SessionStart should not error: %v", err) + } + + // Start a subagent + _, err = harness.Emit(extensions.SubagentStartEvent{ + ToolCallID: "call-1", + Task: "test task", + }) + if err != nil { + t.Fatalf("SubagentStart should not error: %v", err) + } + + // Shutdown - should not panic even with active subagent + _, err = harness.Emit(extensions.SessionShutdownEvent{}) + if err != nil { + t.Fatalf("SessionShutdown should not error: %v", err) + } +} diff --git a/internal/extensions/api.go b/internal/extensions/api.go index b1106f5c..1cf98497 100644 --- a/internal/extensions/api.go +++ b/internal/extensions/api.go @@ -750,6 +750,9 @@ type API struct { registerOption func(OptionDef) registerShortcutFn func(ShortcutDef, func(Context)) registerMessageRendererFn func(MessageRendererConfig) + onSubagentStart func(func(SubagentStartEvent, Context)) + onSubagentChunk func(func(SubagentChunkEvent, Context)) + onSubagentEnd func(func(SubagentEndEvent, Context)) } // OnToolCall registers a handler that fires before a tool executes. @@ -781,6 +784,27 @@ func (a *API) OnToolResult(handler func(ToolResultEvent, Context) *ToolResultRes a.onToolResult(handler) } +// OnSubagentStart registers a handler that fires when a spawn_subagent tool +// call begins executing. Use the ToolCallID to correlate with subsequent +// OnSubagentChunk and OnSubagentEnd events for the same subagent. +func (a *API) OnSubagentStart(handler func(SubagentStartEvent, Context)) { + a.onSubagentStart(handler) +} + +// OnSubagentChunk registers a handler for real-time events from a running +// subagent. ChunkType identifies the kind of event ("text", "tool_call", +// "tool_result", "tool_execution_start", "tool_execution_end", etc.). +// Correlate with OnSubagentStart via the ToolCallID field. +func (a *API) OnSubagentChunk(handler func(SubagentChunkEvent, Context)) { + a.onSubagentChunk(handler) +} + +// OnSubagentEnd registers a handler that fires when a spawn_subagent call +// completes. ErrorMsg is non-empty when the subagent failed. +func (a *API) OnSubagentEnd(handler func(SubagentEndEvent, Context)) { + a.onSubagentEnd(handler) +} + // OnInput registers a handler that fires when user input is received. // Return a non-nil InputResult to transform or handle the input. func (a *API) OnInput(handler func(InputEvent, Context) *InputResult) { @@ -1781,9 +1805,65 @@ type BeforeCompactResult struct { func (BeforeCompactResult) isResult() {} // --------------------------------------------------------------------------- -// Theme types (exposed to Yaegi — concrete structs, string hex colors) +// Subagent lifecycle events (exposed to Yaegi — concrete structs) // --------------------------------------------------------------------------- +// SubagentStartEvent fires when a spawn_subagent tool call begins executing. +type SubagentStartEvent struct { + // ToolCallID is the LLM-assigned ID of the spawn_subagent tool call. + // Use this to correlate SubagentChunkEvent and SubagentEndEvent. + ToolCallID string + // Task is the task description passed to the subagent. + Task string +} + +func (e SubagentStartEvent) Type() EventType { return SubagentStart } + +// SubagentChunkEvent fires for each real-time event from a running subagent. +// Type field indicates the kind of event; read the relevant fields accordingly. +type SubagentChunkEvent struct { + // ToolCallID matches the SubagentStartEvent.ToolCallID for this subagent. + ToolCallID string + // Task is the task description (repeated for convenience). + Task string + // ChunkType identifies the event kind: + // "text" — LLM text chunk (read Content) + // "reasoning" — reasoning/thinking delta (read Content) + // "tool_call" — subagent called a tool (read ToolName, ToolArgs) + // "tool_result" — tool returned a result (read ToolName, ToolResult, IsError) + // "tool_execution_start" — tool began executing (read ToolName) + // "tool_execution_end" — tool finished executing (read ToolName) + // "turn_start" — subagent turn began + // "turn_end" — subagent turn ended + ChunkType string + // Content carries text for "text" and "reasoning" chunk types. + Content string + // ToolName is set on tool-related chunk types. + ToolName string + // ToolArgs is the JSON-encoded tool arguments for "tool_call" chunks. + ToolArgs string + // ToolResult is the tool output for "tool_result" chunks. + ToolResult string + // IsError is true when a "tool_result" chunk represents an error. + IsError bool +} + +func (e SubagentChunkEvent) Type() EventType { return SubagentChunk } + +// SubagentEndEvent fires when a spawn_subagent tool call completes. +type SubagentEndEvent struct { + // ToolCallID matches the SubagentStartEvent.ToolCallID for this subagent. + ToolCallID string + // Task is the task description. + Task string + // Response is the subagent's final text response (empty on error). + Response string + // ErrorMsg is non-empty when the subagent failed. + ErrorMsg string +} + +func (e SubagentEndEvent) Type() EventType { return SubagentEnd } + // ThemeColor is an adaptive color pair with light and dark hex values. // Either field may be empty to inherit from the default theme. type ThemeColor struct { diff --git a/internal/extensions/events.go b/internal/extensions/events.go index f834d8e2..2291b76f 100644 --- a/internal/extensions/events.go +++ b/internal/extensions/events.go @@ -71,6 +71,18 @@ const ( // BeforeCompact fires before context compaction runs. Handlers can // cancel compaction by returning Cancel=true. BeforeCompact EventType = "before_compact" + + // SubagentStart fires when a spawn_subagent tool call begins executing. + // Carries the tool call ID and the task description. + SubagentStart EventType = "subagent_start" + + // SubagentChunk fires for each real-time event emitted by a running + // subagent: text chunks, tool calls, tool results, etc. + SubagentChunk EventType = "subagent_chunk" + + // SubagentEnd fires when a spawn_subagent tool call completes (success + // or error). Carries the final response and any error message. + SubagentEnd EventType = "subagent_end" ) // AllEventTypes returns every supported event type. @@ -82,6 +94,7 @@ func AllEventTypes() []EventType { SessionStart, SessionShutdown, ModelChange, ContextPrepare, BeforeFork, BeforeSessionSwitch, BeforeCompact, + SubagentStart, SubagentChunk, SubagentEnd, } } diff --git a/internal/extensions/events_test.go b/internal/extensions/events_test.go index f13228ca..3e0c1954 100644 --- a/internal/extensions/events_test.go +++ b/internal/extensions/events_test.go @@ -4,8 +4,8 @@ import "testing" func TestAllEventTypes_Count(t *testing.T) { all := AllEventTypes() - if len(all) != 18 { - t.Fatalf("expected 18 event types, got %d", len(all)) + if len(all) != 21 { + t.Fatalf("expected 21 event types, got %d", len(all)) } } @@ -55,6 +55,9 @@ func TestEventType_TypeMethod(t *testing.T) { {BeforeForkEvent{TargetID: "abc"}, BeforeFork}, {BeforeSessionSwitchEvent{Reason: "new"}, BeforeSessionSwitch}, {BeforeCompactEvent{EstimatedTokens: 1000}, BeforeCompact}, + {SubagentStartEvent{ToolCallID: "x", Task: "t"}, SubagentStart}, + {SubagentChunkEvent{ToolCallID: "x", ChunkType: "text"}, SubagentChunk}, + {SubagentEndEvent{ToolCallID: "x"}, SubagentEnd}, } for _, tt := range tests { diff --git a/internal/extensions/loader.go b/internal/extensions/loader.go index 7a6d6563..324fa627 100644 --- a/internal/extensions/loader.go +++ b/internal/extensions/loader.go @@ -580,6 +580,24 @@ func loadSingleExtension(path string) (*LoadedExtension, error) { registerShortcutFn: func(def ShortcutDef, handler func(Context)) { ext.Shortcuts = append(ext.Shortcuts, ShortcutEntry{Def: def, Handler: handler}) }, + onSubagentStart: func(h func(SubagentStartEvent, Context)) { + reg(SubagentStart, func(e Event, c Context) Result { + h(e.(SubagentStartEvent), c) + return nil + }) + }, + onSubagentChunk: func(h func(SubagentChunkEvent, Context)) { + reg(SubagentChunk, func(e Event, c Context) Result { + h(e.(SubagentChunkEvent), c) + return nil + }) + }, + onSubagentEnd: func(h func(SubagentEndEvent, Context)) { + reg(SubagentEnd, func(e Event, c Context) Result { + h(e.(SubagentEndEvent), c) + return nil + }) + }, } // Call Init — the extension registers its handlers, tools, commands. diff --git a/internal/extensions/symbols.go b/internal/extensions/symbols.go index bec8bf78..4148f1bd 100644 --- a/internal/extensions/symbols.go +++ b/internal/extensions/symbols.go @@ -119,6 +119,11 @@ func Symbols() interp.Exports { "SubagentHandle": reflect.ValueOf((*SubagentHandle)(nil)), "SubagentEvent": reflect.ValueOf((*SubagentEvent)(nil)), + // Subagent lifecycle events + "SubagentStartEvent": reflect.ValueOf((*SubagentStartEvent)(nil)), + "SubagentChunkEvent": reflect.ValueOf((*SubagentChunkEvent)(nil)), + "SubagentEndEvent": reflect.ValueOf((*SubagentEndEvent)(nil)), + // Theme types "ThemeColor": reflect.ValueOf((*ThemeColor)(nil)), "ThemeColorConfig": reflect.ValueOf((*ThemeColorConfig)(nil)), diff --git a/internal/extensions/test_api.go b/internal/extensions/test_api.go index 60939a1e..222c4cbc 100644 --- a/internal/extensions/test_api.go +++ b/internal/extensions/test_api.go @@ -171,5 +171,23 @@ func NewTestAPI(ext *LoadedExtension) API { registerMessageRendererFn: func(config MessageRendererConfig) { ext.MessageRenderers = append(ext.MessageRenderers, config) }, + onSubagentStart: func(h func(SubagentStartEvent, Context)) { + reg(SubagentStart, func(e Event, c Context) Result { + h(e.(SubagentStartEvent), c) + return nil + }) + }, + onSubagentChunk: func(h func(SubagentChunkEvent, Context)) { + reg(SubagentChunk, func(e Event, c Context) Result { + h(e.(SubagentChunkEvent), c) + return nil + }) + }, + onSubagentEnd: func(h func(SubagentEndEvent, Context)) { + reg(SubagentEnd, func(e Event, c Context) Result { + h(e.(SubagentEndEvent), c) + return nil + }) + }, } } diff --git a/pkg/kit/extensions_bridge.go b/pkg/kit/extensions_bridge.go index 2b26415d..0587f278 100644 --- a/pkg/kit/extensions_bridge.go +++ b/pkg/kit/extensions_bridge.go @@ -2,6 +2,7 @@ package kit import ( "strings" + "sync" "charm.land/fantasy" "github.com/mark3labs/kit/internal/extensions" @@ -119,6 +120,125 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) { }) } + // --- Subagent lifecycle events --- + // When an extension registers OnSubagentStart/Chunk/End handlers, bridge + // the SDK's per-subagent event stream (SubscribeSubagent) into the + // extension runner. + // + // Flow: + // ToolExecutionStartEvent(spawn_subagent) → emit SubagentStartEvent + // → SubscribeSubagent → emit SubagentChunkEvents + // ToolResultEvent(spawn_subagent) → emit SubagentEndEvent + // + // We use ToolExecutionStart (not ToolCall) for SubagentStart because that + // is when the subagent actually begins running. We use ToolResult for + // SubagentEnd because that carries the final response text. + wantsSubagent := runner.HasHandlers(extensions.SubagentStart) || + runner.HasHandlers(extensions.SubagentChunk) || + runner.HasHandlers(extensions.SubagentEnd) + + if wantsSubagent { + // taskByCallID tracks the task description extracted from ToolCall input, + // keyed by toolCallID. Populated on ToolCall, consumed on ToolResult. + taskByCallID := make(map[string]string) + var taskMu = &taskMutex{} + + // Intercept ToolCall to capture the task and subscribe to child events. + m.Subscribe(func(e Event) { + ev, ok := e.(ToolCallEvent) + if !ok || ev.ToolName != "spawn_subagent" { + return + } + + // Extract task from parsed args. + task := "" + if ev.ParsedArgs != nil { + if t, ok := ev.ParsedArgs["task"].(string); ok { + task = t + } + } + taskMu.set(taskByCallID, ev.ToolCallID, task) + + // Subscribe to child events so we can forward them as SubagentChunkEvents. + if runner.HasHandlers(extensions.SubagentChunk) { + m.SubscribeSubagent(ev.ToolCallID, func(childEvent Event) { + chunk := extensions.SubagentChunkEvent{ + ToolCallID: ev.ToolCallID, + Task: task, + } + switch ce := childEvent.(type) { + case MessageUpdateEvent: + chunk.ChunkType = "text" + chunk.Content = ce.Chunk + case TurnStartEvent: + chunk.ChunkType = "turn_start" + case TurnEndEvent: + chunk.ChunkType = "turn_end" + case ToolCallEvent: + chunk.ChunkType = "tool_call" + chunk.ToolName = ce.ToolName + chunk.ToolArgs = ce.ToolArgs + case ToolExecutionStartEvent: + chunk.ChunkType = "tool_execution_start" + chunk.ToolName = ce.ToolName + case ToolExecutionEndEvent: + chunk.ChunkType = "tool_execution_end" + chunk.ToolName = ce.ToolName + case ToolResultEvent: + chunk.ChunkType = "tool_result" + chunk.ToolName = ce.ToolName + chunk.ToolResult = ce.Result + chunk.IsError = ce.IsError + default: + return // skip unknown event types + } + _, _ = runner.Emit(chunk) + }) + } + }) + + // Emit SubagentStartEvent when execution begins. + if runner.HasHandlers(extensions.SubagentStart) { + m.Subscribe(func(e Event) { + ev, ok := e.(ToolExecutionStartEvent) + if !ok || ev.ToolName != "spawn_subagent" { + return + } + task := taskMu.get(taskByCallID, ev.ToolCallID) + _, _ = runner.Emit(extensions.SubagentStartEvent{ + ToolCallID: ev.ToolCallID, + Task: task, + }) + }) + } + + // Emit SubagentEndEvent when the tool result arrives. + if runner.HasHandlers(extensions.SubagentEnd) { + m.Subscribe(func(e Event) { + ev, ok := e.(ToolResultEvent) + if !ok || ev.ToolName != "spawn_subagent" { + return + } + task := taskMu.get(taskByCallID, ev.ToolCallID) + taskMu.del(taskByCallID, ev.ToolCallID) + errMsg := "" + if ev.IsError { + errMsg = ev.Result + } + response := "" + if !ev.IsError { + response = ev.Result + } + _, _ = runner.Emit(extensions.SubagentEndEvent{ + ToolCallID: ev.ToolCallID, + Task: task, + Response: response, + ErrorMsg: errMsg, + }) + }) + } + } + // --- Context filtering hook --- // Extension ContextPrepare → SDK ContextPrepare hook. if runner.HasHandlers(extensions.ContextPrepare) { @@ -204,3 +324,27 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) { }) } } + +// taskMutex is a simple mutex-protected map helper used by bridgeExtensions. +// It lives in this file to avoid polluting the kit package with unexported types. +type taskMutex struct { + mu sync.Mutex +} + +func (t *taskMutex) set(m map[string]string, key, val string) { + t.mu.Lock() + m[key] = val + t.mu.Unlock() +} + +func (t *taskMutex) get(m map[string]string, key string) string { + t.mu.Lock() + defer t.mu.Unlock() + return m[key] +} + +func (t *taskMutex) del(m map[string]string, key string) { + t.mu.Lock() + delete(m, key) + t.mu.Unlock() +}