mirror of
https://github.com/docker/docker-agent.git
synced 2026-11-04 10:02:15 +00:00
refactor(runtime): extract cooldownManager from fallback.go
The fallback cooldown logic was scattered across three methods on *LocalRuntime (getCooldownState/setCooldownState/clearCooldownState), each repeating the same map+mutex+slog ceremony, plus the eviction-on-read invariant lived inline in getCooldownState. Reasoning about that invariant required reading 60 lines of fallback.go and was untestable without spinning up an agent + team + runtime. Move the map, mutex, eviction logic, and slog emissions into a small cooldownManager type with Get/Set/Clear and an injectable clock. The three runtime methods become one-line delegations. The clock is wired through r.now so WithClock(fakeClock) makes the eviction window deterministic and the new test exercises the contract directly under -race without going through tryModelWithFallback. Net effect on fallback.go: -50 lines, +50 lines of standalone testable type, +8 unit tests covering Set/Get/Clear, eviction, replacement, isolation, nil-clock default, and concurrent safety. Assisted-By: docker-agent
This commit is contained in:
4 files changed
+257
-54
No files matched your search
@@ -0,0 +1,84 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// cooldownManager owns the per-agent fallback-cooldown map. The runtime
|
||||
// activates a cooldown when the primary model fails with a non-retryable
|
||||
// error and a fallback succeeds; subsequent calls for that agent skip the
|
||||
// primary until the window expires.
|
||||
//
|
||||
// Pulled out of fallback.go so the eviction-on-read invariant, the slog
|
||||
// emissions, and the mutex contract live in one small, race-tested place
|
||||
// instead of being threaded through three methods on *LocalRuntime.
|
||||
//
|
||||
// All methods are safe for concurrent use.
|
||||
type cooldownManager struct {
|
||||
mu sync.RWMutex
|
||||
now func() time.Time
|
||||
by map[string]*fallbackCooldownState
|
||||
}
|
||||
|
||||
// newCooldownManager constructs a cooldownManager that reads time via now.
|
||||
// Tests pass a fake clock so they can verify the eviction window without
|
||||
// real wall-clock advancement.
|
||||
func newCooldownManager(now func() time.Time) *cooldownManager {
|
||||
if now == nil {
|
||||
now = time.Now
|
||||
}
|
||||
return &cooldownManager{
|
||||
now: now,
|
||||
by: make(map[string]*fallbackCooldownState),
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns the active cooldown state for agentName, or nil when no
|
||||
// cooldown is active. Expired entries are evicted as a side-effect of the
|
||||
// read so a long-running runtime serving many agents stays bounded.
|
||||
func (c *cooldownManager) Get(agentName string) *fallbackCooldownState {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
state := c.by[agentName]
|
||||
if state == nil {
|
||||
return nil
|
||||
}
|
||||
if c.now().After(state.until) {
|
||||
delete(c.by, agentName)
|
||||
return nil
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
// Set activates a cooldown for agentName, pinning the run loop to the
|
||||
// fallback at fallbackIndex for duration. Replaces any existing entry.
|
||||
func (c *cooldownManager) Set(agentName string, fallbackIndex int, duration time.Duration) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
state := &fallbackCooldownState{
|
||||
fallbackIndex: fallbackIndex,
|
||||
until: c.now().Add(duration),
|
||||
}
|
||||
c.by[agentName] = state
|
||||
|
||||
slog.Info("Fallback cooldown activated",
|
||||
"agent", agentName,
|
||||
"fallback_index", fallbackIndex,
|
||||
"cooldown", duration,
|
||||
"until", state.until.Format(time.RFC3339))
|
||||
}
|
||||
|
||||
// Clear evicts any cooldown entry for agentName. No-op if none is active.
|
||||
func (c *cooldownManager) Clear(agentName string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if _, exists := c.by[agentName]; exists {
|
||||
delete(c.by, agentName)
|
||||
slog.Debug("Fallback cooldown cleared", "agent", agentName)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCooldownManager_GetMissingReturnsNil(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
m := newCooldownManager(time.Now)
|
||||
assert.Nil(t, m.Get("nobody"))
|
||||
}
|
||||
|
||||
func TestCooldownManager_SetAndGet(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := newFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC))
|
||||
m := newCooldownManager(clock.Now)
|
||||
|
||||
m.Set("agent-A", 2, time.Minute)
|
||||
|
||||
state := m.Get("agent-A")
|
||||
require.NotNil(t, state)
|
||||
assert.Equal(t, 2, state.fallbackIndex)
|
||||
assert.True(t, state.until.After(clock.Now()), "until must be in the future")
|
||||
}
|
||||
|
||||
func TestCooldownManager_GetEvictsExpired(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := newFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC))
|
||||
m := newCooldownManager(clock.Now)
|
||||
|
||||
m.Set("agent-A", 0, time.Minute)
|
||||
require.NotNil(t, m.Get("agent-A"), "cooldown should be active immediately after Set")
|
||||
|
||||
clock.Advance(59 * time.Second)
|
||||
assert.NotNil(t, m.Get("agent-A"), "cooldown should still be active just under the window")
|
||||
|
||||
clock.Advance(2 * time.Second)
|
||||
assert.Nil(t, m.Get("agent-A"), "expired cooldown should be evicted")
|
||||
|
||||
// Eviction is observable: a second Get also returns nil.
|
||||
assert.Nil(t, m.Get("agent-A"), "expired cooldown should stay evicted")
|
||||
}
|
||||
|
||||
func TestCooldownManager_SetReplacesExistingEntry(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := newFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC))
|
||||
m := newCooldownManager(clock.Now)
|
||||
|
||||
m.Set("agent-A", 0, time.Minute)
|
||||
m.Set("agent-A", 3, 5*time.Minute)
|
||||
|
||||
state := m.Get("agent-A")
|
||||
require.NotNil(t, state)
|
||||
assert.Equal(t, 3, state.fallbackIndex, "second Set must replace fallback index")
|
||||
|
||||
expectedUntil := clock.Now().Add(5 * time.Minute)
|
||||
assert.Equal(t, expectedUntil, state.until, "second Set must reset the until window")
|
||||
}
|
||||
|
||||
func TestCooldownManager_Clear(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := newFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC))
|
||||
m := newCooldownManager(clock.Now)
|
||||
|
||||
m.Set("agent-A", 0, time.Minute)
|
||||
m.Clear("agent-A")
|
||||
assert.Nil(t, m.Get("agent-A"))
|
||||
|
||||
// Clearing a non-existent entry is a no-op (no panic).
|
||||
assert.NotPanics(t, func() { m.Clear("never-set") })
|
||||
}
|
||||
|
||||
func TestCooldownManager_NilClockDefaultsToTimeNow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
m := newCooldownManager(nil)
|
||||
require.NotNil(t, m.now, "nil clock should be replaced by time.Now")
|
||||
|
||||
got := m.now()
|
||||
assert.WithinDuration(t, time.Now(), got, time.Second)
|
||||
}
|
||||
|
||||
func TestCooldownManager_PerAgentIsolation(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clock := newFakeClock(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC))
|
||||
m := newCooldownManager(clock.Now)
|
||||
|
||||
m.Set("agent-A", 0, time.Minute)
|
||||
m.Set("agent-B", 1, 2*time.Minute)
|
||||
|
||||
stateA := m.Get("agent-A")
|
||||
stateB := m.Get("agent-B")
|
||||
require.NotNil(t, stateA)
|
||||
require.NotNil(t, stateB)
|
||||
assert.Equal(t, 0, stateA.fallbackIndex)
|
||||
assert.Equal(t, 1, stateB.fallbackIndex)
|
||||
|
||||
// Expiring A leaves B intact.
|
||||
clock.Advance(90 * time.Second)
|
||||
assert.Nil(t, m.Get("agent-A"), "A should have expired")
|
||||
assert.NotNil(t, m.Get("agent-B"), "B must not be affected by A's expiry")
|
||||
}
|
||||
|
||||
// TestCooldownManager_ConcurrentSafety drives many concurrent Set/Get/Clear
|
||||
// to confirm the mutex contract under -race.
|
||||
func TestCooldownManager_ConcurrentSafety(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
m := newCooldownManager(time.Now)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := range 8 {
|
||||
wg.Add(3)
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
for j := range 100 {
|
||||
m.Set("agent", j%5, time.Hour)
|
||||
_ = id
|
||||
}
|
||||
}(i)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for range 100 {
|
||||
_ = m.Get("agent")
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for range 100 {
|
||||
m.Clear("agent")
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
+8
-40
@@ -78,53 +78,21 @@ func logRetryBackoff(agentName, modelID string, attempt int, backoffDelay time.D
|
||||
"backoff", backoffDelay)
|
||||
}
|
||||
|
||||
// getCooldownState returns the current cooldown state for an agent (thread-safe).
|
||||
// Returns nil if no cooldown is active or if cooldown has expired.
|
||||
// Expired entries are evicted to prevent stale state accumulation.
|
||||
// getCooldownState returns the active cooldown for an agent, or nil when
|
||||
// no cooldown is active or it has expired. Expired entries are evicted
|
||||
// inside the cooldown manager.
|
||||
func (r *LocalRuntime) getCooldownState(agentName string) *fallbackCooldownState {
|
||||
r.fallbackCooldownsMux.Lock()
|
||||
defer r.fallbackCooldownsMux.Unlock()
|
||||
|
||||
state := r.fallbackCooldowns[agentName]
|
||||
if state == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if cooldown has expired; evict if so
|
||||
if r.now().After(state.until) {
|
||||
delete(r.fallbackCooldowns, agentName)
|
||||
return nil
|
||||
}
|
||||
|
||||
return state
|
||||
return r.cooldowns.Get(agentName)
|
||||
}
|
||||
|
||||
// setCooldownState sets the cooldown state for an agent (thread-safe).
|
||||
// setCooldownState activates a cooldown for an agent.
|
||||
func (r *LocalRuntime) setCooldownState(agentName string, fallbackIndex int, cooldownDuration time.Duration) {
|
||||
r.fallbackCooldownsMux.Lock()
|
||||
defer r.fallbackCooldownsMux.Unlock()
|
||||
|
||||
r.fallbackCooldowns[agentName] = &fallbackCooldownState{
|
||||
fallbackIndex: fallbackIndex,
|
||||
until: r.now().Add(cooldownDuration),
|
||||
}
|
||||
|
||||
slog.Info("Fallback cooldown activated",
|
||||
"agent", agentName,
|
||||
"fallback_index", fallbackIndex,
|
||||
"cooldown", cooldownDuration,
|
||||
"until", r.fallbackCooldowns[agentName].until.Format(time.RFC3339))
|
||||
r.cooldowns.Set(agentName, fallbackIndex, cooldownDuration)
|
||||
}
|
||||
|
||||
// clearCooldownState clears the cooldown state for an agent (thread-safe).
|
||||
// clearCooldownState removes any active cooldown for an agent.
|
||||
func (r *LocalRuntime) clearCooldownState(agentName string) {
|
||||
r.fallbackCooldownsMux.Lock()
|
||||
defer r.fallbackCooldownsMux.Unlock()
|
||||
|
||||
if _, exists := r.fallbackCooldowns[agentName]; exists {
|
||||
delete(r.fallbackCooldowns, agentName)
|
||||
slog.Debug("Fallback cooldown cleared", "agent", agentName)
|
||||
}
|
||||
r.cooldowns.Clear(agentName)
|
||||
}
|
||||
|
||||
// getEffectiveCooldown returns the cooldown duration to use for an agent.
|
||||
|
||||
+19
-14
@@ -224,9 +224,10 @@ type LocalRuntime struct {
|
||||
// Library consumers can enable this via WithRetryOnRateLimit().
|
||||
retryOnRateLimit bool
|
||||
|
||||
// fallbackCooldowns tracks per-agent cooldown state for sticky fallback behavior
|
||||
fallbackCooldowns map[string]*fallbackCooldownState
|
||||
fallbackCooldownsMux sync.RWMutex
|
||||
// cooldowns owns the per-agent fallback-cooldown state. See
|
||||
// cooldown_manager.go for the eviction contract; the manager reads
|
||||
// the runtime's clock so WithClock makes cooldown windows testable.
|
||||
cooldowns *cooldownManager
|
||||
|
||||
currentAgentMu sync.RWMutex
|
||||
|
||||
@@ -407,29 +408,33 @@ func NewLocalRuntime(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
|
||||
}
|
||||
|
||||
r := &LocalRuntime{
|
||||
toolMap: make(map[string]ToolHandlerFunc),
|
||||
team: agents,
|
||||
currentAgent: defaultAgent.Name(),
|
||||
resumeChan: make(chan ResumeRequest),
|
||||
elicitationRequestCh: make(chan ElicitationResult),
|
||||
steerQueue: NewInMemoryMessageQueue(defaultSteerQueueCapacity),
|
||||
followUpQueue: NewInMemoryMessageQueue(defaultFollowUpQueueCapacity),
|
||||
sessionCompaction: true,
|
||||
managedOAuth: true,
|
||||
sessionStore: session.NewInMemorySessionStore(),
|
||||
fallbackCooldowns: make(map[string]*fallbackCooldownState),
|
||||
toolMap: make(map[string]ToolHandlerFunc),
|
||||
team: agents,
|
||||
currentAgent: defaultAgent.Name(),
|
||||
resumeChan: make(chan ResumeRequest),
|
||||
elicitationRequestCh: make(chan ElicitationResult),
|
||||
steerQueue: NewInMemoryMessageQueue(defaultSteerQueueCapacity),
|
||||
followUpQueue: NewInMemoryMessageQueue(defaultFollowUpQueueCapacity),
|
||||
sessionCompaction: true,
|
||||
managedOAuth: true,
|
||||
sessionStore: session.NewInMemorySessionStore(),
|
||||
hooksRegistry: hooksRegistry,
|
||||
builtinsState: builtinsState,
|
||||
now: time.Now,
|
||||
telemetry: defaultTelemetry{},
|
||||
maxOverflowCompactions: defaultMaxOverflowCompactions,
|
||||
}
|
||||
r.cooldowns = newCooldownManager(r.now)
|
||||
r.bgAgents = agenttool.NewHandler(r)
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(r)
|
||||
}
|
||||
|
||||
// If WithClock changed the runtime's clock, the cooldown manager must
|
||||
// see it too — rebuild it now that opts have run.
|
||||
r.cooldowns = newCooldownManager(r.now)
|
||||
|
||||
// Default the runtime's working directory to the process CWD when no
|
||||
// caller supplied one. This matches the session's default and ensures
|
||||
// builtin hooks that look up files (add_prompt_files) can find them
|
||||
|
||||
Reference in new issue
Block a user