diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index 96d18b464..51df9f001 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -192,7 +192,11 @@ 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. + +`specialistStart` fires once the child process has actually started; a launch that fails before exec (missing binary, spawn error) records no public lifecycle hooks. `specialistStop` fires once per terminal stop: completion, error, explicit stop of a recovered task, and a background task reaped by a later `TaskOutput` all share one durable claim, so a hook that notifies or cleans up does not run twice for the same sub-agent run. + +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 ea5b09c20..7e2514442 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -964,6 +964,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) // From the engine the session's commands run through, so the notice // describes the sandbox they actually get. @@ -1221,6 +1222,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 @@ -1232,12 +1236,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 @@ -1256,7 +1274,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 15842b2bb..73bf1cf04 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -655,6 +655,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 fef6fff64..37bd6809c 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..c014b62fc --- /dev/null +++ b/internal/cli/specialist_lifecycle_hooks_test.go @@ -0,0 +1,145 @@ +package cli + +import ( + "bytes" + "context" + "io" + "testing" + + "github.com/Twigpine/zero/internal/config" + "github.com/Twigpine/zero/internal/hooks" + "github.com/Twigpine/zero/internal/mcp" + "github.com/Twigpine/zero/internal/specialist" + "github.com/Twigpine/zero/internal/tools" + "github.com/Twigpine/zero/internal/tui" + "github.com/Twigpine/zero/internal/zeroruntime" +) + +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") + } +} + +// TestRealEntryPointsAttachLifecycleHooks closes the gap the registration-level +// test above cannot see: wiring the bridge when tools register is not enough +// unless the real startup paths actually call attachLifecycleHooks before any +// run. Both surfaces that can launch specialists must do so — the interactive +// TUI and headless exec. attachLifecycleHooks mutates the same +// *specialist.LifecycleHooks the TaskOutput tool holds, so a nil Dispatch on +// that tool after a full startup proves the call never ran and the +// specialistStart/specialistStop hooks would silently never fire. +func TestRealEntryPointsAttachLifecycleHooks(t *testing.T) { + requireTaskOutputDispatchAttached := func(t *testing.T, registry *tools.Registry) { + t.Helper() + if registry == nil { + t.Fatal("startup did not reach the run registry") + } + raw, ok := registry.Get("TaskOutput") + if !ok { + t.Fatal("TaskOutput tool not registered by startup") + } + outputTool, ok := raw.(*specialist.OutputTool) + if !ok { + t.Fatalf("TaskOutput type = %T, want *specialist.OutputTool", raw) + } + if outputTool.LifecycleHooks == nil || outputTool.LifecycleHooks.Dispatch == nil { + t.Fatal("startup left the specialist lifecycle bridge unattached: specialistStart/specialistStop hooks would never fire") + } + } + + t.Run("interactive TUI", func(t *testing.T) { + isolateCLIUserState(t) + cwd := t.TempDir() + var registry *tools.Registry + + // Empty args route runWithDeps to the interactive TUI path; the injected + // runTUI receives the same registry specialist tools registered into. + exitCode := runWithDeps([]string{}, io.Discard, io.Discard, appDeps{ + getwd: func() (string, error) { return cwd, nil }, + resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{ + Provider: config.ProviderProfile{ + Name: "p", + ProviderKind: config.ProviderKindOpenAICompatible, + BaseURL: "http://127.0.0.1/v1", + Model: "m", + }, + }, nil + }, + resolveMCPConfig: func(string, bool) (config.MCPConfig, error) { return config.MCPConfig{}, nil }, + runTUI: func(_ context.Context, options tui.Options) int { + registry = options.Registry + return exitSuccess + }, + }) + if exitCode != exitSuccess { + t.Fatalf("runWithDeps exit = %d, want %d", exitCode, exitSuccess) + } + requireTaskOutputDispatchAttached(t, registry) + }) + + t.Run("headless exec", func(t *testing.T) { + isolateCLIUserState(t) + cwd := t.TempDir() + var registry *tools.Registry + var stdout, stderr bytes.Buffer + + // --auto high is what registers specialist tools on the exec surface. + // MCP registration is the one startup chokepoint that sees the run + // registry before attachLifecycleHooks runs, so capture it there and + // assert after the run completes. + exitCode := runWithDeps([]string{ + "exec", "--auto", "high", "--cwd", cwd, "--init-session-id", "hook_wireup", "hello", + }, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return cwd, nil }, + resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return execResolvedConfig(), nil + }, + newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) { + return echoExecProvider{}, nil + }, + resolveMCPConfig: func(string, bool) (config.MCPConfig, error) { + return config.MCPConfig{Servers: map[string]config.MCPServerConfig{ + "wireup": {Type: "stdio", Command: "wireup-mcp"}, + }}, nil + }, + newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil }, + registerMCPTools: func(_ context.Context, captured *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) { + registry = captured + return closeFunc(func() error { return nil }), nil + }, + }) + if exitCode != exitSuccess { + t.Fatalf("runWithDeps exit = %d, want %d (stderr=%q)", exitCode, exitSuccess, stderr.String()) + } + requireTaskOutputDispatchAttached(t, registry) + }) +} 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 5cdc62c46..265cffeac 100644 --- a/internal/specialist/accounting.go +++ b/internal/specialist/accounting.go @@ -1,7 +1,9 @@ package specialist import ( + "context" "encoding/json" + "os" "strings" "sync" @@ -31,18 +33,41 @@ type specialistAccountingInput struct { Mode string Background bool PID int + // LaunchFailed marks a child that never started (e.g. the command failed to + // exec). The internal start/stop accounting still records the attempt, but no + // public specialistStart/specialistStop hook fires for it: there was no + // sub-agent lifecycle to observe. The zero value is "started", so a caller + // that forgets to set it never silently suppresses hooks. + LaunchFailed bool } +// recordSpecialistStart records the internal attempt accounting. It deliberately +// does NOT dispatch the public specialistStart hook: the hook fires only once +// the child has actually started (dispatchSpecialistStartHook), so a failed +// launch cannot announce a sub-agent that never existed. func (executor Executor) recordSpecialistStart(input specialistAccountingInput) { payload := baseSpecialistPayload(input) _, _ = appendSpecialistSessionEvent(executor.SessionStore, input.ParentSessionID, sessions.EventSpecialistStart, payload) } +// dispatchSpecialistStartHook publishes the specialistStart hook for a child +// whose startup was confirmed. Callers pass the same accounting input used for +// the attempt record so payload and matcher subject stay identical. +func (executor Executor) dispatchSpecialistStartHook(input specialistAccountingInput) { + if input.LaunchFailed || executor.LifecycleHooks == nil || executor.LifecycleHooks.Dispatch == nil { + return + } + executor.dispatchLifecycleHook("specialistStart", input, baseSpecialistPayload(input)) +} + 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) + // childStarted is the durable eligibility flag for the public stop hook: a + // later reaper (TaskOutput poll, recovered TaskStop, onExit in another + // process) derives whether the terminal transition deserves a hook from the + // recorded stop event, since it cannot know whether the child ever started. + payload["childStarted"] = !input.LaunchFailed if summary.RunID != "" { payload["runId"] = summary.RunID } @@ -55,9 +80,120 @@ 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() + executor.dispatchSpecialistStopHook(store, input, payload, summary.RunID) +} + +// specialistStopHookEvent durably records that the specialistStop hook dispatch +// decision for one child/run was settled. It lives in the parent session's event +// log - the same cross-process, lock-guarded ledger the accounting dedup uses - +// so a later Zero process reaping the same task (fresh LifecycleHooks bridge) +// does not re-run the hook. This is internal bookkeeping, not a lifecycle event. +const specialistStopHookEvent sessions.EventType = "specialist_stop_hook" + +// dispatchSpecialistStopHook runs the specialistStop hook at most once per +// child/run. The claim follows the accounting equivalence rules (an empty runId +// is a catch-all for the same child), so concurrent or sequential finishers - +// Task.onExit, a TaskOutput poll in another process, a recovered TaskStop - share +// one durable claim. Swarm members have no parent session (nothing durable to +// key on), so they fall back to this bridge's in-process claim, which mirrors +// the same run-ID equivalence. +// +// Hooks are external side effects, so on a store error the durable claim cannot +// be established; the hook then falls back to the in-process claim rather than +// running unchecked, keeping at-most-once within this process. +func (executor Executor) dispatchSpecialistStopHook(store *sessions.Store, input specialistAccountingInput, payload map[string]any, runID string) { + hooks := executor.LifecycleHooks + if hooks == nil || hooks.Dispatch == nil { + return + } + started := !input.LaunchFailed + parentSessionID := strings.TrimSpace(input.ParentSessionID) + if parentSessionID == "" || !sessions.ValidSessionID(parentSessionID) { + if started && hooks.claimInProcessStopDispatch(input.ChildSessionID, runID) { + executor.dispatchLifecycleHook("specialistStop", input, payload) + } + return + } + if recorded, ok := recordedStopChildStarted(store, parentSessionID, input.ChildSessionID, runID); ok { + started = recorded + } + claimed, err := appendSpecialistEventOnce(store, parentSessionID, specialistStopHookEvent, stopHookClaimPayload(input, runID, started), input.ChildSessionID, runID) + if err != nil { + // The durable ledger is unavailable. Keep the hook advisory but bounded to + // a single dispatch in this process rather than dropping it or risking an + // unbounded duplicate. + if started && hooks.claimInProcessStopDispatch(input.ChildSessionID, runID) { + executor.dispatchLifecycleHook("specialistStop", input, payload) + } + return + } + if claimed && started { + executor.dispatchLifecycleHook("specialistStop", input, payload) + } +} + +// recordedStopChildStarted reads the first recorded stop event matching the +// child/run equivalence rules and reports its childStarted eligibility flag. +// ok is false when no stop event exists yet or it predates the flag, in which +// case the caller keeps its locally known eligibility. +func recordedStopChildStarted(store *sessions.Store, parentSessionID, childSessionID, runID string) (started bool, ok bool) { + store = accountingStore(store) + childSessionID = strings.TrimSpace(childSessionID) + if childSessionID == "" { + return false, false + } + runID = strings.TrimSpace(runID) + events, err := store.ReadEvents(parentSessionID) + if err != nil { + return false, false + } + for _, event := range events { + if event.Type != sessions.EventSpecialistStop { + continue + } + var payload map[string]any + if err := json.Unmarshal(event.Payload, &payload); err != nil { + continue + } + if !specialistPayloadMatches(payload, childSessionID, runID) { + continue + } + if started, ok := payload["childStarted"].(bool); ok { + return started, true + } + return false, false + } + return false, false +} + +func stopHookClaimPayload(input specialistAccountingInput, runID string, started bool) map[string]any { + payload := baseSpecialistPayload(input) + payload["hook"] = "specialistStop" + // childStarted is the resolved eligibility (locally known or read back from + // the durable stop event); dispatched records that a hook was run for it. + payload["childStarted"] = started + payload["dispatched"] = started + if runID = strings.TrimSpace(runID); runID != "" { + payload["runId"] = runID + } + return 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 { @@ -82,6 +218,18 @@ func (executor Executor) recordBackgroundTaskAccounting(task background.Task, su executor.recordSpecialistStop(input, summary, string(task.Status), task.ExitCode, nil, rolledUp) } +// taskStreamSummary builds the stream summary for a finished background task +// from its output file. Every reaper (Task.onExit, TaskOutput, TaskStop) derives +// its runId this way, so the durable stop event and stop-hook claim agree on the +// run identity regardless of which path observed the task first. +func taskStreamSummary(task background.Task) StreamResult { + summary := StreamResult{ExitCode: task.ExitCode} + if data, err := os.ReadFile(task.OutputFile); err == nil { + summary, _ = summarizeTaskData(string(data), task.ExitCode) + } + return summary +} + func appendSpecialistUsageRollup(store *sessions.Store, input specialistAccountingInput, summary StreamResult) (bool, error) { if !summary.Usage.HasUsage() { return false, nil @@ -132,19 +280,7 @@ func specialistEventMatcher(eventType sessions.EventType, childSessionID, runID if err := json.Unmarshal(event.Payload, &payload); err != nil { continue } - if specialistPayloadString(payload, "source") != specialistAccountingSource { - continue - } - if specialistPayloadString(payload, "childSessionId") != childSessionID && specialistPayloadString(payload, "taskId") != childSessionID { - continue - } - // An already-recorded event with NO runId is a catch-all for this child - // (e.g. the immediate stop written when a PID couldn't be registered) and - // must match a later event that does carry a runId — otherwise the same - // child gets two stop/usage events. We match when: querying without a runId, - // the existing event has none, or the runIds are equal. - existingRunID := specialistPayloadString(payload, "runId") - if runID == "" || existingRunID == "" || existingRunID == runID { + if specialistPayloadMatches(payload, childSessionID, runID) { return true } } @@ -152,6 +288,25 @@ func specialistEventMatcher(eventType sessions.EventType, childSessionID, runID } } +// specialistPayloadMatches applies the shared "same specialist stop" identity to +// one event payload. The accounting dedup (specialistEventMatcher) and the +// durable stop-hook claim both use it, so the audit event and the hook side +// effect can never disagree about what counts as a repeat. An already-recorded +// event with NO runId is a catch-all for this child (e.g. the immediate stop +// written when a PID couldn't be registered) and matches a later event that +// carries a runId - otherwise the same child gets two stop/usage events. A query +// without a runId likewise matches any recorded run. +func specialistPayloadMatches(payload map[string]any, childSessionID, runID string) bool { + if specialistPayloadString(payload, "source") != specialistAccountingSource { + return false + } + if specialistPayloadString(payload, "childSessionId") != childSessionID && specialistPayloadString(payload, "taskId") != childSessionID { + return false + } + existingRunID := specialistPayloadString(payload, "runId") + return runID == "" || existingRunID == "" || existingRunID == runID +} + // appendSpecialistEventOnce atomically records a "record once" specialist event: // it appends payload only if no matching event already exists, under the session // store's lock so concurrent stop/usage callers (even cross-process) cannot both diff --git a/internal/specialist/accounting_test.go b/internal/specialist/accounting_test.go index d7d9df1b5..bf08aec44 100644 --- a/internal/specialist/accounting_test.go +++ b/internal/specialist/accounting_test.go @@ -3,10 +3,14 @@ package specialist import ( "context" "encoding/json" + "errors" "os" + "reflect" "strings" "sync" + "sync/atomic" "testing" + "time" "github.com/Twigpine/zero/internal/background" "github.com/Twigpine/zero/internal/sessions" @@ -21,10 +25,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,12 +74,18 @@ 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 { t.Fatalf("ReadEvents returned error: %v", err) } - if got, want := eventTypes(events), []sessions.EventType{sessions.EventSpecialistStart, sessions.EventUsage, sessions.EventSpecialistStop}; strings.Join(eventTypesString(got), ",") != strings.Join(eventTypesString(want), ",") { + if got, want := eventTypes(events), []sessions.EventType{sessions.EventSpecialistStart, sessions.EventUsage, sessions.EventSpecialistStop, specialistStopHookEvent}; strings.Join(eventTypesString(got), ",") != strings.Join(eventTypesString(want), ",") { t.Fatalf("event types = %#v, want %#v", got, want) } startPayload := eventPayload(t, events[0]) @@ -88,6 +109,13 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) { requirePayloadString(t, stopPayload, "status", "success") requirePayloadInt(t, stopPayload, "exitCode", 0) requirePayloadBool(t, stopPayload, "usageRolledUp", true) + requirePayloadBool(t, stopPayload, "childStarted", true) + + claimPayload := eventPayload(t, events[3]) + requirePayloadString(t, claimPayload, "hook", "specialistStop") + requirePayloadString(t, claimPayload, "runId", "run_1") + requirePayloadBool(t, claimPayload, "dispatched", true) + requirePayloadBool(t, claimPayload, "childStarted", true) } func TestExecutorRecordsStartedChildErrorExitCode(t *testing.T) { @@ -286,7 +314,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 +359,551 @@ 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) + } +} + +// A durable stop-hook claim must survive a process boundary: a later Zero +// process (fresh LifecycleHooks bridge over the same session store) reaping the +// same task must not re-run the hook even though its session-event append is +// rejected as a duplicate. +func TestRecordSpecialistStopHookClaimDurableAcrossBridges(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"} + + var firstBridge atomic.Int32 + Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&firstBridge)}.recordSpecialistStop(input, summary, "success", 0, nil, true) + if got := firstBridge.Load(); got != 1 { + t.Fatalf("first bridge stop hooks = %d, want 1", got) + } + + var secondBridge atomic.Int32 + Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&secondBridge)}.recordSpecialistStop(input, summary, "success", 0, nil, true) + if got := secondBridge.Load(); got != 0 { + t.Fatalf("second bridge stop hooks = %d, want 0: the durable claim must suppress a re-dispatch", got) + } + + events, err := store.ReadEvents(parent.SessionID) + if err != nil { + t.Fatalf("ReadEvents returned error: %v", err) + } + if claims := eventsOfType(events, specialistStopHookEvent); len(claims) != 1 { + t.Fatalf("stop-hook claim events = %d, want 1", len(claims)) + } + if stops := eventsOfType(events, sessions.EventSpecialistStop); len(stops) != 1 { + t.Fatalf("stop events = %d, want 1", len(stops)) + } +} + +// The durable claim must follow the accounting run-ID equivalence: an empty +// runId is a catch-all in both directions, while genuinely distinct runs stay +// distinct stop transitions. +func TestRecordSpecialistStopHookClaimRunIDEquivalence(t *testing.T) { + for _, tc := range []struct { + name string + first string + second string + want int32 + }{ + {name: "empty run id is a catch-all for a later run", first: "", second: "run_1", want: 1}, + {name: "a recorded run covers a later empty query", first: "run_1", second: "", want: 1}, + {name: "distinct runs are distinct stops", first: "run_1", second: "run_2", want: 2}, + } { + t.Run(tc.name, func(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) + } + var stops atomic.Int32 + executor := Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&stops)} + input := specialistAccountingInput{ + ParentSessionID: parent.SessionID, + ChildSessionID: "child_task", + SpecialistName: "worker", + Mode: "background", + Background: true, + } + executor.recordSpecialistStop(input, StreamResult{RunID: tc.first}, "success", 0, nil, false) + executor.recordSpecialistStop(input, StreamResult{RunID: tc.second}, "success", 0, nil, false) + if got := stops.Load(); got != tc.want { + t.Fatalf("stop hooks = %d, want %d", got, tc.want) + } + }) + } +} + +// The in-process fallback (swarm members have no parent session) applies the +// same run-ID equivalence, so a later observation carrying a runId cannot +// re-fire a stop that was already claimed without one. +func TestRecordSpecialistStopInProcessClaimRunIDEquivalence(t *testing.T) { + var stops atomic.Int32 + executor := Executor{LifecycleHooks: stopCounterHooks(&stops)} + input := specialistAccountingInput{ChildSessionID: "member_child", SpecialistName: "subagent", Mode: "foreground"} + executor.recordSpecialistStop(input, StreamResult{}, "success", 0, nil, false) + executor.recordSpecialistStop(input, StreamResult{RunID: "run_later"}, "success", 0, nil, false) + if got := stops.Load(); got != 1 { + t.Fatalf("stop hooks = %d, want 1", got) + } +} + +// A child that never started must not publish public lifecycle hooks, while the +// internal attempt/stop accounting is retained and durably marks the stop as +// ineligible so a later reaper cannot fire the hook either. +func TestExecutorSkipsHooksWhenChildNeverStarts(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) + } + var hooks []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, _ string, _ map[string]any) { + hooks = append(hooks, event) + }, + }, + Load: func(LoadOptions) (LoadResult, error) { + return LoadResult{Specialists: []Manifest{{ + Metadata: Metadata{Name: "worker", Description: "Does focused work"}, + SystemPrompt: "Work carefully.", + ResolvedTools: []string{"read_file"}, + }}}, nil + }, + RunChild: func(context.Context, string, []string, func(streamjson.Event)) (ChildRunResult, error) { + return ChildRunResult{}, errors.New("exec: no such file or directory") + }, + } + + if _, err := executor.Run(context.Background(), TaskParameters{Name: "worker", Prompt: "inspect auth"}, TaskRunOptions{ParentSessionID: parent.SessionID}); err == nil { + t.Fatal("Run returned nil error for a failed child start") + } + if len(hooks) != 0 { + t.Fatalf("public lifecycle hooks fired for a child that never started: %#v", hooks) + } + + events, err := store.ReadEvents(parent.SessionID) + if err != nil { + t.Fatalf("ReadEvents returned error: %v", err) + } + if starts := eventsOfType(events, sessions.EventSpecialistStart); len(starts) != 1 { + t.Fatalf("attempt start events = %d, want 1", len(starts)) + } + stopEvents := eventsOfType(events, sessions.EventSpecialistStop) + if len(stopEvents) != 1 { + t.Fatalf("stop events = %#v", events) + } + requirePayloadBool(t, eventPayload(t, stopEvents[0]), "childStarted", false) + + claims := eventsOfType(events, specialistStopHookEvent) + if len(claims) != 1 { + t.Fatalf("stop-hook claim events = %d, want 1", len(claims)) + } + requirePayloadBool(t, eventPayload(t, claims[0]), "dispatched", false) +} + +// A never-started background child leaves an errored task behind; a later +// TaskOutput reaper must not fire the stop hook from the recorded eligibility. +func TestBackgroundLaunchFailureDoesNotFireStopHookOnReap(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) + } + var hooks []string + executor := Executor{ + BinaryPath: "/usr/local/bin/zero", + SessionStore: store, + BackgroundManager: manager, + NewSessionID: func() (string, error) { return "child_task", nil }, + LifecycleHooks: &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + hooks = append(hooks, event) + }, + }, + Load: func(LoadOptions) (LoadResult, error) { + return LoadResult{Specialists: []Manifest{{ + Metadata: Metadata{Name: "worker", Description: "Does focused work"}, + SystemPrompt: "Work carefully.", + ResolvedTools: []string{"read_file"}, + }}}, nil + }, + LaunchBackground: func(string, []string, string, func(int)) (int, error) { + return 0, errors.New("launch failed") + }, + } + + _, err = executor.Run(context.Background(), TaskParameters{Name: "worker", Prompt: "inspect auth", RunInBackground: true}, TaskRunOptions{ParentSessionID: parent.SessionID}) + if err == nil { + t.Fatal("Run returned nil error for a failed background launch") + } + if len(hooks) != 0 { + t.Fatalf("public lifecycle hooks fired for a child that never started: %#v", hooks) + } + + // A separate process reaping the errored task must stay silent too. + var reapHooks atomic.Int32 + tool := NewOutputTool(manager) + tool.SessionStore = store + tool.LifecycleHooks = stopCounterHooks(&reapHooks) + 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 := reapHooks.Load(); got != 0 { + t.Fatalf("stop hooks from reaping a never-started task = %d, want 0", got) + } + + events, err := store.ReadEvents(parent.SessionID) + if err != nil { + t.Fatalf("ReadEvents returned error: %v", err) + } + if stopEvents := eventsOfType(events, sessions.EventSpecialistStop); len(stopEvents) != 1 { + t.Fatalf("stop events = %d, want 1", len(stopEvents)) + } + if claims := eventsOfType(events, specialistStopHookEvent); len(claims) != 1 { + t.Fatalf("stop-hook claim events = %d, want 1", len(claims)) + } +} + +// A background child that exits before the launcher returns must not run its +// terminal accounting - and therefore must not fire the specialistStop hook - +// until the start decision has completed. Otherwise a fast-exiting child +// publishes specialistStop before specialistStart, and the recorded stop can +// miss the pid that SetPID has not stored yet. The synchronous launcher below +// also pins the deadlock requirement: the exit callback may not wait inline, +// because onExit runs before the launcher (and the start decision) can return. +func TestRunBackgroundDefersExitUntilStartDecisionCompletes(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) + } + var hookMu sync.Mutex + var hookEvents []string + stopDispatched := make(chan struct{}, 1) + executor := Executor{ + BinaryPath: "/usr/local/bin/zero", + SessionStore: store, + BackgroundManager: manager, + NewSessionID: func() (string, error) { return "child_task", nil }, + LifecycleHooks: &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + hookMu.Lock() + hookEvents = append(hookEvents, event) + hookMu.Unlock() + if event == "specialistStop" { + select { + case stopDispatched <- struct{}{}: + default: + } + } + }, + }, + Load: func(LoadOptions) (LoadResult, error) { + return LoadResult{Specialists: []Manifest{{ + Metadata: Metadata{Name: "worker", Description: "Does focused work"}, + SystemPrompt: "Work carefully.", + ResolvedTools: []string{"read_file"}, + }}}, nil + }, + // The child exits before the launcher returns, so the exit callback runs + // synchronously - exactly the ordering hazard this test guards. + LaunchBackground: func(_ string, _ []string, _ string, onExit func(exitCode int)) (int, error) { + onExit(0) + return 1234, nil + }, + } + + runErr := make(chan error, 1) + go func() { + _, err := executor.Run(context.Background(), TaskParameters{ + Name: "worker", + Prompt: "inspect auth", + RunInBackground: true, + }, TaskRunOptions{ParentSessionID: parent.SessionID}) + runErr <- err + }() + select { + case err := <-runErr: + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + case <-time.After(10 * time.Second): + t.Fatal("synchronous launcher deadlocked the background start decision") + } + + // Run publishes specialistStart before it returns; the deferred exit handling + // publishes specialistStop afterwards. + select { + case <-stopDispatched: + case <-time.After(10 * time.Second): + t.Fatal("deferred exit accounting never dispatched specialistStop") + } + hookMu.Lock() + gotHooks := strings.Join(hookEvents, ",") + hookMu.Unlock() + if gotHooks != "specialistStart,specialistStop" { + t.Fatalf("lifecycle hooks = %q, want start before stop", gotHooks) + } + + events, err := store.ReadEvents(parent.SessionID) + if err != nil { + t.Fatalf("ReadEvents returned error: %v", err) + } + stopEvents := eventsOfType(events, sessions.EventSpecialistStop) + if len(stopEvents) != 1 { + t.Fatalf("stop events = %#v", events) + } + // The deferred stop must observe the task after SetPID; a stop recorded + // before the start decision completes would omit the pid. + requirePayloadInt(t, eventPayload(t, stopEvents[0]), "pid", 1234) +} + +// A recovered task stopped explicitly (launcher process gone) must record the +// terminal transition and fire the stop hook exactly once through the shared +// durable claim, even when a later TaskOutput poll reaps it again. +func TestStopToolRecordsRecoveredTaskStopOnce(t *testing.T) { + killed := []int{} + manager, err := background.NewManagerWithOptions(background.ManagerOptions{ + RootDir: t.TempDir(), + KillProcess: func(pid int) error { + killed = append(killed, pid) + return nil + }, + }) + if err != nil { + t.Fatalf("NewManagerWithOptions returned error: %v", err) + } + 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) + } + outputFile, err := manager.Register(background.RegisterInput{ + TaskID: "child_task", + Type: specialistTaskType, + SpecialistName: "worker", + Description: "Auth check", + ParentID: parent.SessionID, + PID: 4321, + }) + 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) + } + + var counter atomic.Int32 + tool := &StopTool{manager: manager, SessionStore: store, LifecycleHooks: stopCounterHooks(&counter)} + result := tool.Run(context.Background(), map[string]any{"task_id": "child_task"}) + if result.Status != tools.StatusOK || result.Meta["status"] != string(background.StatusKilled) { + t.Fatalf("TaskStop result = %#v", result) + } + if !reflect.DeepEqual(killed, []int{4321}) { + t.Fatalf("killed pids = %#v", killed) + } + if got := counter.Load(); got != 1 { + t.Fatalf("stop hooks = %d, want 1 from TaskStop", got) + } + + // A later TaskOutput reaper in another process (fresh bridge) must not + // re-dispatch; the durable claim keys on the runId read from the output. + var reapCounter atomic.Int32 + outputTool := NewOutputTool(manager) + outputTool.SessionStore = store + outputTool.LifecycleHooks = stopCounterHooks(&reapCounter) + if result := outputTool.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 := reapCounter.Load(); got != 0 { + t.Fatalf("stop hooks after later reap = %d, want 0", got) + } + + events, err := store.ReadEvents(parent.SessionID) + if err != nil { + t.Fatalf("ReadEvents returned error: %v", err) + } + if stopEvents := eventsOfType(events, sessions.EventSpecialistStop); len(stopEvents) != 1 { + t.Fatalf("stop events = %d, want 1", len(stopEvents)) + } + claims := eventsOfType(events, specialistStopHookEvent) + if len(claims) != 1 { + t.Fatalf("stop-hook claim events = %d, want 1", len(claims)) + } + requirePayloadString(t, eventPayload(t, claims[0]), "runId", "run_1") +} + +func stopCounterHooks(counter *atomic.Int32) *LifecycleHooks { + return &LifecycleHooks{ + Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) { + if event == "specialistStop" { + counter.Add(1) + } + }, + } +} + +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 7bfb6d8eb..47412ac8c 100644 --- a/internal/specialist/exec.go +++ b/internal/specialist/exec.go @@ -15,6 +15,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "github.com/Twigpine/zero/internal/background" "github.com/Twigpine/zero/internal/sessions" @@ -24,6 +25,7 @@ import ( const ( sessionTagSpecialist = "specialist" + specialistTaskType = "specialist" promptFileThresholdBytes = 4 * 1024 maxSpecialistDepth = 8 // hard cap to prevent infinite recursion/resource exhaustion via Task ) @@ -50,6 +52,64 @@ 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) + // stopMu guards stopClaims, the in-process fallback used when no durable + // claim can be written (swarm members have no parent session) or the session + // store is unavailable. + stopMu sync.Mutex + stopClaims map[string]map[string]struct{} +} + +// claimInProcessStopDispatch reports whether this bridge should run +// specialistStop for the given child/run. It is the process-local fallback for +// callers without a durable ledger; the run-ID equivalence mirrors +// specialistPayloadMatches: an empty runId is a catch-all for both the query and +// any recorded claim, so a later observation with a runId cannot re-fire a stop +// that was already claimed without one. An empty child session never dedups +// (there is nothing to key on), matching the accounting matcher. +func (h *LifecycleHooks) claimInProcessStopDispatch(childSessionID, runID string) bool { + if h == nil || h.Dispatch == nil { + return false + } + childSessionID = strings.TrimSpace(childSessionID) + if childSessionID == "" { + return true + } + runID = strings.TrimSpace(runID) + + h.stopMu.Lock() + defer h.stopMu.Unlock() + if h.stopClaims == nil { + h.stopClaims = map[string]map[string]struct{}{} + } + claims := h.stopClaims[childSessionID] + if claims == nil { + claims = map[string]struct{}{} + h.stopClaims[childSessionID] = claims + } + claimed := false + if runID == "" { + claimed = len(claims) > 0 + } else { + _, emptyClaim := claims[""] + _, exactClaim := claims[runID] + claimed = emptyClaim || exactClaim + } + if !claimed { + claims[runID] = struct{}{} + } + return !claimed } type BuildArgsInput struct { @@ -478,7 +538,7 @@ func (executor Executor) runBackground(ctx context.Context, built BuildArgsResul } outputFile, err := manager.Register(background.RegisterInput{ TaskID: built.SessionID, - Type: "specialist", + Type: specialistTaskType, SpecialistName: manifest.Metadata.Name, Description: params.Description, ParentID: options.ParentSessionID, @@ -516,34 +576,48 @@ func (executor Executor) runBackground(ctx context.Context, built BuildArgsResul Background: true, } executor.recordSpecialistStart(accounting) + // The launcher's exit callback runs the terminal accounting, which can claim + // and dispatch the specialistStop hook. A child that exits immediately would + // otherwise publish specialistStop before the start decision below completes + // (and before SetPID records the pid). The callback waits for startSettled + // from its own goroutine: waiting inline would deadlock a synchronous + // injected launcher whose onExit runs before it returns. + startSettled := make(chan struct{}) + defer close(startSettled) pid, err := executor.launchBackground(binaryPath, built.Args, outputFile, func(exitCode int) { - status := background.StatusCompleted - if exitCode != 0 { - status = background.StatusError - } - _ = manager.MarkExited(built.SessionID, status, exitCode) - if task, ok := manager.Get(built.SessionID); ok { - summary := StreamResult{ExitCode: task.ExitCode} - if data, err := os.ReadFile(task.OutputFile); err == nil { - summary, _ = summarizeTaskData(string(data), task.ExitCode) + go func() { + <-startSettled + status := background.StatusCompleted + if exitCode != 0 { + status = background.StatusError } - executor.recordBackgroundTaskAccounting(task, summary) - } - executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile) + _ = manager.MarkExited(built.SessionID, status, exitCode) + if task, ok := manager.Get(built.SessionID); ok { + executor.recordBackgroundTaskAccounting(task, taskStreamSummary(task)) + } + executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile) + }() }) if err != nil { + // The child never started: keep the attempt accounting, but publish no + // public lifecycle hooks for a sub-agent that did not exist. + accounting.LaunchFailed = true _ = manager.UpdateStatus(built.SessionID, background.StatusError, -1) executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile) executor.recordSpecialistStop(accounting, StreamResult{ExitCode: -1}, "error", -1, err, false) return ExecResult{}, err } + // Startup is confirmed (launchBackground only returns without error after the + // child process started), so the specialistStart hook can fire now. + executor.dispatchSpecialistStartHook(accounting) if pid > 0 { if err := manager.SetPID(built.SessionID, pid); err != nil { // The child is running but its PID was never recorded, so nothing can // track or stop it later. Kill it (and its process group) now, mark the // task errored, and record the stop so it cannot become an untracked - // orphan. The dedup in recordSpecialistStop makes the later onExit - // accounting a no-op. + // orphan. The child started, so the stop is a real terminal transition: + // both the accounting dedup and the shared stop-hook claim make the + // later onExit accounting a no-op. if killErr := background.TerminateProcess(pid); killErr != nil { // Don't mask a failed kill: a surviving orphan is worse than the // SetPID error alone, so surface both to the caller. @@ -627,7 +701,23 @@ func (executor Executor) runBuiltArgs(ctx context.Context, built BuildArgsResult Background: false, } executor.recordSpecialistStart(accounting) - run, err := executor.runChild(ctx, binaryPath, built.Args, progress) + startDispatched := false + onStarted := func() { + if startDispatched { + return + } + startDispatched = true + executor.dispatchSpecialistStartHook(accounting) + } + run, err := executor.runChild(ctx, binaryPath, built.Args, progress, onStarted) + started := run.Started || err == nil + if started { + // A seam that returns a completed run without the explicit Started flag + // still proves the child existed; make sure the hook fired exactly once. + onStarted() + } else { + accounting.LaunchFailed = true + } if err != nil { exitCode := run.exitCodeOr(-1) summary := SummarizeStream(run.Events, exitCode) @@ -697,11 +787,20 @@ func (executor Executor) binaryPath() (string, error) { return path, nil } -func (executor Executor) runChild(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event)) (ChildRunResult, error) { +// runChild runs one child process. onStarted, when non-nil, is invoked exactly +// once as soon as the child process has actually started (after os/exec Start +// succeeds), so public lifecycle hooks cannot fire for a launch that failed. +// The injected RunChild seam cannot stream that signal, so a seam result marked +// Started reports it when the seam returns. +func (executor Executor) runChild(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event), onStarted func()) (ChildRunResult, error) { if executor.RunChild != nil { - return executor.RunChild(ctx, binaryPath, append([]string(nil), args...), progress) + result, err := executor.RunChild(ctx, binaryPath, append([]string(nil), args...), progress) + if result.Started && onStarted != nil { + onStarted() + } + return result, err } - return runChildProcess(ctx, binaryPath, args, progress) + return runChildProcess(ctx, binaryPath, args, progress, onStarted) } func (executor Executor) launchBackground(binaryPath string, args []string, outputFile string, onExit func(exitCode int)) (int, error) { @@ -830,7 +929,7 @@ func cleanupPromptFile(promptFile string) { _ = os.Remove(promptFile) } -func runChildProcess(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event)) (ChildRunResult, error) { +func runChildProcess(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event), onStarted func()) (ChildRunResult, error) { var stderr bytes.Buffer command := osexec.CommandContext(ctx, binaryPath, args...) // Put the child in its own process group and group-kill on cancel/timeout, so a @@ -844,6 +943,9 @@ func runChildProcess(ctx context.Context, binaryPath string, args []string, prog if err := command.Start(); err != nil { return ChildRunResult{Stderr: stderr.String()}, fmt.Errorf("start specialist child: %w", err) } + if onStarted != nil { + onStarted() + } events := []streamjson.Event{} reader := bufio.NewReader(stdout) for { diff --git a/internal/specialist/output_tool.go b/internal/specialist/output_tool.go index ee5cdd14e..9bc104024 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 7f0f62449..96f69a741 100644 --- a/internal/specialist/registry.go +++ b/internal/specialist/registry.go @@ -14,8 +14,8 @@ 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(newStopToolWithManagerFunc(runtime.Manager)) + registry.Register(newOutputToolWithManagerFunc(runtime.Manager, executor.SessionStore, executor.LifecycleHooks)) + registry.Register(newStopToolWithManagerFunc(runtime.Manager, executor.SessionStore, executor.LifecycleHooks)) registry.Register(NewGenerateTool(NewStorage(executor.Paths))) return runtime, nil } diff --git a/internal/specialist/stop_tool.go b/internal/specialist/stop_tool.go index 9da2a992f..71269cc98 100644 --- a/internal/specialist/stop_tool.go +++ b/internal/specialist/stop_tool.go @@ -6,20 +6,26 @@ import ( "strings" "github.com/Twigpine/zero/internal/background" + "github.com/Twigpine/zero/internal/sessions" "github.com/Twigpine/zero/internal/tools" ) type StopTool struct { manager *background.Manager managerFunc BackgroundManagerFunc + // SessionStore and LifecycleHooks participate in the shared, deduped stop + // transition so an explicitly stopped recovered task produces exactly one + // specialistStop hook, the same as a launcher-owned exit. + SessionStore *sessions.Store + LifecycleHooks *LifecycleHooks } type stopParameters struct { TaskID string } -func newStopToolWithManagerFunc(managerFunc BackgroundManagerFunc) *StopTool { - return &StopTool{managerFunc: managerFunc} +func newStopToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store, lifecycleHooks *LifecycleHooks) *StopTool { + return &StopTool{managerFunc: managerFunc, SessionStore: sessionStore, LifecycleHooks: lifecycleHooks} } func (tool *StopTool) Name() string { @@ -65,6 +71,18 @@ func (tool *StopTool) Run(ctx context.Context, args map[string]any) tools.Result if err := manager.Kill(params.TaskID); err != nil { return taskError(err) } + // A recovered task (launcher process gone) may only ever be reaped here: the + // in-memory onExit callback does not exist in this process, and a later + // TaskOutput poll is optional. Record the terminal transition now so the + // specialistStop hook fires exactly once through the shared durable claim. + // recordBackgroundTaskAccounting is a no-op for non-specialist tasks and for + // tasks that are still running (Kill raced a status change). + if task, ok := manager.Get(params.TaskID); ok && task.Type == specialistTaskType { + Executor{ + SessionStore: tool.SessionStore, + LifecycleHooks: tool.LifecycleHooks, + }.recordBackgroundTaskAccounting(task, taskStreamSummary(task)) + } return tools.Result{ Status: tools.StatusOK, Output: fmt.Sprintf("task_id: %s\nstatus: killed", params.TaskID),