diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index 18b2c1bc4..00698c1c6 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -192,7 +192,9 @@ Events the agent emits (in dispatch order): | `specialistStart` | A sub-agent is spawned | yes (specialist name) | | `specialistStop` | A sub-agent ends | yes (specialist name) | -A hook's exit code decides what happens next: `0` continues, non-zero blocks the tool call (`beforeTool`) or surfaces an error (`afterTool`). Hook execution is recorded in the audit log; the audit is reachable from the agent's view of past actions, not from a dedicated `zero doctor` check. +For swarm members the matcher subject is the **agent type** (`subagent`, `teammate`, …), not a registered specialist profile name — swarm members run from an inline manifest keyed by agent type. + +A hook's exit code decides what happens next: `0` continues, non-zero blocks the tool call (`beforeTool`) or surfaces an error (`afterTool`). `sessionStart` / `sessionEnd` / `specialistStart` / `specialistStop` hooks are advisory (failures are audited but do not interrupt the run). Hook execution is recorded in the audit log; the audit is reachable from the agent's view of past actions, not from a dedicated `zero doctor` check. > **Roadmap.** An in-UI hooks manager is on the backlog. Today you edit the JSON directly. diff --git a/internal/cli/app.go b/internal/cli/app.go index 80854beb4..7ff7245a2 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -807,6 +807,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a // report can be combined with the plugin activation's, and emit at most one // notice when project hooks/plugins were dropped for an untrusted workspace. hookDispatcher, hookSkip := newHookDispatcherWithExtra(workspaceRoot, pluginActivation.hooks, trustRoot, executionRunner) + specialistRuntime.attachLifecycleHooks(hookDispatcher) emitTrustNotice(stderr, hookSkip, pluginActivation.trustSkip, mcpSkip) return deps.runTUI(context.Background(), tui.Options{ Cwd: workspaceRoot, @@ -1044,6 +1045,9 @@ type agentToolRuntime struct { // the orchestrator's system-prompt delegation section. Populated alongside the // tools, so it is present exactly when the Task tool is. specialists []agent.SpecialistInfo + // lifecycleHooks is shared with the specialist executor so the hooks + // dispatcher (built later in startup) can be attached before any specialist runs. + lifecycleHooks *specialist.LifecycleHooks } // specialistInfos returns the runtime's specialist summaries, nil-safe so the @@ -1055,12 +1059,26 @@ func (r *agentToolRuntime) specialistInfos() []agent.SpecialistInfo { return r.specialists } +func (r *agentToolRuntime) attachLifecycleHooks(dispatcher *hooks.Dispatcher) { + if r == nil || r.lifecycleHooks == nil || dispatcher == nil { + return + } + r.lifecycleHooks.Dispatch = func(ctx context.Context, event string, specialistName string, payload map[string]any) { + dispatcher.Dispatch(ctx, hooks.DispatchInput{ + Event: hooks.Event(event), + ToolName: specialistName, + Payload: payload, + }) + } +} + func registerSpecialistTools(registry *tools.Registry, workspaceRoot string, maxTeamSize int) (*agentToolRuntime, error) { paths, err := specialist.DefaultPaths(workspaceRoot) if err != nil { return nil, err } - executor := specialist.Executor{Paths: paths} + lifecycleHooks := &specialist.LifecycleHooks{} + executor := specialist.Executor{Paths: paths, LifecycleHooks: lifecycleHooks} runtime, err := specialist.RegisterTools(registry, executor) if err != nil { return nil, err @@ -1079,7 +1097,12 @@ func registerSpecialistTools(registry *tools.Registry, workspaceRoot string, max return nil, err } swarm.RegisterTools(registry, sw) - return &agentToolRuntime{specialist: runtime, swarm: sw, specialists: specialistSummaries(paths)}, nil + return &agentToolRuntime{ + specialist: runtime, + swarm: sw, + specialists: specialistSummaries(paths), + lifecycleHooks: lifecycleHooks, + }, nil } // specialistSummaries loads the available specialists (built-ins + user/project diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 2d1fe542a..5de8f736a 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -625,6 +625,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in // can be combined with the plugin activation's, and emit at most one notice when // project hooks/plugins were dropped for an untrusted workspace. hookDispatcher, hookSkip := newHookDispatcherWithExtra(workspaceRoot, pluginActivation.hooks, trustRoot, executionRunner) + specialistRuntime.attachLifecycleHooks(hookDispatcher) emitTrustNotice(stderr, hookSkip, pluginActivation.trustSkip, mcpSkip) result, err := agent.Run(runCtx, agentPrompt, provider, agent.Options{ MaxTurns: resolved.MaxTurns, diff --git a/internal/cli/hooks_manage.go b/internal/cli/hooks_manage.go index e2bc41493..e564f2c9e 100644 --- a/internal/cli/hooks_manage.go +++ b/internal/cli/hooks_manage.go @@ -319,7 +319,7 @@ Flags: --command Command to run (required) --name Human-readable hook name --description Hook description - --matcher Tool matcher (beforeTool/afterTool only) + --matcher Tool/specialist matcher (beforeTool/afterTool/specialistStart/specialistStop) --arg Command argument (repeatable) --user Write to user config instead of the project --json Print command result as JSON diff --git a/internal/cli/specialist_lifecycle_hooks_test.go b/internal/cli/specialist_lifecycle_hooks_test.go new file mode 100644 index 000000000..2a2b1eeff --- /dev/null +++ b/internal/cli/specialist_lifecycle_hooks_test.go @@ -0,0 +1,43 @@ +package cli + +import ( + "io" + "testing" + + "github.com/Gitlawb/zero/internal/hooks" + "github.com/Gitlawb/zero/internal/specialist" + "github.com/Gitlawb/zero/internal/tools" +) + +func TestRegisterSpecialistToolsWiresLifecycleHooksIntoTaskOutput(t *testing.T) { + registry := tools.NewRegistry() + runtime, err := registerSpecialistTools(registry, t.TempDir(), 0) + if err != nil { + t.Fatalf("registerSpecialistTools returned error: %v", err) + } + t.Cleanup(func() { closeSpecialistRuntime(io.Discard, runtime) }) + + if runtime.lifecycleHooks == nil { + t.Fatal("lifecycleHooks is nil after registerSpecialistTools") + } + raw, ok := registry.Get("TaskOutput") + if !ok { + t.Fatal("TaskOutput tool not registered") + } + outputTool, ok := raw.(*specialist.OutputTool) + if !ok { + t.Fatalf("TaskOutput type = %T, want *specialist.OutputTool", raw) + } + if outputTool.LifecycleHooks != runtime.lifecycleHooks { + t.Fatal("TaskOutput LifecycleHooks is not the shared registerSpecialistTools bridge") + } + + dispatcher := hooks.NewDispatcher(hooks.DispatcherOptions{}) + runtime.attachLifecycleHooks(dispatcher) + if runtime.lifecycleHooks.Dispatch == nil { + t.Fatal("attachLifecycleHooks left Dispatch nil") + } + if outputTool.LifecycleHooks.Dispatch == nil { + t.Fatal("TaskOutput did not observe Dispatch after attachLifecycleHooks") + } +} diff --git a/internal/hooks/hooks.go b/internal/hooks/hooks.go index f7dd79cea..8a3e95d50 100644 --- a/internal/hooks/hooks.go +++ b/internal/hooks/hooks.go @@ -779,7 +779,7 @@ func normalizeDefinition(raw any, field string) (Definition, error) { return Definition{}, err } if matcher != "" && !eventSupportsMatcher(event) { - return Definition{}, manifestError{fieldPath: field + ".matcher", message: "matcher is only supported for beforeTool and afterTool hooks."} + return Definition{}, manifestError{fieldPath: field + ".matcher", message: "matcher is only supported for beforeTool, afterTool, specialistStart, and specialistStop hooks."} } command, err := requiredString(obj, field+".command") if err != nil { @@ -843,7 +843,7 @@ func matchesHookMatcher(matcher string, toolName string) bool { } func eventSupportsMatcher(event Event) bool { - return event == EventBeforeTool || event == EventAfterTool + return event == EventBeforeTool || event == EventAfterTool || event == EventSpecialistStart || event == EventSpecialistStop } func toDiagnostic(err error, source ConfigSource, path string) Diagnostic { diff --git a/internal/hooks/hooks_test.go b/internal/hooks/hooks_test.go index a12580512..f24b6d1b4 100644 --- a/internal/hooks/hooks_test.go +++ b/internal/hooks/hooks_test.go @@ -180,7 +180,7 @@ func TestLoadConfigPreservesUserDisabledStateWhenProjectOmitsEnabled(t *testing. } func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) { - for _, event := range []string{"sessionStart", "specialistStart"} { + for _, event := range []string{"sessionStart", "sessionEnd"} { t.Run(event, func(t *testing.T) { dir := t.TempDir() projectConfigPath := filepath.Join(dir, "hooks.json") @@ -210,6 +210,46 @@ func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) { } } +func TestLoadConfigAcceptsMatchersOnSpecialistHooks(t *testing.T) { + for _, tc := range []struct { + event Event + id string + }{ + {event: EventSpecialistStart, id: "zero.explorer-start"}, + {event: EventSpecialistStop, id: "zero.explorer-stop"}, + } { + t.Run(string(tc.event), func(t *testing.T) { + dir := t.TempDir() + projectConfigPath := filepath.Join(dir, "hooks.json") + writeHookJSON(t, projectConfigPath, map[string]any{ + "hooks": []any{map[string]any{ + "id": tc.id, + "event": string(tc.event), + "matcher": "explorer", + "command": "node", + }}, + }) + + result, err := LoadConfig(LoadOptions{ + UserConfigPath: filepath.Join(dir, "missing-user-hooks.json"), + ProjectConfigPath: projectConfigPath, + }) + if err != nil { + t.Fatalf("LoadConfig returned error: %v", err) + } + if len(result.Config.Hooks) != 1 || result.Config.Hooks[0].Matcher != "explorer" { + t.Fatalf("specialist matcher hook = %#v", result.Config.Hooks) + } + if got := hookIDs(Select(result.Config, SelectInput{Event: tc.event, ToolName: "explorer"})); !reflect.DeepEqual(got, []string{tc.id}) { + t.Fatalf("matched selection = %#v", got) + } + if got := Select(result.Config, SelectInput{Event: tc.event, ToolName: "worker"}); len(got) != 0 { + t.Fatalf("unmatched specialist should be empty, got %#v", got) + } + }) + } +} + func TestConfigStorePersistsUpdates(t *testing.T) { configPath := filepath.Join(t.TempDir(), "hooks.json") store, err := NewConfigStore(StoreOptions{ConfigPath: configPath}) diff --git a/internal/specialist/accounting.go b/internal/specialist/accounting.go index 5a072a36b..485a7bcb6 100644 --- a/internal/specialist/accounting.go +++ b/internal/specialist/accounting.go @@ -1,6 +1,7 @@ package specialist import ( + "context" "encoding/json" "strings" "sync" @@ -36,12 +37,11 @@ type specialistAccountingInput struct { func (executor Executor) recordSpecialistStart(input specialistAccountingInput) { payload := baseSpecialistPayload(input) _, _ = appendSpecialistSessionEvent(executor.SessionStore, input.ParentSessionID, sessions.EventSpecialistStart, payload) + executor.dispatchLifecycleHook("specialistStart", input, payload) } func (executor Executor) recordSpecialistStop(input specialistAccountingInput, summary StreamResult, status string, exitCode int, runErr error, usageRolledUp bool) { store := accountingStore(executor.SessionStore) - accountingMu.Lock() - defer accountingMu.Unlock() payload := baseSpecialistPayload(input) if summary.RunID != "" { payload["runId"] = summary.RunID @@ -55,9 +55,30 @@ func (executor Executor) recordSpecialistStop(input specialistAccountingInput, s if len(summary.Errors) > 0 { payload["errors"] = append([]string(nil), summary.Errors...) } - // Atomic check+append under the session lock so a concurrent stop path cannot - // also pass the existence check and write a duplicate stop event. + // Hold accountingMu only for the deduped append so a slow lifecycle hook cannot + // stall concurrent specialist stop/usage accounting. + accountingMu.Lock() _, _ = appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID) + accountingMu.Unlock() + // Dispatch is claimed on this LifecycleHooks bridge, not on session append. + // Swarm members may have an empty parent session (append never succeeds) and + // TaskOutput may race Task.onExit; hooks must still fire exactly once when a + // hook-carrying executor observes the stop. + if executor.LifecycleHooks.claimStopDispatch(input.ChildSessionID, summary.RunID) { + executor.dispatchLifecycleHook("specialistStop", input, payload) + } +} + +func (executor Executor) dispatchLifecycleHook(event string, input specialistAccountingInput, payload map[string]any) { + if executor.LifecycleHooks == nil || executor.LifecycleHooks.Dispatch == nil { + return + } + hookPayload := make(map[string]any, len(payload)+1) + for key, value := range payload { + hookPayload[key] = value + } + hookPayload["event"] = event + executor.LifecycleHooks.Dispatch(context.Background(), event, strings.TrimSpace(input.SpecialistName), hookPayload) } func (executor Executor) rollUpSpecialistUsage(input specialistAccountingInput, summary StreamResult) bool { diff --git a/internal/specialist/accounting_test.go b/internal/specialist/accounting_test.go index 7d9542501..0ba158374 100644 --- a/internal/specialist/accounting_test.go +++ b/internal/specialist/accounting_test.go @@ -6,6 +6,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "testing" "github.com/Gitlawb/zero/internal/background" @@ -21,10 +22,21 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) { t.Fatalf("Create parent returned error: %v", err) } zero := 0 + var hookEvents []string + var hookNames []string executor := Executor{ BinaryPath: "/usr/local/bin/zero", SessionStore: store, NewSessionID: func() (string, error) { return "child_task", nil }, + LifecycleHooks: &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, specialistName string, payload map[string]any) { + hookEvents = append(hookEvents, event) + hookNames = append(hookNames, specialistName) + if payload["event"] != event { + t.Errorf("payload event = %#v, want %q", payload["event"], event) + } + }, + }, Load: func(LoadOptions) (LoadResult, error) { return LoadResult{Specialists: []Manifest{{ Metadata: Metadata{Name: "worker", Description: "Does focused work"}, @@ -59,6 +71,12 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) { if result.Result.Status != tools.StatusOK { t.Fatalf("Run status = %s, output=%s", result.Result.Status, result.Result.Output) } + if got, want := strings.Join(hookEvents, ","), "specialistStart,specialistStop"; got != want { + t.Fatalf("lifecycle hooks = %q, want %q", got, want) + } + if got, want := strings.Join(hookNames, ","), "worker,worker"; got != want { + t.Fatalf("lifecycle specialist names = %q, want %q", got, want) + } events, err := store.ReadEvents(parent.SessionID) if err != nil { @@ -286,7 +304,17 @@ func TestRecordSpecialistStopDedupesUnderConcurrency(t *testing.T) { if err != nil { t.Fatalf("Create parent returned error: %v", err) } - executor := Executor{SessionStore: store} + var hookStops atomic.Int32 + executor := Executor{ + SessionStore: store, + LifecycleHooks: &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + if event == "specialistStop" { + hookStops.Add(1) + } + }, + }, + } input := specialistAccountingInput{ ParentSessionID: parent.SessionID, ChildSessionID: "child_task", @@ -321,4 +349,143 @@ func TestRecordSpecialistStopDedupesUnderConcurrency(t *testing.T) { if stops != 1 { t.Fatalf("expected exactly 1 stop event under concurrency, got %d", stops) } + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d, want 1 under concurrency", got) + } +} + +func TestRecordSpecialistStopFiresWithoutParentSession(t *testing.T) { + var hookStops atomic.Int32 + executor := Executor{ + LifecycleHooks: &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, name string, _ map[string]any) { + if event == "specialistStop" { + hookStops.Add(1) + if name != "subagent" { + t.Errorf("specialist name = %q, want subagent", name) + } + } + }, + }, + } + executor.recordSpecialistStop(specialistAccountingInput{ + // Swarm members omit ParentSessionID; append never succeeds. + ChildSessionID: "member_child", + SpecialistName: "subagent", + Mode: "foreground", + }, StreamResult{RunID: "run_swarm"}, "success", 0, nil, false) + + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d, want 1 with empty parent", got) + } + // Second observation of the same child/run must not re-fire. + executor.recordSpecialistStop(specialistAccountingInput{ + ChildSessionID: "member_child", + SpecialistName: "subagent", + Mode: "foreground", + }, StreamResult{RunID: "run_swarm"}, "success", 0, nil, false) + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d after duplicate, want 1", got) + } +} + +func TestRecordSpecialistStopHookLessWinnerDoesNotSuppressLaterHooks(t *testing.T) { + store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()}) + parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"}) + if err != nil { + t.Fatalf("Create parent returned error: %v", err) + } + input := specialistAccountingInput{ + ParentSessionID: parent.SessionID, + ChildSessionID: "child_task", + SpecialistName: "worker", + Mode: "background", + Background: true, + } + summary := StreamResult{RunID: "run_1"} + + // Hook-less TaskOutput path wins the session append race. + Executor{SessionStore: store}.recordSpecialistStop(input, summary, "success", 0, nil, true) + + var hookStops atomic.Int32 + hooks := &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + if event == "specialistStop" { + hookStops.Add(1) + } + }, + } + Executor{SessionStore: store, LifecycleHooks: hooks}.recordSpecialistStop(input, summary, "success", 0, nil, true) + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d, want 1 after hook-less winner", got) + } +} + +func TestOutputToolSharesLifecycleHooksWithStopAccounting(t *testing.T) { + store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()}) + parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"}) + if err != nil { + t.Fatalf("Create parent returned error: %v", err) + } + manager, err := background.NewManager(t.TempDir()) + if err != nil { + t.Fatalf("NewManager returned error: %v", err) + } + outputFile, err := manager.Register(background.RegisterInput{ + TaskID: "child_task", + Type: "specialist", + SpecialistName: "worker", + Description: "Auth check", + ParentID: parent.SessionID, + }) + if err != nil { + t.Fatalf("Register returned error: %v", err) + } + if err := os.WriteFile(outputFile, []byte(strings.Join([]string{ + `{"schemaVersion":2,"type":"run_start","runId":"run_1","sessionId":"child_task"}`, + `{"schemaVersion":2,"type":"final","runId":"run_1","text":"done"}`, + `{"schemaVersion":2,"type":"run_end","runId":"run_1","status":"success","exitCode":0}`, + "", + }, "\n")), 0o600); err != nil { + t.Fatalf("WriteFile returned error: %v", err) + } + if err := manager.UpdateStatus("child_task", background.StatusCompleted, 0); err != nil { + t.Fatalf("UpdateStatus returned error: %v", err) + } + + var hookStops atomic.Int32 + hooks := &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + if event == "specialistStop" { + hookStops.Add(1) + } + }, + } + tool := NewOutputTool(manager) + tool.SessionStore = store + tool.LifecycleHooks = hooks + + if result := tool.Run(context.Background(), map[string]any{"task_id": "child_task"}); result.Status != tools.StatusOK { + t.Fatalf("TaskOutput status = %s output=%s", result.Status, result.Output) + } + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d, want 1 from TaskOutput", got) + } + // Shared bridge: a later Task.onExit-style caller must not re-fire. + Executor{SessionStore: store, LifecycleHooks: hooks}.recordBackgroundTaskAccounting( + mustGetBackgroundTask(t, manager, "child_task"), + StreamResult{RunID: "run_1"}, + ) + if got := hookStops.Load(); got != 1 { + t.Fatalf("specialistStop hooks = %d after second path, want 1", got) + } +} + +func mustGetBackgroundTask(t *testing.T, manager *background.Manager, taskID string) background.Task { + t.Helper() + task, ok := manager.Get(taskID) + if !ok { + t.Fatalf("background task %q not found", taskID) + } + return task } diff --git a/internal/specialist/exec.go b/internal/specialist/exec.go index 6dac1d259..6efc90994 100644 --- a/internal/specialist/exec.go +++ b/internal/specialist/exec.go @@ -15,6 +15,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "github.com/Gitlawb/zero/internal/background" "github.com/Gitlawb/zero/internal/sessions" @@ -50,6 +51,38 @@ type Executor struct { BackgroundManager *background.Manager BackgroundManagerFunc BackgroundManagerFunc BackgroundRuntime *Runtime + // LifecycleHooks optionally dispatches specialistStart/specialistStop hooks. + // The shared pointer lets the CLI attach the hooks dispatcher after specialist + // tools are registered (hooks are constructed later in startup). + LifecycleHooks *LifecycleHooks +} + +// LifecycleHooks bridges specialist accounting into the hooks dispatcher without +// importing the hooks package into every specialist call site. +type LifecycleHooks struct { + // Dispatch runs configured hooks for event ("specialistStart"/"specialistStop"). + // specialistName is used as the matcher subject. nil Dispatch is a no-op. + Dispatch func(ctx context.Context, event string, specialistName string, payload map[string]any) + // stopOnce dedupes specialistStop dispatch across concurrent finishers that + // share this bridge (Task onExit, TaskOutput poll, swarm members). + stopOnce sync.Map +} + +// claimStopDispatch reports whether this bridge should run specialistStop for +// the given child/run. The first observer wins; later finishers return false so +// hooks are not re-run. Decoupled from session-event append so swarm members +// without a parent session and hook-carrying reap paths still fire once. +func (h *LifecycleHooks) claimStopDispatch(childSessionID, runID string) bool { + if h == nil || h.Dispatch == nil { + return false + } + childSessionID = strings.TrimSpace(childSessionID) + if childSessionID == "" { + return true + } + key := childSessionID + "\x00" + strings.TrimSpace(runID) + _, loaded := h.stopOnce.LoadOrStore(key, struct{}{}) + return !loaded } type BuildArgsInput struct { diff --git a/internal/specialist/output_tool.go b/internal/specialist/output_tool.go index 3c4dd4667..1668e7760 100644 --- a/internal/specialist/output_tool.go +++ b/internal/specialist/output_tool.go @@ -21,9 +21,12 @@ const ( ) type OutputTool struct { - manager *background.Manager - managerFunc BackgroundManagerFunc - SessionStore *sessions.Store + manager *background.Manager + managerFunc BackgroundManagerFunc + SessionStore *sessions.Store + // LifecycleHooks must be the same bridge as Task so onExit and TaskOutput + // poll share stop-hook dedup. nil is safe (accounting still runs). + LifecycleHooks *LifecycleHooks PollInterval time.Duration DefaultTimeout time.Duration MaxTimeout time.Duration @@ -35,8 +38,8 @@ type outputParameters struct { TimeoutMS int } -func newOutputToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store) *OutputTool { - return &OutputTool{managerFunc: managerFunc, SessionStore: sessionStore} +func newOutputToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store, lifecycleHooks *LifecycleHooks) *OutputTool { + return &OutputTool{managerFunc: managerFunc, SessionStore: sessionStore, LifecycleHooks: lifecycleHooks} } func (tool *OutputTool) Name() string { @@ -143,7 +146,10 @@ func (tool *OutputTool) readOutput(task background.Task) tools.Result { dataString := string(data) summary, rawLines := summarizeTaskData(dataString, task.ExitCode) if task.Status != background.StatusRunning { - Executor{SessionStore: tool.SessionStore}.recordBackgroundTaskAccounting(task, summary) + Executor{ + SessionStore: tool.SessionStore, + LifecycleHooks: tool.LifecycleHooks, + }.recordBackgroundTaskAccounting(task, summary) } return tools.Result{ Status: tools.StatusOK, diff --git a/internal/specialist/registry.go b/internal/specialist/registry.go index be1c10bde..e182b5727 100644 --- a/internal/specialist/registry.go +++ b/internal/specialist/registry.go @@ -14,7 +14,7 @@ func RegisterTools(registry *tools.Registry, executor Executor) (*Runtime, error toolExecutor.BackgroundRuntime = runtime toolExecutor.BackgroundManagerFunc = runtime.Manager registry.Register(NewTaskTool(toolExecutor)) - registry.Register(newOutputToolWithManagerFunc(runtime.Manager, executor.SessionStore)) + registry.Register(newOutputToolWithManagerFunc(runtime.Manager, executor.SessionStore, executor.LifecycleHooks)) registry.Register(newStopToolWithManagerFunc(runtime.Manager)) registry.Register(NewGenerateTool(NewStorage(executor.Paths))) return runtime, nil