From 8948ebdb255d1f4fd51ea7186aaaeeb5f2d8952e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Salavisa?= Date: Sun, 26 Apr 2026 21:07:22 +0100 Subject: [PATCH] fix(acp): flush ScheduleWakeup output via synthetic prompt (#706) --- .../server/adapter/transport/acp/adapter.go | 159 ++++++- .../server/adapter/transport/acp/wakeup.go | 197 ++++++++ .../adapter/transport/acp/wakeup_test.go | 446 ++++++++++++++++++ 3 files changed, 801 insertions(+), 1 deletion(-) create mode 100644 apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup.go create mode 100644 apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup_test.go diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go index 824ca51b4c..0cee791d01 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go @@ -11,6 +11,7 @@ import ( "log/slog" "strings" "sync" + "time" "github.com/coder/acp-go-sdk" acpclient "github.com/kandev/kandev/internal/agentctl/server/acp" @@ -40,6 +41,12 @@ const ( contentTypeResource = "resource" ) +// wakeupPromptTimeout bounds how long a synthetic wakeup prompt can run. +// Wakeup turns can perform real work (the model often runs a few tool calls +// before stopping) so we mirror what a normal user-initiated prompt would +// allow rather than a tight RPC deadline. +const wakeupPromptTimeout = 30 * time.Minute + // AgentInfo contains information about the connected agent. type AgentInfo struct { Name string `json:"name"` @@ -104,6 +111,16 @@ type Adapter struct { // and rebuilt during session/load replay. activeMonitors map[string]map[string]string + // ScheduleWakeup tracking. The Claude Agent SDK's ScheduleWakeup tool fires + // its timer inside the SDK's async-iterator, but the upstream + // @agentclientprotocol/claude-agent-acp bridge only drains that iterator + // inside its prompt() handler — so a wakeup that fires while no prompt is + // in flight produces no output. wakeup re-injects the wakeup as a synthetic + // session/prompt at fire time. pendingWakeups tracks per-tool-call info + // (prompt + scheduledFor) since these arrive in separate notifications. + wakeup *wakeupScheduler + pendingWakeups map[string]*pendingWakeup + // OTel tracing: active prompt span context. // Notification spans become children of the prompt span for visual grouping. promptTraceCtx context.Context @@ -124,6 +141,13 @@ type Adapter struct { // Synchronization mu sync.RWMutex closed bool + + // lifetimeCtx is cancelled by Close. Background work that may outlive + // the call site (e.g. the synthetic wakeup prompt goroutine) derives its + // context from this one so it aborts when the adapter shuts down rather + // than continuing to drive a dead subprocess. + lifetimeCtx context.Context + lifetimeCancel context.CancelFunc } // NewAdapter creates a new ACP protocol adapter. @@ -131,7 +155,8 @@ type Adapter struct { // cfg.AgentID is required for debug file naming. func NewAdapter(cfg *shared.Config, log *logger.Logger) *Adapter { l := log.WithFields(zap.String("adapter", "acp"), zap.String("agent_id", cfg.AgentID)) - return &Adapter{ + ctx, cancel := context.WithCancel(context.Background()) + a := &Adapter{ cfg: cfg, logger: l, agentID: cfg.AgentID, @@ -139,8 +164,13 @@ func NewAdapter(cfg *shared.Config, log *logger.Logger) *Adapter { updatesCh: make(chan AgentEvent, 100), activeToolCalls: make(map[string]*streams.NormalizedPayload), activeMonitors: make(map[string]map[string]string), + pendingWakeups: make(map[string]*pendingWakeup), attachMgr: shared.NewAttachmentManager(cfg.WorkDir, l.Zap()), + lifetimeCtx: ctx, + lifetimeCancel: cancel, } + a.wakeup = newWakeupScheduler(l, a.fireWakeup) + return a } // PrepareEnvironment is a no-op for ACP. @@ -252,6 +282,15 @@ func (a *Adapter) NewSession(ctx context.Context, mcpServers []types.McpServer) return "", fmt.Errorf("adapter not initialized") } + // A fresh session invalidates any pending wakeup keyed to the prior + // session. Reset pendingWakeups and cancel the scheduler under one + // a.mu critical section so a concurrent handleWakeupEvent can't slip + // a stale entry between the two operations. + a.mu.Lock() + a.pendingWakeups = make(map[string]*pendingWakeup) + a.wakeup.cancel() + a.mu.Unlock() + ctx, span := shared.TraceProtocolRequest(ctx, shared.ProtocolACP, a.agentID, "session.new") defer span.End() @@ -423,6 +462,15 @@ func (a *Adapter) LoadSession(ctx context.Context, sessionID string, mcpServers return fmt.Errorf("agent does not support session loading (LoadSession capability is false)") } + // Loading a different session invalidates any pending wakeup keyed to the + // prior session — same reset block as NewSession to avoid leaving an armed + // timer for a session id that's about to change and accumulating stale + // pendingWakeups entries across reloads. + a.mu.Lock() + a.pendingWakeups = make(map[string]*pendingWakeup) + a.wakeup.cancel() + a.mu.Unlock() + ctx, span := shared.TraceProtocolRequest(ctx, shared.ProtocolACP, a.agentID, "session.load") defer span.End() @@ -653,6 +701,96 @@ func (a *Adapter) Prompt(ctx context.Context, message string, attachments []v1.M return nil } +// fireWakeup is invoked by wakeupScheduler when a ScheduleWakeup timer +// elapses. It issues a synthetic session/prompt so the upstream +// @agentclientprotocol/claude-agent-acp bridge drains the SDK's queued wakeup +// turn and emits visible ACP frames. The session must still match (the user +// hasn't started a fresh session) and the adapter must not be closed. +// +// Concurrent-prompt safety: if a user prompt is already in flight when this +// runs, both end up calling conn.Prompt() on the same ClientSideConnection. +// That's safe at the wire level — the ACP SDK's Connection.sendMessage +// holds a write mutex, so request frames never interleave on stdin, and +// JSON-RPC pairs each response back to its originating request via id — +// but it does mean two prompts can be in flight against the bridge at +// once. The bridge serialises them in the order it receives them, which +// is exactly what we want for a wakeup that races a user message. +func (a *Adapter) fireWakeup(sessionID, prompt string) { + a.mu.RLock() + closed := a.closed + currentSession := a.sessionID + a.mu.RUnlock() + + if closed { + a.logger.Debug("skipping wakeup fire: adapter closed", + zap.String("session_id", sessionID)) + return + } + if currentSession != sessionID { + a.logger.Info("skipping wakeup fire: session changed", + zap.String("scheduled_for", sessionID), + zap.String("current", currentSession)) + return + } + + a.logger.Info("injecting synthetic wakeup prompt", + zap.String("session_id", sessionID), + zap.Int("prompt_len", len(prompt))) + + go func() { + // Derive from lifetimeCtx so a concurrent Close aborts the in-flight + // prompt instead of letting it run against a dead subprocess. + ctx, cancel := context.WithTimeout(a.lifetimeCtx, wakeupPromptTimeout) + defer cancel() + if err := a.Prompt(ctx, prompt, nil); err != nil { + a.logger.Error("synthetic wakeup prompt failed", + zap.String("session_id", sessionID), + zap.Error(err)) + } + }() +} + +// handleWakeupEvent inspects a tool-call meta + rawInput pair, accumulates +// pending state per toolCallID, and schedules a wakeup once both the prompt +// and scheduledFor timestamp are known. terminal=true means the tool call has +// reached a terminal state, so any pending entry should be cleaned up. +func (a *Adapter) handleWakeupEvent(sessionID, toolCallID string, meta any, rawInput any, terminal bool) { + if toolCallID == "" { + return + } + + scheduledForMs, isWakeup := extractScheduleWakeup(meta) + + a.mu.Lock() + pw, tracked := a.pendingWakeups[toolCallID] + if !tracked { + if !isWakeup { + a.mu.Unlock() + return + } + pw = &pendingWakeup{} + a.pendingWakeups[toolCallID] = pw + } + + if scheduledForMs > 0 { + pw.scheduledForMs = scheduledForMs + } + if prompt, ok := extractWakeupPrompt(rawInput); ok { + pw.prompt = prompt + } + + prompt := pw.prompt + stamp := pw.scheduledForMs + if (prompt != "" && stamp > 0) || terminal { + delete(a.pendingWakeups, toolCallID) + } + a.mu.Unlock() + + if prompt != "" && stamp > 0 { + a.wakeup.schedule(sessionID, prompt, stamp) + } +} + // SetPendingContext sets the context to be injected into the next prompt. // This is used by the fork_session pattern for ACP agents that don't support session/load. // The context will be prepended to the first prompt sent to this session. @@ -785,6 +923,15 @@ func (a *Adapter) Close() error { a.logger.Info("closing ACP adapter") + // Stop any pending ScheduleWakeup timer so it doesn't fire after close, + // and cancel the lifetime context so any in-flight wakeup prompt aborts. + if a.wakeup != nil { + a.wakeup.cancel() + } + if a.lifetimeCancel != nil { + a.lifetimeCancel() + } + // Clean up any saved attachments a.attachMgr.Cleanup() @@ -1432,6 +1579,11 @@ func (a *Adapter) convertToolCallUpdate(sessionID string, tc *acp.SessionUpdateT a.activeToolCalls[toolCallID] = normalizedPayload a.mu.Unlock() + // ScheduleWakeup tracking: meta carries `_meta.claudeCode.toolName` + // on the initial tool_call; rawInput is usually empty here but record + // the prompt eagerly when it does arrive in the same notification. + a.handleWakeupEvent(sessionID, toolCallID, tc.Meta, tc.RawInput, false) + // Detect tool type for logging toolType := DetectToolOperationType(toolKind, args) _ = toolType // Used for normalization @@ -1557,6 +1709,11 @@ func (a *Adapter) convertToolCallResultUpdate(sessionID string, tcu *acp.Session } a.mu.Unlock() + // ScheduleWakeup tracking: tool_call_update is where rawInput.prompt and + // `_meta.claudeCode.toolResponse.scheduledFor` typically arrive. Once both + // are known, schedule the synthetic prompt; on terminal status, clean up. + a.handleWakeupEvent(sessionID, toolCallID, tcu.Meta, tcu.RawInput, isTerminal) + // When a switch_mode tool carries a plan (e.g. ExitPlanMode), emit it // as an agent_plan event so the orchestrator creates a visible plan message. if tcu.RawInput != nil { diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup.go new file mode 100644 index 0000000000..39b43a8067 --- /dev/null +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup.go @@ -0,0 +1,197 @@ +package acp + +import ( + "sync" + "time" + + "github.com/kandev/kandev/internal/common/logger" + "go.uber.org/zap" +) + +// wakeupScheduler holds at most one pending ScheduleWakeup-driven prompt for +// the adapter. The Claude Agent SDK exposes `ScheduleWakeup` as a tool inside +// its harness; when its timer fires, the SDK queues a new turn on its +// async-iterator (`session.query`). The upstream +// @agentclientprotocol/claude-agent-acp bridge only iterates that channel +// inside its `prompt()` handler, so a wakeup that fires while no +// `session/prompt` is in flight produces no ACP output at all — the buffered +// turn only drains on the next user message. +// +// This scheduler closes the gap: when we observe a `ScheduleWakeup` tool +// completion, we record `scheduledFor` and the prompt text, and at fire time +// issue a synthetic `session/prompt` to the bridge. That triggers the bridge +// to drain the buffered output and produce visible ACP frames for the wakeup +// turn — matching what standalone Claude Code does, where the wakeup prompt +// is injected as a synthesized user turn. +type wakeupScheduler struct { + logger *logger.Logger + fire func(sessionID, prompt string) + + mu sync.Mutex + timer *time.Timer + sessionID string + prompt string + // gen is incremented on every schedule/cancel. fireOnce captures the gen + // at scheduling time and refuses to fire if it doesn't match the current + // gen — that closes the race where time.AfterFunc has already launched a + // goroutine for the old timer (timer.Stop() returns false) and the new + // schedule() then writes new state before the stale fireOnce acquires the + // lock. Without gen, the stale fireOnce would consume the new state and + // fire the new wakeup immediately while clobbering w.timer. + gen uint64 +} + +// pendingWakeup accumulates partial state for an in-flight ScheduleWakeup tool +// call. The bridge emits the wakeup metadata across multiple notifications: +// the initial tool_call carries `_meta.claudeCode.toolName="ScheduleWakeup"` +// (no rawInput yet); subsequent tool_call_updates fill in `rawInput.prompt` +// and `_meta.claudeCode.toolResponse.scheduledFor`. We only have enough to +// schedule once both the prompt and scheduledFor timestamp are present. +type pendingWakeup struct { + prompt string + scheduledForMs int64 +} + +func newWakeupScheduler(log *logger.Logger, fire func(sessionID, prompt string)) *wakeupScheduler { + return &wakeupScheduler{ + logger: log.WithFields(zap.String("component", "wakeup-scheduler")), + fire: fire, + } +} + +// schedule registers a wakeup for the given session, replacing any prior +// pending wakeup. unixMs is the Unix-epoch millisecond timestamp at which the +// wakeup should fire (matches the `_meta.claudeCode.toolResponse.scheduledFor` +// field reported by the bridge). If the timestamp is in the past or zero, the +// wakeup is dropped without firing — that matches the ACP semantics of "the +// turn already happened" and avoids unbounded immediate-fire loops on stale +// timestamps. +func (w *wakeupScheduler) schedule(sessionID, prompt string, unixMs int64) { + if sessionID == "" || prompt == "" || unixMs <= 0 { + w.logger.Warn("ignoring wakeup with missing fields", + zap.String("session_id", sessionID), + zap.Int("prompt_len", len(prompt)), + zap.Int64("unix_ms", unixMs)) + return + } + + delay := time.Until(time.UnixMilli(unixMs)) + if delay <= 0 { + // Warn rather than Debug: a stale wakeup typically signals system + // suspension, container pause, or significant clock skew — those are + // production-relevant and shouldn't be hidden at default log levels. + w.logger.Warn("dropping stale wakeup", + zap.String("session_id", sessionID), + zap.Duration("past_by", -delay)) + return + } + + w.mu.Lock() + if w.timer != nil { + w.timer.Stop() + } + w.gen++ + myGen := w.gen + w.sessionID = sessionID + w.prompt = prompt + w.timer = time.AfterFunc(delay, func() { w.fireOnce(myGen) }) + w.mu.Unlock() + + w.logger.Info("scheduled wakeup", + zap.String("session_id", sessionID), + zap.Time("fires_at", time.UnixMilli(unixMs)), + zap.Duration("delay", delay)) +} + +// fireOnce is invoked by the timer and dispatches to the configured callback. +// myGen identifies the schedule call that armed this timer; if a later +// schedule or cancel has bumped w.gen, this fire is stale and is dropped. +// It clears the pending state before calling the callback so a re-entry from +// inside the callback (e.g. the synthetic prompt schedules another wakeup) +// can overwrite cleanly. The gen check guarantees sessionID is non-empty +// here — schedule() rejects empty sessionID before bumping gen, and any +// cancel would have bumped gen too — so we don't re-check it. +func (w *wakeupScheduler) fireOnce(myGen uint64) { + w.mu.Lock() + if myGen != w.gen { + w.mu.Unlock() + return + } + sessionID := w.sessionID + prompt := w.prompt + w.timer = nil + w.sessionID = "" + w.prompt = "" + w.gen++ + w.mu.Unlock() + + w.logger.Info("firing wakeup", + zap.String("session_id", sessionID)) + w.fire(sessionID, prompt) +} + +// cancel stops any pending wakeup. Safe to call multiple times. Bumping gen +// invalidates any in-flight fireOnce that has already been launched but is +// blocked on the lock — without that, a stale fireOnce after Stop()=false +// would still observe the cleared state and no-op cleanly, but cancel races +// with concurrent schedule could otherwise let a stale fire consume newly +// scheduled state. +func (w *wakeupScheduler) cancel() { + w.mu.Lock() + defer w.mu.Unlock() + if w.timer != nil { + w.timer.Stop() + w.timer = nil + } + w.sessionID = "" + w.prompt = "" + w.gen++ +} + +// extractScheduleWakeup inspects an ACP tool-call meta map and returns +// (scheduledForUnixMs, isScheduleWakeup). It expects the bridge's +// `_meta.claudeCode.toolName` and `_meta.claudeCode.toolResponse.scheduledFor` +// shape that @agentclientprotocol/claude-agent-acp emits; missing/malformed +// entries are reported as not-a-wakeup rather than as errors so unrelated +// tools pass through unaffected. +func extractScheduleWakeup(meta any) (scheduledForMs int64, isWakeup bool) { + m, ok := meta.(map[string]any) + if !ok { + return 0, false + } + cc, ok := m["claudeCode"].(map[string]any) + if !ok { + return 0, false + } + if name, _ := cc["toolName"].(string); name != "ScheduleWakeup" { + return 0, false + } + resp, ok := cc["toolResponse"].(map[string]any) + if !ok { + return 0, true // ScheduleWakeup tool, but response not yet present + } + switch v := resp["scheduledFor"].(type) { + case float64: + return int64(v), true + case int64: + return v, true + case int: + return int64(v), true + } + return 0, true +} + +// extractWakeupPrompt pulls the prompt string out of a ScheduleWakeup +// tool-call rawInput. Returns ("", false) when the field is absent or the +// wrong type. +func extractWakeupPrompt(rawInput any) (string, bool) { + m, ok := rawInput.(map[string]any) + if !ok { + return "", false + } + s, ok := m["prompt"].(string) + if !ok || s == "" { + return "", false + } + return s, true +} diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup_test.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup_test.go new file mode 100644 index 0000000000..deb2bd4220 --- /dev/null +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/wakeup_test.go @@ -0,0 +1,446 @@ +package acp + +import ( + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/kandev/kandev/internal/common/logger" +) + +func newTestWakeupLogger(t *testing.T) *logger.Logger { + t.Helper() + log, err := logger.NewLogger(logger.LoggingConfig{ + Level: "error", + Format: "json", + OutputPath: "stderr", + }) + if err != nil { + t.Fatalf("logger init failed: %v", err) + } + return log +} + +func TestExtractScheduleWakeup_NotAWakeup(t *testing.T) { + cases := []struct { + name string + meta any + }{ + {"nil", nil}, + {"non-map", "string"}, + {"empty map", map[string]any{}}, + {"missing claudeCode", map[string]any{"other": "x"}}, + {"claudeCode not a map", map[string]any{"claudeCode": "x"}}, + {"different tool", map[string]any{ + "claudeCode": map[string]any{"toolName": "Bash"}, + }}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + ms, ok := extractScheduleWakeup(c.meta) + if ok { + t.Errorf("expected isWakeup=false for %s, got true (ms=%d)", c.name, ms) + } + }) + } +} + +func TestExtractScheduleWakeup_NoResponseYet(t *testing.T) { + meta := map[string]any{ + "claudeCode": map[string]any{ + "toolName": "ScheduleWakeup", + }, + } + ms, ok := extractScheduleWakeup(meta) + if !ok { + t.Fatal("expected isWakeup=true") + } + if ms != 0 { + t.Errorf("expected ms=0 when toolResponse missing, got %d", ms) + } +} + +func TestExtractScheduleWakeup_WithResponse(t *testing.T) { + cases := []struct { + name string + val any + want int64 + }{ + {"float64", float64(1700000000000), 1700000000000}, + {"int", int(1700000000000), 1700000000000}, + {"int64", int64(1700000000000), 1700000000000}, + {"unsupported type", "1700000000000", 0}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + meta := map[string]any{ + "claudeCode": map[string]any{ + "toolName": "ScheduleWakeup", + "toolResponse": map[string]any{ + "scheduledFor": c.val, + }, + }, + } + ms, ok := extractScheduleWakeup(meta) + if !ok { + t.Fatal("expected isWakeup=true") + } + if ms != c.want { + t.Errorf("ms=%d, want %d", ms, c.want) + } + }) + } +} + +func TestExtractWakeupPrompt(t *testing.T) { + cases := []struct { + name string + input any + wantOK bool + wantText string + }{ + {"nil input", nil, false, ""}, + {"non-map", "x", false, ""}, + {"empty map", map[string]any{}, false, ""}, + {"empty prompt", map[string]any{"prompt": ""}, false, ""}, + {"non-string prompt", map[string]any{"prompt": 42}, false, ""}, + {"valid prompt", map[string]any{"prompt": "hello"}, true, "hello"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got, ok := extractWakeupPrompt(c.input) + if ok != c.wantOK || got != c.wantText { + t.Errorf("got (%q, %v), want (%q, %v)", got, ok, c.wantText, c.wantOK) + } + }) + } +} + +type wakeupCapture struct { + mu sync.Mutex + calls int32 + lastSess string + lastPrmpt string +} + +func newWakeupCapture() *wakeupCapture { + return &wakeupCapture{} +} + +func (c *wakeupCapture) fire(sessionID, prompt string) { + c.mu.Lock() + c.lastSess = sessionID + c.lastPrmpt = prompt + c.mu.Unlock() + atomic.AddInt32(&c.calls, 1) +} + +// Timer-driven tests use testing/synctest so time.AfterFunc is intercepted by +// synctest's fake clock. To advance fake time past a scheduled timer we +// time.Sleep in the test goroutine — under synctest, all bubble goroutines +// are then blocked and the runtime jumps fake time to the next pending event. +// synctest.Wait() afterwards ensures any goroutine the timer spawned has +// finished its work before the assertion. + +func TestWakeupScheduler_FiresAfterDelay(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + at := time.Now().Add(40 * time.Millisecond).UnixMilli() + sched.schedule("sess-1", "wake-prompt", at) + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 1 { + t.Fatalf("expected 1 fire, got %d", got) + } + cap.mu.Lock() + defer cap.mu.Unlock() + if cap.lastSess != "sess-1" || cap.lastPrmpt != "wake-prompt" { + t.Errorf("unexpected fire args: sess=%q prompt=%q", cap.lastSess, cap.lastPrmpt) + } + }) +} + +func TestWakeupScheduler_CancelPreventsFire(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + at := time.Now().Add(50 * time.Millisecond).UnixMilli() + sched.schedule("sess-1", "p", at) + sched.cancel() + + // Advance fake time past the would-be fire instant; nothing should run. + time.Sleep(200 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("expected 0 fires, got %d", got) + } + }) +} + +func TestWakeupScheduler_StaleTimestampDropped(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + // Past timestamp: should not arm a timer at all. + sched.schedule("sess-1", "p", time.Now().Add(-time.Hour).UnixMilli()) + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("expected 0 fires, got %d", got) + } + }) +} + +func TestWakeupScheduler_RescheduleReplacesPrior(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + // First schedule far in the future, then reschedule soon — only the + // second should fire, and only once. + sched.schedule("sess-1", "first", time.Now().Add(time.Hour).UnixMilli()) + sched.schedule("sess-1", "second", time.Now().Add(40*time.Millisecond).UnixMilli()) + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 1 { + t.Fatalf("expected 1 fire, got %d", got) + } + cap.mu.Lock() + defer cap.mu.Unlock() + if cap.lastPrmpt != "second" { + t.Errorf("expected last prompt %q, got %q", "second", cap.lastPrmpt) + } + }) +} + +// TestWakeupScheduler_StaleFireOnceDoesNotConsumeNewState reproduces the race +// where time.AfterFunc has already launched a fireOnce goroutine for the old +// timer before schedule() runs. The stale fireOnce must observe the gen +// mismatch and bail out, not consume the newly scheduled wakeup. +func TestWakeupScheduler_StaleFireOnceDoesNotConsumeNewState(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + // Manually drive the race: arm a timer, capture its gen, then simulate + // a stale fire by calling fireOnce with the captured gen *after* a + // rescheduling has bumped gen. + sched.mu.Lock() + sched.gen++ + staleGen := sched.gen + sched.sessionID = "old-sess" + sched.prompt = "old-prompt" + sched.mu.Unlock() + + // Reschedule before stale fire executes. + sched.schedule("new-sess", "new-prompt", time.Now().Add(time.Hour).UnixMilli()) + + // Stale fire arrives now — should be a no-op. + sched.fireOnce(staleGen) + + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("stale fireOnce fired anyway: calls=%d", got) + } + + // New schedule should still fire normally when its timer elapses. + sched.cancel() + sched.schedule("new-sess", "new-prompt", time.Now().Add(40*time.Millisecond).UnixMilli()) + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 1 { + t.Fatalf("expected 1 fire after reschedule, got %d", got) + } + cap.mu.Lock() + defer cap.mu.Unlock() + if cap.lastSess != "new-sess" || cap.lastPrmpt != "new-prompt" { + t.Errorf("fired with wrong args: sess=%q prompt=%q", cap.lastSess, cap.lastPrmpt) + } + }) +} + +func TestWakeupScheduler_IgnoresMissingFields(t *testing.T) { + log := newTestWakeupLogger(t) + cap := newWakeupCapture() + sched := newWakeupScheduler(log, cap.fire) + + // Missing each required field in turn — none should arm a timer. + sched.schedule("", "p", time.Now().Add(time.Minute).UnixMilli()) + sched.schedule("sess", "", time.Now().Add(time.Minute).UnixMilli()) + sched.schedule("sess", "p", 0) + + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("expected 0 fires, got %d", got) + } +} + +func TestHandleWakeupEvent_SchedulesOnceBothFieldsArrive(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + a := newTestAdapter() + t.Cleanup(func() { _ = a.Close() }) + + cap := newWakeupCapture() + a.wakeup = newWakeupScheduler(a.logger, cap.fire) + + at := time.Now().Add(40 * time.Millisecond).UnixMilli() + + // Initial tool_call: only the toolName is known. + metaInitial := map[string]any{ + "claudeCode": map[string]any{"toolName": "ScheduleWakeup"}, + } + a.handleWakeupEvent("sess-1", "tc-1", metaInitial, nil, false) + + // Mid-update: rawInput arrives with the prompt. + a.handleWakeupEvent("sess-1", "tc-1", metaInitial, map[string]any{ + "prompt": "wake-now", + "delaySeconds": 60, + }, false) + + // Verify nothing fired yet — scheduledFor still missing. + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Fatalf("fired before scheduledFor known: calls=%d", got) + } + + // Final update: scheduledFor arrives in meta. + metaWithResp := map[string]any{ + "claudeCode": map[string]any{ + "toolName": "ScheduleWakeup", + "toolResponse": map[string]any{"scheduledFor": float64(at)}, + }, + } + a.handleWakeupEvent("sess-1", "tc-1", metaWithResp, nil, false) + + time.Sleep(100 * time.Millisecond) + synctest.Wait() + + if got := atomic.LoadInt32(&cap.calls); got != 1 { + t.Fatalf("expected 1 fire, got %d", got) + } + cap.mu.Lock() + defer cap.mu.Unlock() + if cap.lastPrmpt != "wake-now" { + t.Errorf("prompt=%q, want %q", cap.lastPrmpt, "wake-now") + } + }) +} + +func TestHandleWakeupEvent_NonWakeupIgnored(t *testing.T) { + a := newTestAdapter() + t.Cleanup(func() { _ = a.Close() }) + + cap := newWakeupCapture() + a.wakeup = newWakeupScheduler(a.logger, cap.fire) + + // Bash tool — should be ignored entirely. + meta := map[string]any{ + "claudeCode": map[string]any{ + "toolName": "Bash", + "toolResponse": map[string]any{"scheduledFor": float64(time.Now().Add(time.Hour).UnixMilli())}, + }, + } + a.handleWakeupEvent("sess-1", "tc-1", meta, map[string]any{"prompt": "x"}, false) + + a.mu.Lock() + if _, present := a.pendingWakeups["tc-1"]; present { + t.Error("non-wakeup tool should not be tracked") + } + a.mu.Unlock() + + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("expected 0 fires, got %d", got) + } +} + +func TestHandleWakeupEvent_TerminalCleansUpPending(t *testing.T) { + a := newTestAdapter() + t.Cleanup(func() { _ = a.Close() }) + + cap := newWakeupCapture() + a.wakeup = newWakeupScheduler(a.logger, cap.fire) + + meta := map[string]any{ + "claudeCode": map[string]any{"toolName": "ScheduleWakeup"}, + } + a.handleWakeupEvent("sess-1", "tc-1", meta, nil, false) + + a.mu.Lock() + _, present := a.pendingWakeups["tc-1"] + a.mu.Unlock() + if !present { + t.Fatal("expected pending wakeup before terminal") + } + + // Terminal arrives without prompt or scheduledFor — should clean up. + a.handleWakeupEvent("sess-1", "tc-1", meta, nil, true) + + a.mu.Lock() + _, present = a.pendingWakeups["tc-1"] + a.mu.Unlock() + if present { + t.Error("expected pending wakeup cleared on terminal") + } + if got := atomic.LoadInt32(&cap.calls); got != 0 { + t.Errorf("terminal-without-data should not fire, got %d calls", got) + } +} + +// TestFireWakeup_SkipsOnSessionChange exercises fireWakeup's session-guard +// branch — when the adapter's current session has rotated away from the one +// the wakeup was scheduled for, the wakeup must drop instead of trying to +// drive an unrelated session. +// TestFireWakeup_SkipsOnSessionChange exercises fireWakeup's session-guard +// branch — when the adapter's current session has rotated away from the one +// the wakeup was scheduled for, the wakeup must drop instead of trying to +// drive an unrelated session. +// +// We can't observe Prompt directly (it returns early on nil acpConn without +// any side effect on the updates channel) so we use a.pendingContext as a +// canary: Prompt's first action under a.mu is to read-and-clear it. If the +// guard works, fireWakeup never spawns the goroutine that calls Prompt and +// the canary survives. If the guard regresses, Prompt runs and clears it. +func TestFireWakeup_SkipsOnSessionChange(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + a := newTestAdapter() + t.Cleanup(func() { _ = a.Close() }) + + // Adapter currently tracks a different session than the one the wakeup + // was scheduled against, and pendingContext holds a canary string. + const canary = "guard-canary" + a.mu.Lock() + a.sessionID = "new-sess" + a.pendingContext = canary + a.mu.Unlock() + + a.fireWakeup("old-sess", "should-not-fire") + + // Let any spawned goroutine run to completion. Inside synctest, after + // Wait() returns the bubble has settled and we can inspect state. + synctest.Wait() + + a.mu.Lock() + got := a.pendingContext + a.mu.Unlock() + if got != canary { + t.Errorf("pendingContext=%q, want %q — Prompt was reached, guard regressed", got, canary) + } + }) +}