diff --git a/docs/community/contributing/index.md b/docs/community/contributing/index.md index fc923819e6..0f9ff6ac0b 100644 --- a/docs/community/contributing/index.md +++ b/docs/community/contributing/index.md @@ -474,3 +474,32 @@ implementer, even to verify the skip path. Validation is limited to pure gate units, source placement review, compile-only checks, static analysis and builds. Future fixture additions must be guarded before server, proxy, transport, SDK flow or browser-session setup, with pure cases and non-Gemini rows kept separate. + +## Toolset worker lifetimes + +Toolsets with a background loop and no startup handshake can implement +`tools.Runnable.Run(ctx)` instead of `tools.Startable`. `Run` blocks until +cancellation and joins its work before returning; timers and other run resources +stay local. Callers must not run the same toolset concurrently. + +Agent, additional-toolset, Code Mode, and fork-skill paths use the canonical +`tools.NewStartable` wrapper. It owns one run, preserves startup context values +without inheriting startup cancellation, and cancels and joins on shutdown—even +with an already-canceled shutdown context. `Startable` takes precedence over +`Runnable` so existing lifecycle wrappers retain ownership. Fork-skill toolsets +are started lazily and owned until team shutdown, not one skill invocation. + +For direct scheduler users, replace `Start`/`Stop` with a caller-owned `Run(ctx)` +and cancellation/join, or wrap it once with `tools.NewStartable`. Two independent +wrappers around the same raw worker are not a shared lifecycle. The wrapper is +an intentional bridge to legacy host `Start`/`Stop`; it does not make those host +lifetimes structured. + +Raw runnable session extras supplied programmatically are instead adopted once +by their runtime and joined by `Close`; existing lifecycle wrappers remain +borrowed. Keep a session's extras fixed during execution. Share one explicit +wrapper when sharing a worker across runtimes or using a noncomparable runnable +value across sessions. Runtime adoption deduplicates comparable worker providers +across decorated references; noncomparable workers are scoped to a session slot. +Runnable-only workers do not support the separate `Restartable` capability; +stop and start their canonical wrapper to replace a run. diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 06001c360a..68842773e7 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -28,6 +28,7 @@ type Agent struct { welcomeMessage string instruction string toolsets []*tools.StartableToolSet + ownedToolsets []*tools.StartableToolSet models []provider.Provider fallbackModels []provider.Provider // Fallback models to try if primary fails fallbackRetries int // Number of retries per fallback model with exponential backoff @@ -699,7 +700,7 @@ func (a *Agent) DrainWarnings() []string { func (a *Agent) StopToolSets(ctx context.Context) error { var errs []error - for _, toolSet := range a.toolsets { + for _, toolSet := range slices.Concat(a.toolsets, a.ownedToolsets) { // StopIfStarted checks-and-stops atomically under the toolset's // lifecycle lock, so a toolset that was never started is left // untouched, an in-flight Start that settles in time is stopped diff --git a/pkg/agent/opts.go b/pkg/agent/opts.go index 528aebc78b..adafe7b284 100644 --- a/pkg/agent/opts.go +++ b/pkg/agent/opts.go @@ -34,6 +34,14 @@ func WithToolSets(toolSet ...tools.ToolSet) Opt { } } +// WithOwnedToolSets registers lazily started toolsets for agent shutdown without +// exposing or starting them with the agent's ordinary tools. +func WithOwnedToolSets(toolsets ...*tools.StartableToolSet) Opt { + return func(a *Agent) { + a.ownedToolsets = append(a.ownedToolsets, toolsets...) + } +} + // WithAdditionalToolSets appends toolsets without replacing existing lifecycle wrappers. // Like other options, apply it before the agent is used by a runtime. func WithAdditionalToolSets(toolsets ...tools.ToolSet) Opt { diff --git a/pkg/agent/runnable_test.go b/pkg/agent/runnable_test.go new file mode 100644 index 0000000000..2c0607dc56 --- /dev/null +++ b/pkg/agent/runnable_test.go @@ -0,0 +1,80 @@ +package agent + +import ( + "context" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/scheduler" + "github.com/docker/docker-agent/pkg/tools/codemode" +) + +type schedulerRecallRuntime struct { + tools.NopRuntime + + recalls chan string +} + +func (r *schedulerRecallRuntime) Supports(c tools.Capability) bool { + return c == tools.CapabilityRecall +} + +func (r *schedulerRecallRuntime) Recall(ctx context.Context, prompt string) error { + select { + case r.recalls <- prompt: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func TestAgentOwnsSchedulerRun(t *testing.T) { + t.Parallel() + for _, mode := range []string{"ordinary", "additional", "code_mode"} { + t.Run(mode, func(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + raw := scheduler.New() + ts := tools.WithName(raw, "scheduler") + a := New("root", "") + switch mode { + case "ordinary": + WithToolSets(ts)(a) + case "additional": + WithAdditionalToolSets(ts)(a) + case "code_mode": + WithToolSets(codemode.Wrap(ts))(a) + } + t.Cleanup(func() { require.NoError(t, a.StopToolSets(t.Context())) }) + ctx, cancel := context.WithCancel(t.Context()) + for _, outcome := range a.StartToolSets(ctx, time.Second) { + require.Equal(t, tools.StartReady, (<-outcome).Kind) + } + cancel() + all, err := raw.Tools(t.Context()) + require.NoError(t, err) + rt := &schedulerRecallRuntime{recalls: make(chan string, 2)} + result, err := all[0].Handler(t.Context(), tools.ToolCall{Function: tools.FunctionCall{ + Name: scheduler.ToolNameCreateSchedule, Arguments: `{"prompt":"check build","when":"every:1m"}`, + }}, rt) + require.NoError(t, err) + require.False(t, result.IsError) + synctest.Sleep(time.Minute) + select { + case prompt := <-rt.recalls: + assert.Contains(t, prompt, "check build") + default: + t.Fatal("scheduler did not survive startup cancellation") + } + require.NoError(t, a.StopToolSets(ctx)) + synctest.Sleep(time.Minute) + assert.Empty(t, rt.recalls) + }) + }) + } +} diff --git a/pkg/runtime/extra_toolsets.go b/pkg/runtime/extra_toolsets.go new file mode 100644 index 0000000000..cde131fce5 --- /dev/null +++ b/pkg/runtime/extra_toolsets.go @@ -0,0 +1,82 @@ +package runtime + +import ( + "context" + "errors" + "reflect" + "sync" + + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/tools" +) + +// extraToolsetRuns owns raw runnable extras; prewrapped extras remain borrowed +// from their team or embedder. Only lifecycle is cached, never filtered tools. +type extraToolsetRuns struct { + mu sync.Mutex + entries []extraToolsetRun + closed bool + done chan struct{} + err error +} + +type extraToolsetRun struct { + provider tools.Runnable + sess *session.Session + slot int + managed *tools.StartableToolSet +} + +func (e *extraToolsetRuns) start(ctx context.Context, sess *session.Session, slot int, ts tools.ToolSet) error { + if startable, ok := tools.As[tools.Startable](ts); ok { + return startable.Start(ctx) + } + provider, ok := tools.As[tools.Runnable](ts) + if !ok { + return nil + } + e.mu.Lock() + defer e.mu.Unlock() + if e.closed { + return errors.New("runtime is closed") + } + keyComparable := reflect.ValueOf(provider).Comparable() + for _, entry := range e.entries { + if keyComparable && reflect.ValueOf(entry.provider).Comparable() { + if entry.provider == provider { + return entry.managed.Start(ctx) + } + } else if !keyComparable && entry.sess == sess && entry.slot == slot { + return entry.managed.Start(ctx) + } + } + managed := tools.NewStartable(ts) + e.entries = append(e.entries, extraToolsetRun{provider: provider, sess: sess, slot: slot, managed: managed}) + // Runnable has no startup handshake, so Start never waits for Run here. + return managed.Start(ctx) +} + +func (e *extraToolsetRuns) close() error { + e.mu.Lock() + if e.closed { + done := e.done + e.mu.Unlock() + <-done + return e.err + } + e.closed = true + e.done = make(chan struct{}) + entries := e.entries + e.entries = nil + e.mu.Unlock() + + var errs []error + for _, entry := range entries { + if err := entry.managed.Stop(context.Background()); err != nil { + errs = append(errs, err) + } + } + e.err = errors.Join(errs...) + close(e.done) + return e.err +} diff --git a/pkg/runtime/extra_toolsets_test.go b/pkg/runtime/extra_toolsets_test.go new file mode 100644 index 0000000000..db42b30627 --- /dev/null +++ b/pkg/runtime/extra_toolsets_test.go @@ -0,0 +1,148 @@ +package runtime + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/scheduler" +) + +type extraRecallRuntime struct { + tools.NopRuntime + + messages chan string +} + +func (*extraRecallRuntime) Supports(c tools.Capability) bool { return c == tools.CapabilityRecall } +func (r *extraRecallRuntime) Recall(ctx context.Context, msg string) error { + select { + case r.messages <- msg: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func TestExtraSchedulerLifecycle(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + root := agent.New("root", "test", agent.WithModel(&mockProvider{id: "test/mock-model", stream: &mockStream{}})) + r, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithModelStore(mockModelStore{})) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, r.Close()) }) + raw := scheduler.New() + first := session.New(session.WithExtraToolSets([]tools.ToolSet{raw})) + second := session.New(session.WithExtraToolSets([]tools.ToolSet{tools.WithName(raw, "schedule")})) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + var listed []tools.Tool + for range 3 { + listed = r.skillSubSessionTools(ctx, first, root, nil) + require.Len(t, listed, 3) + require.Len(t, r.skillSubSessionTools(ctx, second, root, nil), 3) + } + assert.Len(t, r.extraToolsetRuns.entries, 1, "decorated references share one run") + host := &extraRecallRuntime{messages: make(chan string, 2)} + result, err := listed[0].Handler(ctx, tools.ToolCall{Function: tools.FunctionCall{ + Name: scheduler.ToolNameCreateSchedule, Arguments: `{"prompt":"check","when":"every:1m"}`, + }}, host) + require.NoError(t, err) + require.False(t, result.IsError) + cancel() + synctest.Sleep(time.Minute) + assert.Len(t, host.messages, 1) + <-host.messages + require.NoError(t, r.Close()) + synctest.Sleep(time.Minute) + assert.Empty(t, host.messages) + assert.Empty(t, r.skillSubSessionTools(t.Context(), first, root, nil), "closed runtime cannot restart raw workers") + }) +} + +type extraWorker struct { + run func(context.Context) +} + +func (*extraWorker) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (w *extraWorker) Run(ctx context.Context) { w.run(ctx) } + +func TestExtraRunnableBorrowedOwnership(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + started := make(chan context.Context, 1) + inner := &extraWorker{run: func(ctx context.Context) { started <- ctx; <-ctx.Done() }} + owner := tools.NewStartable(inner) + t.Cleanup(func() { require.NoError(t, owner.Stop(t.Context())) }) + var runs extraToolsetRuns + sess := session.New() + require.NoError(t, runs.start(t.Context(), sess, 0, tools.WithName(owner, "worker"))) + runCtx := <-started + require.NoError(t, runs.close()) + require.NoError(t, runCtx.Err(), "runtime must not stop team/embedder-owned workers") + require.NoError(t, owner.Stop(t.Context())) + require.ErrorIs(t, runCtx.Err(), context.Canceled) + }) +} + +func TestExtraRunnableCloseRacesStarts(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + var active, launches atomic.Int32 + inner := &extraWorker{run: func(ctx context.Context) { + active.Add(1) + launches.Add(1) + defer active.Add(-1) + <-ctx.Done() + }} + var runs extraToolsetRuns + sess := session.New() + require.NoError(t, runs.start(t.Context(), sess, 0, inner)) + synctest.Wait() + var callers sync.WaitGroup + for range 20 { + callers.Go(func() { _ = runs.start(t.Context(), sess, 0, inner) }) + callers.Go(func() { assert.NoError(t, runs.close()) }) + } + callers.Wait() + assert.Zero(t, active.Load()) + assert.EqualValues(t, 1, launches.Load()) + require.ErrorContains(t, runs.start(t.Context(), sess, 0, inner), "closed") + }) +} + +type noncomparableExtra struct { + start []func(context.Context) +} + +func (w noncomparableExtra) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (w noncomparableExtra) Run(ctx context.Context) { w.start[0](ctx) } + +func TestExtraRunnableNoncomparable(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + var launches atomic.Int32 + inner := noncomparableExtra{start: []func(context.Context){func(ctx context.Context) { + launches.Add(1) + <-ctx.Done() + }}} + var runs extraToolsetRuns + sess := session.New() + for range 3 { + require.NoError(t, runs.start(t.Context(), sess, 0, inner)) + } + synctest.Wait() + assert.EqualValues(t, 1, launches.Load()) + require.NoError(t, runs.close()) + }) +} diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 4a30e44e61..f493608e84 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -2014,7 +2014,7 @@ func (r *LocalRuntime) skillSubSessionTools(ctx context.Context, sess *session.S agentTools = filterAllowedTools(agentTools, sess.AllowedTools) - for _, ts := range sess.ExtraToolSets { + for i, ts := range sess.ExtraToolSets { tools.ConfigureHandlers(ts, tools.ScopedElicitationHandler, tools.SamplingScopeHandler, @@ -2023,12 +2023,10 @@ func (r *LocalRuntime) skillSubSessionTools(ctx context.Context, sess *session.S r.managedOAuth, r.unmanagedOAuthRedirectURI, ) - if startable, ok := tools.As[tools.Startable](ts); ok { - if err := startable.Start(ctx); err != nil { - slog.WarnContext(ctx, "Skill toolset failed to start; skipping", - "agent", a.Name(), "toolset", tools.DescribeToolSet(ts), "error", err) - continue - } + if err := r.extraToolsetRuns.start(ctx, sess, i, ts); err != nil { + slog.WarnContext(ctx, "Skill toolset failed to start; skipping", + "agent", a.Name(), "toolset", tools.DescribeToolSet(ts), "error", err) + continue } extra, err := ts.Tools(ctx) if err != nil { diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 7226046dba..f9c0c60966 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -385,6 +385,8 @@ type LocalRuntime struct { bgAgents *agenttool.Handler + extraToolsetRuns extraToolsetRuns + // dmrModelLister lists the models pulled locally in Docker Model Runner, // used to populate DMR entries in the model picker. Defaults to // dmrmodels.ListModelsWithMetadata in NewLocalRuntime; left nil by runtimes built directly @@ -1580,7 +1582,7 @@ func (r *LocalRuntime) SessionStore() session.Store { func (r *LocalRuntime) Close() error { r.bgAgents.StopAll() r.OnToolsChanged(nil) - return nil + return r.extraToolsetRuns.close() } // UpdateSessionTitle persists the session title via the session store. diff --git a/pkg/session/session.go b/pkg/session/session.go index 4bbd5a68a2..2cb3db3f34 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -1794,7 +1794,10 @@ func WithAllowedTools(names []string) Opt { // WithExtraToolSets injects additional toolsets into the session on top of // the agent's own toolsets. Used by fork-mode skill sub-sessions that declare -// assistive toolsets. +// assistive toolsets. Raw Runnable workers are owned until runtime Close; +// prewrapped toolsets remain owned by their team or embedder. Set extras before +// execution. Shared workers require one owner: use an explicit NewStartable +// wrapper to share across runtimes or for noncomparable runnable values. func WithExtraToolSets(toolSets []tools.ToolSet) Opt { return func(s *Session) { s.ExtraToolSets = toolSets diff --git a/pkg/teamloader/skills_toolsets_test.go b/pkg/teamloader/skills_toolsets_test.go index 446849939d..b9791cad53 100644 --- a/pkg/teamloader/skills_toolsets_test.go +++ b/pkg/teamloader/skills_toolsets_test.go @@ -1,12 +1,15 @@ package teamloader import ( + "context" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/tools" skillstool "github.com/docker/docker-agent/pkg/tools/builtin/skills" ) @@ -161,3 +164,78 @@ agents: assert.NotEqual(t, "run_shell_background", tool.Name, "readonly agent must not expose mutating background-shell tools via a fork skill") } } + +func TestForkSkillToolsetsOwnedUntilTeamShutdown(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "dummy") + data := []byte(`toolsets: + worker: + type: worker +agents: + root: + model: openai/gpt-4o + instruction: test + skills: + - name: builder + description: Build things. + context: fork + toolsets: [worker] + instructions: Build. +`) + started, finished := make(chan context.Context, 1), make(chan struct{}) + inner := &forkWorker{started: started, finished: finished} + registry := NewToolsetRegistry(map[string]ToolsetCreator{ + "worker": func(context.Context, latest.Toolset, string, *config.RuntimeConfig, string) (tools.ToolSet, error) { + return inner, nil + }, + }) + tm, err := Load(t.Context(), config.NewBytesSource("fork_worker.yaml", data), &config.RuntimeConfig{}, + withTestProviderRegistry(WithToolsetRegistry(registry))...) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, tm.StopToolSets(t.Context())) }) + root, err := tm.Agent("root") + require.NoError(t, err) + for _, outcome := range root.StartToolSets(t.Context(), time.Second) { + require.Equal(t, tools.StartReady, (<-outcome).Kind) + } + assert.Empty(t, started, "ordinary startup must not start assistive toolsets") + var skillSet *skillstool.ToolSet + for _, ts := range root.ToolSets() { + if s, ok := skillSetFrom(ts); ok { + skillSet = s + } + } + require.NotNil(t, skillSet) + prepare := func() *skillstool.PreparedSkillFork { + prepared, failure, err := skillSet.PrepareForkSubSession(t.Context(), skillstool.RunSkillArgs{Name: "builder"}, tools.NopRuntime{}) + require.NoError(t, err) + require.Nil(t, failure) + return prepared + } + first, second := prepare(), prepare() + require.Len(t, first.ToolSets, 1) + require.Same(t, first.ToolSets[0], second.ToolSets[0], "invocations share the team-owned lifecycle") + wrapper, ok := tools.As[*tools.StartableToolSet](first.ToolSets[0]) + require.True(t, ok) + ctx, cancel := context.WithCancel(t.Context()) + require.NoError(t, wrapper.Start(ctx)) + runCtx := <-started + cancel() + require.NoError(t, wrapper.Start(t.Context())) + require.NoError(t, runCtx.Err(), "ending one invocation must not stop the shared worker") + require.NoError(t, tm.StopToolSets(ctx)) + <-finished + require.ErrorIs(t, runCtx.Err(), context.Canceled) + assert.False(t, wrapper.IsStarted()) +} + +type forkWorker struct { + started chan<- context.Context + finished chan<- struct{} +} + +func (*forkWorker) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (w *forkWorker) Run(ctx context.Context) { + defer close(w.finished) + w.started <- ctx + <-ctx.Done() +} diff --git a/pkg/teamloader/teamloader.go b/pkg/teamloader/teamloader.go index 80f03b7936..809d7477b4 100644 --- a/pkg/teamloader/teamloader.go +++ b/pkg/teamloader/teamloader.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "log/slog" + "maps" "path/filepath" "slices" "strings" @@ -556,6 +557,12 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c } if len(forkToolSets) > 0 { skillSet.SetForkToolSets(forkToolSets) + for _, name := range slices.Sorted(maps.Keys(forkToolSets)) { + for _, ts := range forkToolSets[name] { + owned, _ := tools.As[*tools.StartableToolSet](ts) + opts = append(opts, agent.WithOwnedToolSets(owned)) + } + } } if len(forkWarnings) > 0 { opts = append(opts, agent.WithLoadTimeWarnings(forkWarnings)) diff --git a/pkg/tools/builtin/scheduler/loop_test.go b/pkg/tools/builtin/scheduler/loop_test.go index fbd1070cdf..007a5c5741 100644 --- a/pkg/tools/builtin/scheduler/loop_test.go +++ b/pkg/tools/builtin/scheduler/loop_test.go @@ -1,6 +1,7 @@ package scheduler import ( + "context" "errors" "testing" "testing/synctest" @@ -22,8 +23,7 @@ func TestLoopFiresAtDeadline(t *testing.T) { res, err := ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "check build", When: when}, rt) require.NoError(t, err) require.False(t, res.IsError, res.Output) - require.NoError(t, ts.Start(t.Context())) - t.Cleanup(func() { require.NoError(t, ts.Stop(t.Context())) }) + runScheduler(t, ts) synctest.Sleep(time.Minute - time.Nanosecond) require.Empty(t, rt.messages()) @@ -55,8 +55,7 @@ func TestLoopWakesForEarlierSchedule(t *testing.T) { res, err := ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "later", When: "in:1m"}, rt) require.NoError(t, err) require.False(t, res.IsError, res.Output) - require.NoError(t, ts.Start(t.Context())) - t.Cleanup(func() { require.NoError(t, ts.Stop(t.Context())) }) + runScheduler(t, ts) synctest.Wait() res, err = ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "earlier", When: "in:10s"}, rt) @@ -91,8 +90,7 @@ func TestLoopRetriesFailedRecall(t *testing.T) { res, err := ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "retry", When: "in:10s"}, rt) require.NoError(t, err) require.False(t, res.IsError, res.Output) - require.NoError(t, ts.Start(t.Context())) - t.Cleanup(func() { require.NoError(t, ts.Stop(t.Context())) }) + runScheduler(t, ts) synctest.Sleep(10 * time.Second) require.Len(t, rt.messages(), 1) @@ -117,12 +115,44 @@ func TestLoopStopPreventsRecall(t *testing.T) { res, err := ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "stopped", When: "in:10s"}, rt) require.NoError(t, err) require.False(t, res.IsError, res.Output) - require.NoError(t, ts.Start(t.Context())) - t.Cleanup(func() { require.NoError(t, ts.Stop(t.Context())) }) + stop := runScheduler(t, ts) synctest.Wait() - require.NoError(t, ts.Stop(t.Context())) + stop() synctest.Sleep(time.Minute) require.Empty(t, rt.messages()) }) } + +func runScheduler(t *testing.T, ts *ToolSet) func() { + t.Helper() + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + go func() { + defer close(done) + ts.Run(ctx) + }() + stop := func() { + cancel() + <-done + } + t.Cleanup(stop) + return stop +} + +func TestRunCanceledBeforeStart(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + ts := New() + rt := &fakeRuntime{recall: true} + res, err := ts.createSchedule(t.Context(), CreateScheduleArgs{Prompt: "stopped", When: "in:1s"}, rt) + require.NoError(t, err) + require.False(t, res.IsError) + synctest.Sleep(time.Second) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + ts.Run(ctx) + require.Empty(t, rt.messages()) + require.Len(t, ts.store.list(), 1) + }) +} diff --git a/pkg/tools/builtin/scheduler/scheduler.go b/pkg/tools/builtin/scheduler/scheduler.go index ce986cab8a..1221208d92 100644 --- a/pkg/tools/builtin/scheduler/scheduler.go +++ b/pkg/tools/builtin/scheduler/scheduler.go @@ -30,15 +30,13 @@ type ToolSet struct { now func() time.Time wake chan struct{} - mu sync.Mutex - rt tools.Runtime - cancel context.CancelFunc - wg sync.WaitGroup + mu sync.Mutex + rt tools.Runtime } var ( _ tools.ToolSet = (*ToolSet)(nil) - _ tools.Startable = (*ToolSet)(nil) + _ tools.Runnable = (*ToolSet)(nil) _ tools.Instructable = (*ToolSet)(nil) ) @@ -107,29 +105,10 @@ func (t *ToolSet) fireDue(ctx context.Context, now time.Time) { } } -func (t *ToolSet) Start(context.Context) error { - ctx, cancel := context.WithCancel(context.Background()) - t.mu.Lock() - t.cancel = cancel - t.mu.Unlock() - - t.wg.Go(func() { t.loop(ctx) }) - return nil -} - -func (t *ToolSet) Stop(context.Context) error { - t.mu.Lock() - cancel := t.cancel - t.mu.Unlock() - if cancel != nil { - cancel() - } - t.wg.Wait() - return nil -} - -func (t *ToolSet) loop(ctx context.Context) { - for { +// Run delivers due schedules until ctx is canceled. The caller owns the run +// and must wait for it to return before starting another. +func (t *ToolSet) Run(ctx context.Context) { + for ctx.Err() == nil { wait := loopMaxWait if d, ok := t.store.untilNext(t.now()); ok && d < wait { wait = d @@ -142,6 +121,9 @@ func (t *ToolSet) loop(ctx context.Context) { case <-t.wake: timer.Stop() case <-timer.C: + if ctx.Err() != nil { + return + } t.fireDue(ctx, t.now()) } } diff --git a/pkg/tools/builtin/scheduler/scheduler_test.go b/pkg/tools/builtin/scheduler/scheduler_test.go index efd6c1d845..b45bde0315 100644 --- a/pkg/tools/builtin/scheduler/scheduler_test.go +++ b/pkg/tools/builtin/scheduler/scheduler_test.go @@ -254,8 +254,10 @@ func TestToolSetImplementsInterfaces(t *testing.T) { var ts tools.ToolSet = New() - _, ok := ts.(tools.Startable) - require.True(t, ok, "must implement tools.Startable") + _, ok := ts.(tools.Runnable) + require.True(t, ok, "must implement tools.Runnable") + _, ok = ts.(tools.Startable) + require.False(t, ok, "Run must be owned by the caller") _, ok = ts.(tools.Instructable) require.True(t, ok, "must implement tools.Instructable") @@ -267,7 +269,7 @@ func TestToolSetImplementsInterfaces(t *testing.T) { func TestStartStopIsClean(t *testing.T) { t.Parallel() - ts := New() + ts := tools.NewStartable(New()) require.NoError(t, ts.Start(t.Context())) require.NoError(t, ts.Stop(t.Context())) } diff --git a/pkg/tools/capabilities.go b/pkg/tools/capabilities.go index 710b2912f6..b4f6ae0e1f 100644 --- a/pkg/tools/capabilities.go +++ b/pkg/tools/capabilities.go @@ -14,6 +14,14 @@ type Startable interface { Stop(ctx context.Context) error } +// Runnable is implemented by toolsets with background work but no startup handshake. +// Run blocks until cancellation and joins all work before returning. It must not +// run concurrently on the same toolset. NewStartable owns a detached run until +// Stop; Startable takes precedence when both capabilities are present. +type Runnable interface { + Run(ctx context.Context) +} + // StartReporter is implemented by toolsets whose live lifecycle state can // be queried independently of the StartableToolSet wrapper's latched state // (e.g. an MCP toolset whose supervisor lost the session in the background). diff --git a/pkg/tools/runnable.go b/pkg/tools/runnable.go new file mode 100644 index 0000000000..54900b6b2e --- /dev/null +++ b/pkg/tools/runnable.go @@ -0,0 +1,57 @@ +package tools + +import "context" + +// toolsetRun bridges blocking workers to the host's Start/Stop lifecycle. +// Only the canonical wrapper owns cancellation and joining; workers keep +// their resources local to Run. +type toolsetRun struct { + cancel context.CancelFunc + done <-chan struct{} +} + +func startToolsetRun(ctx context.Context, runnable Runnable) *toolsetRun { + ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + done := make(chan struct{}) + go func() { + defer close(done) + defer cancel() + runnable.Run(ctx) + }() + return &toolsetRun{cancel: cancel, done: done} +} + +func (r *toolsetRun) finished() bool { + select { + case <-r.done: + return true + default: + return false + } +} + +func (r *toolsetRun) stop() { + r.cancel() + <-r.done +} + +// Blocking reporters must settle inner lifecycle operations before recovery. +func (s *StartableToolSet) isStartedLocked(wait bool) bool { + if !s.started || (s.run != nil && s.run.finished()) { + return false + } + if owner, ok := As[*StartableToolSet](s.ToolSet); ok { + if wait { + return owner.IsStarted() + } + return owner.TryIsStarted() + } + return true +} + +func nonblockingReporterStarted(reporter StartReporter) bool { + if owner, ok := reporter.(*StartableToolSet); ok { + return owner.TryIsHealthy() + } + return reporter.IsStarted() +} diff --git a/pkg/tools/runnable_probe_test.go b/pkg/tools/runnable_probe_test.go new file mode 100644 index 0000000000..969a1feda6 --- /dev/null +++ b/pkg/tools/runnable_probe_test.go @@ -0,0 +1,72 @@ +package tools_test + +import ( + "sync/atomic" + "testing" + "testing/synctest" + + "github.com/docker/docker-agent/pkg/tools" +) + +type probeBoundary struct { + tools.ToolSet + + calls atomic.Int32 + enabled atomic.Bool + reached chan struct{} + proceed chan struct{} +} + +func (w *probeBoundary) Unwrap() tools.ToolSet { + if w.enabled.Load() && w.calls.Add(1) == 2 { + close(w.reached) + <-w.proceed + } + return w.ToolSet +} + +func TestNestedLifecycleProbeRace(t *testing.T) { + t.Parallel() + for _, kind := range []string{"healthy", "available"} { + t.Run(kind, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + raw := &blockingRestartToolSet{entered: make(chan struct{}, 2), release: make(chan struct{})} + owner := tools.NewStartable(raw) + boundary := &probeBoundary{ToolSet: owner, reached: make(chan struct{}), proceed: make(chan struct{})} + middle := tools.NewStartable(boundary) + outer := tools.NewStartable(middle) + if err := outer.Start(t.Context()); err != nil { + t.Fatal(err) + } + boundary.enabled.Store(true) + done := make(chan bool, 1) + go func() { + if kind == "healthy" { + done <- outer.TryIsHealthy() + } else { + done <- outer.TryIsAvailable() + } + }() + <-boundary.reached + restarted := make(chan error, 1) + go func() { restarted <- owner.RestartIfSupported(t.Context()) }() + <-raw.entered + close(boundary.proceed) + synctest.Wait() + select { + case <-done: + default: + t.Error("nonblocking probe waits on a nested lifecycle operation after its successful precheck") + } + close(raw.release) + if err := <-restarted; err != nil { + t.Error(err) + } + synctest.Wait() + if err := outer.Stop(t.Context()); err != nil { + t.Error(err) + } + }) + }) + } +} diff --git a/pkg/tools/runnable_test.go b/pkg/tools/runnable_test.go new file mode 100644 index 0000000000..5cf7745ce2 --- /dev/null +++ b/pkg/tools/runnable_test.go @@ -0,0 +1,332 @@ +package tools_test + +import ( + "context" + "errors" + "strconv" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/tools" +) + +type runnableToolSet struct { + run func(context.Context) +} + +func (*runnableToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (r *runnableToolSet) Run(ctx context.Context) { r.run(ctx) } + +func TestRunnableDetachedStartup(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + type contextKey struct{} + started := make(chan context.Context, 1) + finished := make(chan struct{}) + s := tools.NewStartable(tools.WithName(&runnableToolSet{run: func(ctx context.Context) { + defer close(finished) + started <- ctx + <-ctx.Done() + }}, "worker")) + t.Cleanup(func() { require.NoError(t, s.Stop(t.Context())) }) + ctx, cancel := context.WithTimeout(context.WithValue(t.Context(), contextKey{}, "value"), time.Second) + defer cancel() + ready, err := s.TryStartWithTimeout(ctx, time.Second) + require.NoError(t, err) + require.True(t, ready) + runCtx := <-started + cancel() + synctest.Wait() + require.NoError(t, runCtx.Err()) + assert.Equal(t, "value", runCtx.Value(contextKey{})) + _, hasDeadline := runCtx.Deadline() + assert.False(t, hasDeadline) + require.NoError(t, s.Stop(ctx)) + <-finished + require.ErrorIs(t, runCtx.Err(), context.Canceled) + assert.False(t, s.IsStarted()) + }) +} + +func TestRunnableConcurrentLifecycle(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + var runs, active atomic.Int32 + var overlap atomic.Bool + inner := &runnableToolSet{run: func(ctx context.Context) { + if active.Add(1) != 1 { + overlap.Store(true) + } + defer active.Add(-1) + runs.Add(1) + <-ctx.Done() + }} + // Nested lifecycle wrappers must delegate instead of starting a second run. + s := tools.NewStartable(tools.WithName(tools.NewStartable(inner), "worker")) + t.Cleanup(func() { require.NoError(t, s.Stop(t.Context())) }) + for range 3 { + var callers sync.WaitGroup + for range 20 { + callers.Go(func() { assert.NoError(t, s.Start(t.Context())) }) + } + callers.Wait() + synctest.Wait() + assert.EqualValues(t, 1, active.Load()) + for range 20 { + callers.Go(func() { assert.NoError(t, s.Stop(t.Context())) }) + } + callers.Wait() + assert.Zero(t, active.Load()) + } + assert.EqualValues(t, 3, runs.Load()) + assert.False(t, overlap.Load()) + }) +} + +func TestRunnableStopJoinsBeforeRestart(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + started := make(chan struct{}, 2) + cancelled, release := make(chan struct{}), make(chan struct{}) + var runs atomic.Int32 + s := tools.NewStartable(&runnableToolSet{run: func(ctx context.Context) { + generation := runs.Add(1) + started <- struct{}{} + <-ctx.Done() + if generation == 1 { + close(cancelled) + <-release + } + }}) + t.Cleanup(func() { require.NoError(t, s.Stop(t.Context())) }) + require.NoError(t, s.Start(t.Context())) + <-started + ctx, cancel := context.WithCancel(t.Context()) + cancel() + stopped := make(chan error, 1) + go func() { stopped <- s.Stop(ctx) }() + <-cancelled + restarted := make(chan error, 1) + go func() { restarted <- s.Start(t.Context()) }() + synctest.Wait() + select { + case err := <-stopped: + t.Fatalf("Stop returned before the worker exited: %v", err) + default: + } + assert.EqualValues(t, 1, runs.Load()) + close(release) + require.NoError(t, <-stopped) + require.NoError(t, <-restarted) + <-started + assert.EqualValues(t, 2, runs.Load()) + }) +} + +func TestRunnableCompletionUpdatesState(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + var runs atomic.Int32 + release := make(chan struct{}) + s := tools.NewStartable(&runnableToolSet{run: func(context.Context) { + runs.Add(1) + <-release + }}) + require.NoError(t, s.StopIfStarted(t.Context())) + assert.Zero(t, runs.Load()) + require.NoError(t, s.Start(t.Context())) + synctest.Wait() + require.True(t, s.IsStarted()) + close(release) + synctest.Wait() + assert.False(t, s.IsStarted()) + assert.False(t, s.TryIsStarted()) + assert.False(t, s.TryIsHealthy()) + assert.False(t, s.TryIsAvailable()) + started, inFlight := s.TryState() + assert.False(t, started) + assert.False(t, inFlight) + require.NoError(t, s.Start(t.Context())) + synctest.Wait() + assert.EqualValues(t, 2, runs.Load()) + require.NoError(t, s.StopIfStarted(t.Context())) + require.NoError(t, s.StopIfStarted(t.Context())) + }) +} + +type bothLifecycleToolSet struct { + runnableToolSet + + starts, stops int +} + +func (b *bothLifecycleToolSet) Start(context.Context) error { b.starts++; return nil } +func (b *bothLifecycleToolSet) Stop(context.Context) error { b.stops++; return nil } + +func TestRunnablePreservesStartableOwnership(t *testing.T) { + t.Parallel() + inner := &bothLifecycleToolSet{runnableToolSet: runnableToolSet{run: func(context.Context) { + t.Error("Run must not bypass Startable") + }}} + s := tools.NewStartable(tools.WithName(inner, "worker")) + require.NoError(t, s.Start(t.Context())) + require.NoError(t, s.Start(t.Context())) + require.NoError(t, s.Stop(t.Context())) + assert.Equal(t, 1, inner.starts) + assert.Equal(t, 1, inner.stops) +} + +func TestRunnablePendingStopHandoff(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + started := make(chan struct{}) + finished := make(chan struct{}) + s := tools.NewStartable(&runnableToolSet{run: func(ctx context.Context) { + defer close(finished) + close(started) + <-ctx.Done() + }}) + require.NoError(t, s.Start(t.Context())) + <-started + s.ExportedPublishStopRequest(t.Context()) + assert.False(t, s.TryIsStarted()) + <-finished + require.NoError(t, s.StopIfStarted(t.Context())) + }) +} + +type restartableRunnable struct { + runnableToolSet + + restarts atomic.Int32 +} + +func (r *restartableRunnable) Restart(context.Context) error { + r.restarts.Add(1) + return errors.New("restart failed") +} + +func TestRunnableRejectsUnownedRestart(t *testing.T) { + t.Parallel() + for _, pending := range []bool{false, true} { + t.Run(strconv.FormatBool(pending), func(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + started := make(chan context.Context, 2) + inner := &restartableRunnable{runnableToolSet: runnableToolSet{run: func(ctx context.Context) { + started <- ctx + <-ctx.Done() + }}} + s := tools.NewStartable(inner) + t.Cleanup(func() { require.NoError(t, s.Stop(t.Context())) }) + require.NoError(t, s.Start(t.Context())) + runCtx := <-started + assert.False(t, s.CanRestart()) + require.ErrorContains(t, s.RestartIfSupported(t.Context()), "does not support restart") + assert.Zero(t, inner.restarts.Load()) + assert.True(t, s.IsStarted(), "rejected restart must not lose the running worker") + if pending { + s.ExportedPublishStopRequest(t.Context()) + assert.False(t, s.TryIsStarted()) + } else { + require.NoError(t, s.StopIfStarted(t.Context())) + } + require.ErrorIs(t, runCtx.Err(), context.Canceled) + require.NoError(t, s.Start(t.Context())) + runCtx = <-started + require.NoError(t, runCtx.Err()) + }) + }) + } +} + +func TestRunnableNestedRestartOwnership(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + started := make(chan context.Context, 2) + release := make(chan struct{}) + var runs atomic.Int32 + inner := &restartableRunnable{runnableToolSet: runnableToolSet{run: func(ctx context.Context) { + generation := runs.Add(1) + started <- ctx + if generation == 1 { + select { + case <-release: + case <-ctx.Done(): + } + return + } + <-ctx.Done() + }}} + owner := tools.NewStartable(inner) + middle := tools.NewStartable(tools.WithName(owner, "worker")) + outer := tools.NewStartable(middle) + t.Cleanup(func() { require.NoError(t, outer.Stop(t.Context())) }) + require.NoError(t, outer.Start(t.Context())) + first := <-started + assert.False(t, owner.CanRestart()) + assert.False(t, outer.CanRestart()) + require.ErrorContains(t, outer.RestartIfSupported(t.Context()), "does not support restart") + assert.True(t, outer.IsStarted()) + assert.Zero(t, inner.restarts.Load()) + close(release) + synctest.Wait() + require.ErrorIs(t, first.Err(), context.Canceled) + assert.False(t, middle.IsStarted()) + assert.False(t, outer.IsStarted()) + assert.False(t, outer.TryIsHealthy()) + require.NoError(t, outer.Start(t.Context())) + second := <-started + assert.True(t, owner.IsStarted()) + assert.True(t, outer.IsStarted()) + assert.EqualValues(t, 2, runs.Load()) + assert.Zero(t, inner.restarts.Load()) + require.NoError(t, outer.StopIfStarted(t.Context())) + require.ErrorIs(t, second.Err(), context.Canceled) + }) +} + +type blockingRestartToolSet struct { + entered chan struct{} + release chan struct{} + calls atomic.Int32 +} + +func (*blockingRestartToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (*blockingRestartToolSet) Start(context.Context) error { return nil } +func (*blockingRestartToolSet) Stop(context.Context) error { return nil } +func (r *blockingRestartToolSet) Restart(context.Context) error { + r.calls.Add(1) + r.entered <- struct{}{} + <-r.release + return nil +} + +func TestNestedLifecycleBusyIsNotLoss(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + inner := &blockingRestartToolSet{entered: make(chan struct{}, 2), release: make(chan struct{})} + owner := tools.NewStartable(inner) + middle := tools.NewStartable(owner) + outer := tools.NewStartable(middle) + require.NoError(t, outer.Start(t.Context())) + restarted := make(chan error, 1) + go func() { restarted <- owner.RestartIfSupported(t.Context()) }() + <-inner.entered + started := make(chan error, 1) + go func() { started <- outer.Start(t.Context()) }() + synctest.Wait() + assert.False(t, middle.TryIsHealthy(), "nonblocking probes must not wait for the inner lock") + close(inner.release) + require.NoError(t, <-restarted) + require.NoError(t, <-started) + assert.EqualValues(t, 1, inner.calls.Load(), "ordinary Start must not queue a redundant restart") + require.NoError(t, outer.StopIfStarted(t.Context())) + }) +} diff --git a/pkg/tools/startable.go b/pkg/tools/startable.go index e33e253117..4cfee712fd 100644 --- a/pkg/tools/startable.go +++ b/pkg/tools/startable.go @@ -277,6 +277,7 @@ type StartableToolSet struct { started bool attempted bool recovering bool + run *toolsetRun stopErr error // Retained until a successful start; pending-stop handoffs must not lose failures. // stopRequestMu guards the stop request that StopIfStarted publishes @@ -339,7 +340,7 @@ func (s *StartableToolSet) nowFn() time.Time { } // IsStarted returns whether the toolset has been successfully started. -// For toolsets that don't implement Startable, this always returns true. +// Runnable toolsets report false after Run returns. // It waits for any in-flight lifecycle operation (Start/Stop) to settle, // so callers deciding whether a Stop is needed observe the final state; // use TryIsStarted on paths that must never block. Releasing the lifecycle @@ -348,7 +349,7 @@ func (s *StartableToolSet) nowFn() time.Time { func (s *StartableToolSet) IsStarted() bool { s.mu.Lock() defer s.unlock() - return s.started + return s.isStartedLocked(true) } // TryIsStarted is the non-blocking variant of IsStarted: when another @@ -385,7 +386,7 @@ func (s *StartableToolSet) TryState() (started, inFlight bool) { if !s.mu.TryLock() { return false, true } - started = s.started + started = s.isStartedLocked(false) if s.unlock() { return false, false } @@ -399,10 +400,10 @@ func (s *StartableToolSet) TryIsHealthy() bool { if !s.mu.TryLock() { return false } - healthy := s.started + healthy := s.isStartedLocked(false) if healthy { if reporter, ok := As[StartReporter](s.ToolSet); ok { - healthy = reporter.IsStarted() + healthy = nonblockingReporterStarted(reporter) } } if s.unlock() { @@ -419,10 +420,10 @@ func (s *StartableToolSet) TryIsAvailable() bool { if !s.mu.TryLock() { return false } - available := !s.attempted || s.started + available := !s.attempted || s.isStartedLocked(false) if available && s.started { if reporter, ok := As[StartReporter](s.ToolSet); ok { - available = reporter.IsStarted() + available = nonblockingReporterStarted(reporter) } } if s.unlock() { @@ -434,7 +435,8 @@ func (s *StartableToolSet) TryIsAvailable() bool { // Start starts the toolset with single-flight semantics. // Concurrent callers block until the start attempt completes. // If start fails, a future call will retry. -// If the underlying toolset doesn't implement Startable, this is a no-op. +// Runnable toolsets run until Stop, independent of the startup context. +// Without Startable or Runnable, this is a no-op. func (s *StartableToolSet) Start(ctx context.Context) error { s.mu.Lock() defer s.unlock() @@ -566,6 +568,13 @@ func (s *StartableToolSet) tryStartLocked(ctx context.Context) error { // startLocked implements the start sequence shared by Start and TryStart. // s.mu must be held. func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { + if s.run != nil { + if !s.run.finished() { + return nil + } + s.run = nil + s.started = false + } recovering := s.recovering if recovering { if reporter, ok := As[StartReporter](s.ToolSet); ok && reporter.IsStarted() { @@ -605,7 +614,7 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { if u, ok := inner.(Unwrapper); ok { inner = u.Unwrap() } - if restarter, hasRestarter := As[Restartable](s.ToolSet); recovering && hasRestarter { + if restart, hasRestarter := s.restartFunc(); recovering && hasRestarter { s.attempted = true ctx, span := otel.Tracer("github.com/docker/docker-agent/pkg/tools").Start( ctx, @@ -620,7 +629,7 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { } span.End() }() - if err := restarter.Restart(ctx); err != nil { + if err := restart(ctx); err != nil { s.recovering = true s.startStreak.fail() s.recoveryStreak.fail() @@ -682,6 +691,9 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { } return err } + } else if runnable, ok := As[Runnable](s.ToolSet); ok { + s.attempted = true + s.run = startToolsetRun(ctx, runnable) } // Successful start: clear streaks and backoff so any future failure is @@ -719,10 +731,33 @@ func (s *StartableToolSet) Tools(ctx context.Context) ([]Tool, error) { // Callers must use Restart on this wrapper so the lifecycle latch, failure // streaks, and retry gate stay synchronized with the underlying toolset. func (s *StartableToolSet) CanRestart() bool { - _, ok := As[Restartable](s.ToolSet) + _, ok := s.restartFunc() return ok } +// restartFunc must not discover capabilities past another lifecycle owner. +func (s *StartableToolSet) restartFunc() (func(context.Context) error, bool) { + if _, startable := As[Startable](s.ToolSet); !startable { + if _, runnable := As[Runnable](s.ToolSet); runnable { + return nil, false + } + } + for ts := s.ToolSet; ts != nil; { + if owner, ok := ts.(*StartableToolSet); ok { + return owner.RestartIfSupported, owner.CanRestart() + } + if restarter, ok := ts.(Restartable); ok { + return restarter.Restart, true + } + unwrapper, ok := ts.(Unwrapper) + if !ok { + break + } + ts = unwrapper.Unwrap() + } + return nil, false +} + // RestartIfSupported restarts the underlying toolset under the canonical // lifecycle lock. It bypasses the retry gate because an explicit operator // action must run immediately, then synchronizes all wrapper-owned state. @@ -730,7 +765,7 @@ func (s *StartableToolSet) RestartIfSupported(ctx context.Context) error { s.mu.Lock() defer s.unlock() - restarter, ok := As[Restartable](s.ToolSet) + restart, ok := s.restartFunc() if !ok { return errors.New("toolset does not support restart") } @@ -738,7 +773,7 @@ func (s *StartableToolSet) RestartIfSupported(ctx context.Context) error { wasStarted := s.started s.attempted = true s.started = false - if err := restarter.Restart(ctx); err != nil { + if err := restart(ctx); err != nil { s.recovering = wasStarted || s.recovering s.startStreak.fail() if wasStarted { @@ -758,7 +793,7 @@ func (s *StartableToolSet) RestartIfSupported(ctx context.Context) error { return nil } -// Stop stops the toolset if it implements Startable and resets +// Stop stops and joins the toolset and resets // the started flag so that a subsequent Start will re-initialize. func (s *StartableToolSet) Stop(ctx context.Context) error { s.mu.Lock() @@ -800,7 +835,7 @@ func (s *StartableToolSet) StopIfStarted(ctx context.Context) error { s.stopRequestCtx = nil s.stopRequestMu.Unlock() - if !s.started { + if !s.started && s.run == nil { return s.stopErr } return s.stopLocked(ctx) @@ -817,7 +852,10 @@ func (s *StartableToolSet) stopLocked(ctx context.Context) error { s.recoveryStreak.reset() s.resetStartBackoff() s.stopErr = nil - if startable, ok := As[Startable](s.ToolSet); ok { + if s.run != nil { + s.run.stop() + s.run = nil + } else if startable, ok := As[Startable](s.ToolSet); ok { s.stopErr = startable.Stop(ctx) } return s.stopErr @@ -863,7 +901,7 @@ func (s *StartableToolSet) settleStopRequestAndRelease() bool { ctx := context.WithoutCancel(s.stopRequestCtx) s.stopRequestCtx = nil s.stopRequestMu.Unlock() - if s.started { + if s.started || s.run != nil { if err := s.stopLocked(ctx); err != nil { slog.WarnContext(ctx, "Failed to stop toolset for a pending shutdown request", "toolset", DescribeToolSet(s), "error", err) }