refactor(scheduler): scope worker lifetime to blocking Run(ctx)

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.

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

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 {

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)
+			})
+		})
+	}
+}

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
+}

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())
+	})
+}