From 97910f559322fa8d8688ccf7add49aadd5530609 Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Fri, 24 Jul 2026 02:59:20 +0530 Subject: [PATCH 1/3] fix(hooks): dispatch specialistStart/Stop and allow matchers These events were documented and accepted by hooks add, but never fired. Wire advisory dispatch from specialist accounting and allow specialist name matchers so EXTENDING.md matches runtime behavior. --- docs/EXTENDING.md | 2 +- internal/cli/app.go | 27 ++++++++++++++++++++-- internal/cli/exec.go | 1 + internal/cli/hooks_manage.go | 2 +- internal/hooks/hooks.go | 4 ++-- internal/hooks/hooks_test.go | 32 +++++++++++++++++++++++++- internal/specialist/accounting.go | 15 ++++++++++++ internal/specialist/accounting_test.go | 17 ++++++++++++++ internal/specialist/exec.go | 12 ++++++++++ 9 files changed, 105 insertions(+), 7 deletions(-) diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index 18b2c1bc4..bf1f046c4 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -192,7 +192,7 @@ 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. +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/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..1939b6ae2 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,36 @@ func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) { } } +func TestLoadConfigAcceptsMatchersOnSpecialistHooks(t *testing.T) { + dir := t.TempDir() + projectConfigPath := filepath.Join(dir, "hooks.json") + writeHookJSON(t, projectConfigPath, map[string]any{ + "hooks": []any{map[string]any{ + "id": "zero.explorer-start", + "event": "specialistStart", + "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: EventSpecialistStart, ToolName: "explorer"})); !reflect.DeepEqual(got, []string{"zero.explorer-start"}) { + t.Fatalf("matched selection = %#v", got) + } + if got := Select(result.Config, SelectInput{Event: EventSpecialistStart, 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..72770ae3a 100644 --- a/internal/specialist/accounting.go +++ b/internal/specialist/accounting.go @@ -1,6 +1,7 @@ package specialist import ( + "context" "encoding/json" "strings" "sync" @@ -36,6 +37,7 @@ 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) { @@ -58,6 +60,19 @@ func (executor Executor) recordSpecialistStop(input specialistAccountingInput, s // Atomic check+append under the session lock so a concurrent stop path cannot // also pass the existence check and write a duplicate stop event. _, _ = appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, 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..d6bdcd773 100644 --- a/internal/specialist/accounting_test.go +++ b/internal/specialist/accounting_test.go @@ -21,10 +21,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 +70,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 { diff --git a/internal/specialist/exec.go b/internal/specialist/exec.go index 6dac1d259..69141ea9a 100644 --- a/internal/specialist/exec.go +++ b/internal/specialist/exec.go @@ -50,6 +50,18 @@ 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) } type BuildArgsInput struct { From 19f98c2b3e7992ebdedc7c0e0c96f1b83287514e Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Sun, 26 Jul 2026 16:21:29 +0530 Subject: [PATCH 2/3] fix(hooks): unlock before specialistStop hooks and fire once Release accountingMu before running lifecycle hooks, and dispatch specialistStop only when the stop event is newly recorded so concurrent finishers cannot re-run arbitrary hook commands. --- internal/hooks/hooks_test.go | 60 +++++++++++++++----------- internal/specialist/accounting.go | 16 ++++--- internal/specialist/accounting_test.go | 16 ++++++- 3 files changed, 60 insertions(+), 32 deletions(-) diff --git a/internal/hooks/hooks_test.go b/internal/hooks/hooks_test.go index 1939b6ae2..f24b6d1b4 100644 --- a/internal/hooks/hooks_test.go +++ b/internal/hooks/hooks_test.go @@ -211,32 +211,42 @@ func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) { } func TestLoadConfigAcceptsMatchersOnSpecialistHooks(t *testing.T) { - dir := t.TempDir() - projectConfigPath := filepath.Join(dir, "hooks.json") - writeHookJSON(t, projectConfigPath, map[string]any{ - "hooks": []any{map[string]any{ - "id": "zero.explorer-start", - "event": "specialistStart", - "matcher": "explorer", - "command": "node", - }}, - }) + 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: EventSpecialistStart, ToolName: "explorer"})); !reflect.DeepEqual(got, []string{"zero.explorer-start"}) { - t.Fatalf("matched selection = %#v", got) - } - if got := Select(result.Config, SelectInput{Event: EventSpecialistStart, ToolName: "worker"}); len(got) != 0 { - t.Fatalf("unmatched specialist should be empty, got %#v", got) + 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) + } + }) } } diff --git a/internal/specialist/accounting.go b/internal/specialist/accounting.go index 72770ae3a..cab68c90c 100644 --- a/internal/specialist/accounting.go +++ b/internal/specialist/accounting.go @@ -42,8 +42,6 @@ func (executor Executor) recordSpecialistStart(input specialistAccountingInput) 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 @@ -57,10 +55,16 @@ 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. - _, _ = appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID) - executor.dispatchLifecycleHook("specialistStop", input, payload) + // Hold accountingMu only for the deduped append so a slow lifecycle hook cannot + // stall concurrent specialist stop/usage accounting. + accountingMu.Lock() + appended, _ := appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID) + accountingMu.Unlock() + // Duplicate stop paths (onExit + TaskOutput poll, swarm races) must not re-run + // arbitrary configured hooks for the same specialist stop. + if appended { + executor.dispatchLifecycleHook("specialistStop", input, payload) + } } func (executor Executor) dispatchLifecycleHook(event string, input specialistAccountingInput, payload map[string]any) { diff --git a/internal/specialist/accounting_test.go b/internal/specialist/accounting_test.go index d6bdcd773..e011e00da 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" @@ -303,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", @@ -338,4 +349,7 @@ 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) + } } From 7d55f5680ec40aa674fa030b6e40918bc29a294b Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Mon, 27 Jul 2026 00:00:10 +0530 Subject: [PATCH 3/3] fix(hooks): fire specialistStop without relying on session append Decouple stop-hook dispatch from parent-session event append so swarm members (empty parent) and TaskOutput still get exactly one hook via a shared LifecycleHooks claim, and wire that bridge into TaskOutput. --- docs/EXTENDING.md | 2 + .../cli/specialist_lifecycle_hooks_test.go | 43 ++++++ internal/specialist/accounting.go | 10 +- internal/specialist/accounting_test.go | 136 ++++++++++++++++++ internal/specialist/exec.go | 21 +++ internal/specialist/output_tool.go | 18 ++- internal/specialist/registry.go | 2 +- 7 files changed, 221 insertions(+), 11 deletions(-) create mode 100644 internal/cli/specialist_lifecycle_hooks_test.go diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index bf1f046c4..00698c1c6 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -192,6 +192,8 @@ Events the agent emits (in dispatch order): | `specialistStart` | A sub-agent is spawned | yes (specialist name) | | `specialistStop` | A sub-agent ends | yes (specialist name) | +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/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/specialist/accounting.go b/internal/specialist/accounting.go index cab68c90c..485a7bcb6 100644 --- a/internal/specialist/accounting.go +++ b/internal/specialist/accounting.go @@ -58,11 +58,13 @@ func (executor Executor) recordSpecialistStop(input specialistAccountingInput, s // Hold accountingMu only for the deduped append so a slow lifecycle hook cannot // stall concurrent specialist stop/usage accounting. accountingMu.Lock() - appended, _ := appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID) + _, _ = appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID) accountingMu.Unlock() - // Duplicate stop paths (onExit + TaskOutput poll, swarm races) must not re-run - // arbitrary configured hooks for the same specialist stop. - if appended { + // 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) } } diff --git a/internal/specialist/accounting_test.go b/internal/specialist/accounting_test.go index e011e00da..0ba158374 100644 --- a/internal/specialist/accounting_test.go +++ b/internal/specialist/accounting_test.go @@ -353,3 +353,139 @@ func TestRecordSpecialistStopDedupesUnderConcurrency(t *testing.T) { 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 69141ea9a..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" @@ -62,6 +63,26 @@ 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