From 175e05287a496e4b96c0bbd4608370c38a1be03d Mon Sep 17 00:00:00 2001 From: viettranx Date: Sat, 28 Mar 2026 19:02:53 +0700 Subject: [PATCH] fix(cache): heartbeat/cron used UUID in session keys causing stale cache after agent update MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Session keys must use agentKey (human-readable, e.g. "default") not UUID. Cache invalidation via InvalidateAgent(agentKey) only matches agentKey suffix, so UUID-keyed entries were never cleared — causing stale provider/model after agent updates from the web UI. - heartbeat: use agentKey (already resolved) instead of UUID for session key and HasActiveSessionsForAgent check - cron: resolve agentKey from UUID via agentStore before building session key - agents_update: also invalidate by UUID as safety net for legacy entries - startup migration: auto-fix existing UUID-based session keys in DB - tests: 7 new tests covering session key format and cache invalidation --- cmd/gateway.go | 2 +- cmd/gateway_cron.go | 16 ++++- internal/agent/router_invalidate_test.go | 72 +++++++++++++++++++ internal/gateway/methods/agents_update.go | 3 + internal/heartbeat/ticker.go | 5 +- internal/sessions/key.go | 8 ++- internal/sessions/key_test.go | 87 +++++++++++++++++++++++ internal/store/pg/sessions.go | 38 ++++++++++ 8 files changed, 224 insertions(+), 7 deletions(-) create mode 100644 internal/agent/router_invalidate_test.go create mode 100644 internal/sessions/key_test.go diff --git a/cmd/gateway.go b/cmd/gateway.go index 8cceba25..e571a662 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -862,7 +862,7 @@ func runGateway() { defer sched.Stop() // Start cron service with job handler (routes through scheduler's cron lane) - pgStores.Cron.SetOnJob(makeCronJobHandler(sched, msgBus, cfg, channelMgr, pgStores.Sessions)) + pgStores.Cron.SetOnJob(makeCronJobHandler(sched, msgBus, cfg, channelMgr, pgStores.Sessions, pgStores.Agents)) pgStores.Cron.SetOnEvent(func(event store.CronEvent) { server.BroadcastEvent(*protocol.NewEvent(protocol.EventCron, event)) }) diff --git a/cmd/gateway_cron.go b/cmd/gateway_cron.go index 12e33fd3..8bb4d920 100644 --- a/cmd/gateway_cron.go +++ b/cmd/gateway_cron.go @@ -6,6 +6,8 @@ import ( "log/slog" "strings" + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/agent" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/channels" @@ -22,11 +24,18 @@ import ( // Safe because cron jobs only fire after Start(), well after this is set. var cronHeartbeatWakeFn func(agentID string) -func makeCronJobHandler(sched *scheduler.Scheduler, msgBus *bus.MessageBus, cfg *config.Config, channelMgr *channels.Manager, sessionMgr store.SessionStore) func(job *store.CronJob) (*store.CronJobResult, error) { +func makeCronJobHandler(sched *scheduler.Scheduler, msgBus *bus.MessageBus, cfg *config.Config, channelMgr *channels.Manager, sessionMgr store.SessionStore, agentStore store.AgentStore) func(job *store.CronJob) (*store.CronJobResult, error) { return func(job *store.CronJob) (*store.CronJobResult, error) { agentID := job.AgentID if agentID == "" { agentID = cfg.ResolveDefaultAgentID() + } else if id, err := uuid.Parse(agentID); err == nil && agentStore != nil { + // Resolve agentKey from UUID so session key uses agentKey + // (consistent with chat/WS/team paths, fixes cache invalidation mismatch). + cronCtx := store.WithTenantID(context.Background(), job.TenantID) + if ag, err := agentStore.GetByID(cronCtx, id); err == nil { + agentID = ag.AgentKey + } } else { agentID = config.NormalizeAgentID(agentID) } @@ -119,9 +128,10 @@ func makeCronJobHandler(sched *scheduler.Scheduler, msgBus *bus.MessageBus, cfg cronResult.OutputTokens = result.Usage.CompletionTokens } - // wakeMode: trigger heartbeat after cron job completes + // wakeMode: trigger heartbeat after cron job completes. + // Use original job.AgentID (UUID) — cronHeartbeatWakeFn expects UUID for ticker.Wake(). if job.Payload.WakeHeartbeat && cronHeartbeatWakeFn != nil { - cronHeartbeatWakeFn(agentID) + cronHeartbeatWakeFn(job.AgentID) } return cronResult, nil diff --git a/internal/agent/router_invalidate_test.go b/internal/agent/router_invalidate_test.go new file mode 100644 index 00000000..d5090ce9 --- /dev/null +++ b/internal/agent/router_invalidate_test.go @@ -0,0 +1,72 @@ +package agent + +import ( + "context" + "testing" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +// TestInvalidateAgent_MatchesAgentKeyNotUUID verifies that InvalidateAgent(agentKey) +// clears entries cached under agentKey but NOT entries cached under UUID. +// This documents the root cause of the stale-cache bug: heartbeat/cron paths +// previously used UUID in session keys, creating cache entries keyed by UUID +// that InvalidateAgent(agentKey) could never match. +func TestInvalidateAgent_MatchesAgentKeyNotUUID(t *testing.T) { + r := NewRouter() + agentKey := "my-agent" + agentUUID := uuid.New().String() + tenantID := uuid.New() + + // Simulate chat path: cache entry keyed by agentKey + ctx := store.WithTenantID(context.Background(), tenantID) + chatKey := agentCacheKey(ctx, agentKey) + r.agents[chatKey] = &agentEntry{} + + // Simulate old heartbeat path: cache entry keyed by UUID (the bug) + uuidKey := agentCacheKey(ctx, agentUUID) + r.agents[uuidKey] = &agentEntry{} + + // InvalidateAgent with agentKey should clear agentKey entry + r.InvalidateAgent(agentKey) + + if _, ok := r.agents[chatKey]; ok { + t.Error("InvalidateAgent(agentKey) should have cleared the agentKey-based cache entry") + } + + // UUID entry is NOT cleared — this is why the bug existed + if _, ok := r.agents[uuidKey]; !ok { + t.Error("InvalidateAgent(agentKey) should NOT match UUID-based entries (documenting the bug)") + } + + // InvalidateAgent with UUID clears the UUID entry (belt-and-suspenders fix) + r.InvalidateAgent(agentUUID) + if _, ok := r.agents[uuidKey]; ok { + t.Error("InvalidateAgent(UUID) should have cleared the UUID-based cache entry") + } +} + +// TestInvalidateAgent_TenantScoped verifies that InvalidateAgent clears entries +// across all tenants for the same agentKey (suffix match). +func TestInvalidateAgent_TenantScoped(t *testing.T) { + r := NewRouter() + agentKey := "default" + tenantA := uuid.New() + tenantB := uuid.New() + + ctxA := store.WithTenantID(context.Background(), tenantA) + ctxB := store.WithTenantID(context.Background(), tenantB) + + keyA := agentCacheKey(ctxA, agentKey) + keyB := agentCacheKey(ctxB, agentKey) + r.agents[keyA] = &agentEntry{} + r.agents[keyB] = &agentEntry{} + + r.InvalidateAgent(agentKey) + + if len(r.agents) != 0 { + t.Errorf("InvalidateAgent should clear all tenant-scoped entries, got %d remaining", len(r.agents)) + } +} diff --git a/internal/gateway/methods/agents_update.go b/internal/gateway/methods/agents_update.go index 1ce4f36b..62f88917 100644 --- a/internal/gateway/methods/agents_update.go +++ b/internal/gateway/methods/agents_update.go @@ -157,6 +157,9 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client } m.agents.InvalidateAgent(params.AgentID) + // Also invalidate by UUID — heartbeat/cron sessions cached under UUID key + // before the agentKey fix may still be in the router cache. + m.agents.InvalidateAgent(ag.ID.String()) } else { // --- Fallback: config.json --- spec, ok := m.cfg.Agents.List[params.AgentID] diff --git a/internal/heartbeat/ticker.go b/internal/heartbeat/ticker.go index e92d819c..885f7f06 100644 --- a/internal/heartbeat/ticker.go +++ b/internal/heartbeat/ticker.go @@ -178,7 +178,7 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) { } // [2] Queue-aware: skip if agent is busy (active runs in scheduler). - if t.sched != nil && t.sched.HasActiveSessionsForAgent(agentIDStr) { + if t.sched != nil && t.sched.HasActiveSessionsForAgent(agentKey) { t.logSkipped(ctx, hb, "queue_busy", agentKey) // Don't advance next_run_at — retry on next poll. return @@ -216,7 +216,8 @@ func (t *Ticker) runOne(ctx context.Context, hb store.AgentHeartbeat) { agentKey, checklistContent, ) - sessionKey := sessions.BuildHeartbeatSessionKey(agentIDStr, hb.IsolatedSession) + // Use agentKey (not UUID) — session keys must use agentKey for cache invalidation consistency. + sessionKey := sessions.BuildHeartbeatSessionKey(agentKey, hb.IsolatedSession) channel := "heartbeat" if hb.Channel != nil && *hb.Channel != "" { diff --git a/internal/sessions/key.go b/internal/sessions/key.go index 2f07853b..8633819f 100644 --- a/internal/sessions/key.go +++ b/internal/sessions/key.go @@ -2,7 +2,13 @@ // // Session keys follow the TS OpenClaw canonical format: // -// agent:{agentId}:{rest} +// agent:{agentKey}:{rest} +// +// IMPORTANT: {agentKey} is the human-readable agent key (e.g. "default", "my-agent"), +// NOT the agent's UUID. All Build*SessionKey functions expect agentKey, not UUID. +// This is critical for cache invalidation — the agent router cache uses agentKey +// as the lookup key, and InvalidateAgent() matches by agentKey suffix. +// Tenant isolation is handled separately via context (tenant_id in DB, prefix in cache). // // Where {rest} depends on the session type: // diff --git a/internal/sessions/key_test.go b/internal/sessions/key_test.go new file mode 100644 index 00000000..3d9eb8ee --- /dev/null +++ b/internal/sessions/key_test.go @@ -0,0 +1,87 @@ +package sessions + +import ( + "strings" + "testing" +) + +// TestBuildHeartbeatSessionKey_UsesAgentKey verifies that heartbeat session keys +// use agentKey (human-readable, e.g. "my-agent"), NOT the agent's UUID. +// This is critical: the agent router cache is invalidated by agentKey suffix match, +// so UUID-based keys would never be invalidated on agent config updates. +func TestBuildHeartbeatSessionKey_UsesAgentKey(t *testing.T) { + agentKey := "my-agent" + key := BuildHeartbeatSessionKey(agentKey, false) + + if !strings.HasPrefix(key, "agent:my-agent:") { + t.Errorf("expected session key to use agentKey, got %q", key) + } + if key != "agent:my-agent:heartbeat" { + t.Errorf("expected agent:my-agent:heartbeat, got %q", key) + } +} + +// TestBuildHeartbeatSessionKey_Isolated verifies isolated heartbeat sessions +// include a timestamp suffix. +func TestBuildHeartbeatSessionKey_Isolated(t *testing.T) { + key := BuildHeartbeatSessionKey("my-agent", true) + + if !strings.HasPrefix(key, "agent:my-agent:heartbeat:") { + t.Errorf("expected agent:my-agent:heartbeat:{timestamp}, got %q", key) + } +} + +// TestBuildCronSessionKey_UsesAgentKey verifies that cron session keys use agentKey. +// CronJob.AgentID is a UUID from DB — callers must resolve it to agentKey before calling. +func TestBuildCronSessionKey_UsesAgentKey(t *testing.T) { + agentKey := "my-agent" + key := BuildCronSessionKey(agentKey, "job-123") + + if key != "agent:my-agent:cron:job-123" { + t.Errorf("expected agent:my-agent:cron:job-123, got %q", key) + } +} + +// TestBuildSessionKey_UUIDIsWrong documents that passing a UUID instead of agentKey +// produces a session key that won't match cache invalidation patterns. +// This test exists as a guard: if someone passes a UUID, the key will contain +// it literally, which breaks InvalidateAgent(agentKey) suffix matching. +func TestBuildSessionKey_UUIDIsWrong(t *testing.T) { + uuid := "550e8400-e29b-41d4-a716-446655440000" + agentKey := "my-agent" + + uuidKey := BuildHeartbeatSessionKey(uuid, false) + correctKey := BuildHeartbeatSessionKey(agentKey, false) + + // UUID-based key would never be matched by InvalidateAgent("my-agent") + if strings.Contains(uuidKey, agentKey) { + t.Error("UUID-based key should NOT contain agentKey") + } + if !strings.Contains(correctKey, agentKey) { + t.Error("agentKey-based key MUST contain agentKey") + } +} + +// TestParseSessionKey_ExtractsAgentKey verifies ParseSessionKey returns the agentKey +// (the second segment), which should be a human-readable key, not a UUID. +func TestParseSessionKey_ExtractsAgentKey(t *testing.T) { + tests := []struct { + key string + wantAgent string + wantRest string + }{ + {"agent:my-agent:heartbeat", "my-agent", "heartbeat"}, + {"agent:my-agent:cron:job-1", "my-agent", "cron:job-1"}, + {"agent:default:ws:direct:conv-1", "default", "ws:direct:conv-1"}, + {"agent:my-agent:team:t1:c1", "my-agent", "team:t1:c1"}, + } + for _, tt := range tests { + agentID, rest := ParseSessionKey(tt.key) + if agentID != tt.wantAgent { + t.Errorf("ParseSessionKey(%q) agentID = %q, want %q", tt.key, agentID, tt.wantAgent) + } + if rest != tt.wantRest { + t.Errorf("ParseSessionKey(%q) rest = %q, want %q", tt.key, rest, tt.wantRest) + } + } +} diff --git a/internal/store/pg/sessions.go b/internal/store/pg/sessions.go index e9415ddd..7c7a3760 100644 --- a/internal/store/pg/sessions.go +++ b/internal/store/pg/sessions.go @@ -33,6 +33,7 @@ func NewPGSessionStore(db *sql.DB) *PGSessionStore { cache: make(map[string]*store.SessionData), } s.migrateLegacyWSKeys() + s.migrateUUIDSessionKeys() return s } @@ -59,6 +60,43 @@ func (s *PGSessionStore) migrateLegacyWSKeys() { } } +// migrateUUIDSessionKeys fixes legacy heartbeat/cron session keys that used the agent's +// UUID instead of agentKey. The old format "agent:{UUID}:heartbeat" or "agent:{UUID}:cron:..." +// is replaced with "agent:{agentKey}:..." by JOINing with the agents table. +// Idempotent — no-op if no UUID-based keys exist. +func (s *PGSessionStore) migrateUUIDSessionKeys() { + // UUID pattern: 8-4-4-4-12 hex chars. Matches session keys where the agent segment is a UUID. + // Rewrites to use agents.agent_key instead. + // Build the target key and skip rows where the target already exists (avoids unique constraint violation + // when both UUID-keyed and agentKey-keyed sessions coexist for the same agent). + res, err := s.db.ExecContext(context.Background(), ` + UPDATE sessions s + SET session_key = 'agent:' || a.agent_key || ':' || split_part(s.session_key, ':', 3) + || CASE WHEN array_length(string_to_array(s.session_key, ':'), 1) > 3 + THEN ':' || (SELECT string_agg(part, ':') FROM unnest((string_to_array(s.session_key, ':'))[4:]) AS part) + ELSE '' END + FROM agents a + WHERE s.session_key ~ '^agent:[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}:' + AND a.id = (split_part(s.session_key, ':', 2))::uuid + AND a.deleted_at IS NULL + AND NOT EXISTS ( + SELECT 1 FROM sessions s2 + WHERE s2.tenant_id = s.tenant_id + AND s2.session_key = 'agent:' || a.agent_key || ':' || split_part(s.session_key, ':', 3) + || CASE WHEN array_length(string_to_array(s.session_key, ':'), 1) > 3 + THEN ':' || (SELECT string_agg(p, ':') FROM unnest((string_to_array(s.session_key, ':'))[4:]) AS p) + ELSE '' END + ) + `) + if err != nil { + slog.Warn("sessions.migrate_uuid_keys", "error", err) + return + } + if n, _ := res.RowsAffected(); n > 0 { + slog.Info("sessions.migrate_uuid_keys", "migrated", n) + } +} + // sessionCacheKey prefixes session key with tenant UUID to prevent cross-tenant cache collisions. // Two tenants with the same agent_key produce different cache keys. func sessionCacheKey(ctx context.Context, key string) string {