fix: propagate recovery through nested tool owners

Assisted-By: OpenAI

pkg/teamloader/workspace_recovery_test.go

@@ -0,0 +1,88 @@
+package teamloader
+
+import (
+	"context"
+	"testing"
+
+	"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"
+	"github.com/docker/docker-agent/pkg/tools/codemode"
+)
+
+type workspaceRecoveringTool struct {
+	healthy  bool
+	starts   int
+	restarts int
+}
+
+func (*workspaceRecoveringTool) Tools(context.Context) ([]tools.Tool, error) {
+	return []tools.Tool{{Name: "recovering", Annotations: tools.ToolAnnotations{ReadOnlyHint: true}}}, nil
+}
+func (s *workspaceRecoveringTool) IsStarted() bool { return s.healthy }
+
+func (s *workspaceRecoveringTool) Start(context.Context) error {
+	s.starts++
+	s.healthy = true
+	return nil
+}
+
+func (s *workspaceRecoveringTool) Restart(context.Context) error {
+	s.restarts++
+	s.healthy = true
+	return nil
+}
+
+func (s *workspaceRecoveringTool) Stop(context.Context) error { s.healthy = false; return nil }
+
+func TestWorkspaceNestedOwnersRecoverUnhealthyTools(t *testing.T) {
+	t.Setenv("OPENAI_API_KEY", "dummy")
+	for _, codeMode := range []bool{false, true} {
+		name := "plain"
+		if codeMode {
+			name = "code mode"
+		}
+		t.Run(name, func(t *testing.T) {
+			var raw []*workspaceRecoveringTool
+			registry := NewToolsetRegistry(map[string]ToolsetCreator{
+				"recovering": func(context.Context, latest.Toolset, string, *config.RuntimeConfig, string) (tools.ToolSet, error) {
+					ts := &workspaceRecoveringTool{}
+					raw = append(raw, ts)
+					return ts, nil
+				},
+			})
+			source := config.NewBytesSource("recovery.yaml", []byte(`agents:
+  root:
+    model: openai/gpt-4o
+    toolsets:
+      - type: recovering
+        tools: [recovering]
+        readonly: true
+`))
+			tm, err := Load(t.Context(), source, &config.RuntimeConfig{Config: config.Config{GlobalCodeMode: codeMode}},
+				withTestProviderRegistry(WithToolsetRegistry(registry), WithCodeMode(codemode.Wrap))...)
+			require.NoError(t, err)
+			t.Cleanup(func() { assert.NoError(t, tm.Dispose(context.WithoutCancel(t.Context()))) })
+			a, err := tm.DefaultAgent()
+			require.NoError(t, err)
+			rebuilt, err := a.ForWorkingDir(t.Context(), t.TempDir())
+			require.NoError(t, err)
+			t.Cleanup(func() { assert.NoError(t, rebuilt.Dispose(context.WithoutCancel(t.Context()))) })
+			require.Len(t, raw, 2)
+			for i, toolsets := range [][]tools.ToolSet{a.ToolSets(), rebuilt.ToolSets()} {
+				require.Len(t, toolsets, 1)
+				owner := toolsets[0].(*tools.StartableToolSet)
+				require.NoError(t, owner.Start(t.Context()))
+				raw[i].healthy = false
+				assert.False(t, owner.TryIsHealthy())
+				require.NoError(t, owner.Start(t.Context()))
+				assert.True(t, owner.TryIsHealthy())
+				assert.Equal(t, 1, raw[i].starts)
+				assert.Equal(t, 1, raw[i].restarts)
+			}
+		})
+	}
+}

pkg/tools/runnable.go

@@ -55,3 +55,20 @@ func nonblockingReporterStarted(reporter StartReporter) bool {
 	}
 	return reporter.IsStarted()
 }
+
+// Recovery waits for an inner owner and then checks health, not just its latch.
+func blockingReporterStarted(reporter StartReporter) bool {
+	owner, ok := reporter.(*StartableToolSet)
+	if !ok {
+		return reporter.IsStarted()
+	}
+	owner.mu.Lock()
+	defer owner.unlock()
+	if owner.disposed || !owner.isStartedLocked(true) {
+		return false
+	}
+	if inner, ok := As[StartReporter](owner.ToolSet); ok {
+		return blockingReporterStarted(inner)
+	}
+	return true
+}

pkg/tools/startable.go

@@ -589,7 +589,7 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) {
 	}
 	recovering := s.recovering
 	if recovering {
-		if reporter, ok := As[StartReporter](s.ToolSet); ok && reporter.IsStarted() {
+		if reporter, ok := As[StartReporter](s.ToolSet); ok && blockingReporterStarted(reporter) {
 			s.started = true
 			s.stopErr = nil
 			s.recovering = false
@@ -600,7 +600,7 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) {
 		}
 	}
 	if s.started {
-		if reporter, ok := As[StartReporter](s.ToolSet); !ok || reporter.IsStarted() {
+		if reporter, ok := As[StartReporter](s.ToolSet); !ok || blockingReporterStarted(reporter) {
 			return nil
 		}
 		s.started = false

pkg/tools/startable_test.go

@@ -1240,3 +1240,57 @@ func TestStartableToolSet_ZeroValueIsUsable(t *testing.T) {
 	assert.NilError(t, s.StopIfStarted(t.Context()))
 	assert.Check(t, is.Equal(s.IsStarted(), false))
 }
+
+func TestNestedStartableRecoversUnhealthyTool(t *testing.T) {
+	t.Parallel()
+	raw := &reportingToolSet{}
+	inner := tools.NewStartable(raw)
+	outer := tools.NewStartable(tools.WithName(tools.NewStartable(inner), "nested"))
+	assert.NilError(t, outer.Start(t.Context()))
+	raw.started = false
+	assert.Check(t, outer.IsStarted(), "latch remains set until recovery")
+	assert.Check(t, !outer.TryIsHealthy())
+	assert.NilError(t, outer.Start(t.Context()))
+	assert.Check(t, is.Equal(raw.startCalls, 1))
+	assert.Check(t, is.Equal(raw.restartCalls, 1))
+	assert.Check(t, outer.TryIsHealthy())
+	failure := errors.New("restart failed")
+	raw.started, raw.restartErr = false, failure
+	assert.ErrorIs(t, outer.Start(t.Context()), failure)
+	raw.restartErr = nil
+	assert.NilError(t, outer.Start(t.Context()))
+	assert.Check(t, is.Equal(raw.restartCalls, 3))
+	assert.Check(t, outer.TryIsHealthy())
+}
+
+func TestNestedStartableRecoversWithStartFallback(t *testing.T) {
+	t.Parallel()
+	raw := &reportingStartOnlyToolSet{}
+	outer := tools.NewStartable(tools.NewStartable(raw))
+	assert.NilError(t, outer.Start(t.Context()))
+	raw.started = false
+	assert.NilError(t, outer.Start(t.Context()))
+	assert.Check(t, is.Equal(raw.startCalls, 2))
+	assert.Check(t, outer.TryIsHealthy())
+}
+
+func TestNestedStartableRetriesPartialStart(t *testing.T) {
+	t.Parallel()
+	partial := tools.NewPartialStartError(errors.New("unavailable"))
+	raw := &partialFlappyToolSet{errs: []error{partial, partial, nil}}
+	outer := tools.NewStartable(tools.NewStartable(raw))
+	assert.ErrorIs(t, outer.Start(t.Context()), partial)
+	assert.Check(t, outer.IsStarted())
+	assert.Check(t, !outer.TryIsHealthy())
+	listed, err := outer.Tools(t.Context())
+	assert.NilError(t, err)
+	assert.Check(t, is.Len(listed, 1))
+	assert.Check(t, !outer.ShouldReportRecoveryFailure())
+	assert.ErrorIs(t, outer.Start(t.Context()), partial)
+	assert.Check(t, is.Equal(raw.starts, 2))
+	assert.Check(t, !outer.ShouldReportRecoveryFailure())
+	assert.Check(t, outer.IsStarted())
+	assert.NilError(t, outer.Start(t.Context()))
+	assert.Check(t, is.Equal(raw.starts, 3))
+	assert.Check(t, outer.TryIsHealthy())
+}