From 03e3ec2be46dcff7f8758f865e2fdebcbb8e55f4 Mon Sep 17 00:00:00 2001 From: nikhil achale Date: Wed, 8 Jul 2026 20:29:31 +0530 Subject: [PATCH] fix: harden after-start prompt delivery --- backend/internal/session_manager/manager.go | 66 +++++++++- .../internal/session_manager/manager_test.go | 115 +++++++++++++++++- 2 files changed, 170 insertions(+), 11 deletions(-) diff --git a/backend/internal/session_manager/manager.go b/backend/internal/session_manager/manager.go index 4581ca771f..42df3824e7 100644 --- a/backend/internal/session_manager/manager.go +++ b/backend/internal/session_manager/manager.go @@ -58,6 +58,12 @@ const ( EnvDataDir = "AO_DATA_DIR" ) +const ( + afterStartPromptInitialDelay = 500 * time.Millisecond + afterStartPromptRetryDelay = 250 * time.Millisecond + afterStartPromptAttempts = 4 +) + // hookBinaryName is the executable name the workspace hook commands invoke: // every agent adapter installs a bare `ao hooks `. The session // PATH pin (hookPATH) only works when the daemon's own executable carries this @@ -129,6 +135,10 @@ type Manager struct { // workspace hook commands resolve back to this daemon. Tests inject a stub. executable func() (string, error) logger *slog.Logger + + afterStartPromptInitialDelay time.Duration + afterStartPromptRetryDelay time.Duration + afterStartPromptAttempts int } // Deps are the collaborators a Session Manager needs; New wires them together. @@ -171,6 +181,10 @@ func New(d Deps) *Manager { lookPath: d.LookPath, executable: d.Executable, logger: d.Logger, + + afterStartPromptInitialDelay: afterStartPromptInitialDelay, + afterStartPromptRetryDelay: afterStartPromptRetryDelay, + afterStartPromptAttempts: afterStartPromptAttempts, } if m.clock == nil { // UTC so spawn-stamped CreatedAt/UpdatedAt match every other session @@ -315,15 +329,55 @@ func (m *Manager) Spawn(ctx context.Context, cfg ports.SpawnConfig) (domain.Sess m.markSpawnFailedTerminated(ctx, id) return domain.SessionRecord{}, fmt.Errorf("spawn %s: completed: %w", id, err) } - if delivery == ports.PromptDeliveryAfterStart && prompt != "" { + if err := m.deliverAfterStartPrompt(ctx, id, prompt, delivery); err != nil { + _ = m.runtime.Destroy(ctx, handle) + m.destroySpawnWorkspace(ctx, ws, workspaceProject) + m.markSpawnFailedTerminatedWithoutWorkspace(ctx, id) + return domain.SessionRecord{}, fmt.Errorf("spawn %s: deliver prompt: %w", id, err) + } + return m.getRecord(ctx, id) +} + +func (m *Manager) deliverAfterStartPrompt(ctx context.Context, id domain.SessionID, prompt string, strategy ports.PromptDeliveryStrategy) error { + if strategy != ports.PromptDeliveryAfterStart || prompt == "" { + return nil + } + if err := waitContext(ctx, m.afterStartPromptInitialDelay); err != nil { + return err + } + attempts := m.afterStartPromptAttempts + if attempts <= 0 { + attempts = 1 + } + var lastErr error + for i := 0; i < attempts; i++ { if err := m.messenger.Send(ctx, id, prompt); err != nil { - _ = m.runtime.Destroy(ctx, handle) - m.destroySpawnWorkspace(ctx, ws, workspaceProject) - m.markSpawnFailedTerminatedWithoutWorkspace(ctx, id) - return domain.SessionRecord{}, fmt.Errorf("spawn %s: deliver prompt: %w", id, err) + lastErr = err + } else { + return nil + } + if i == attempts-1 { + break + } + if err := waitContext(ctx, m.afterStartPromptRetryDelay); err != nil { + return err } } - return m.getRecord(ctx, id) + return lastErr +} + +func waitContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } } // loadProject loads the project record so spawn can resolve its per-project diff --git a/backend/internal/session_manager/manager_test.go b/backend/internal/session_manager/manager_test.go index 61cecd88b7..2af847b2f4 100644 --- a/backend/internal/session_manager/manager_test.go +++ b/backend/internal/session_manager/manager_test.go @@ -424,11 +424,31 @@ func fakeWorkspaceRepoName(info ports.WorkspaceInfo) string { } type fakeMessenger struct { - msgs []string - err error -} - -func (m *fakeMessenger) Send(_ context.Context, _ domain.SessionID, msg string) error { + msgs []string + ids []domain.SessionID + err error + failErr error + failCount int + onSend func(domain.SessionID, string) error +} + +func (m *fakeMessenger) Send(_ context.Context, id domain.SessionID, msg string) error { + if m.onSend != nil { + if err := m.onSend(id, msg); err != nil { + return err + } + } + if m.failCount > 0 { + m.failCount-- + if m.failErr != nil { + return m.failErr + } + return m.err + } + if m.err != nil { + return m.err + } + m.ids = append(m.ids, id) m.msgs = append(m.msgs, msg) return m.err } @@ -576,6 +596,8 @@ func TestSpawn_DeliversPromptAfterStartWhenAgentRequestsIt(t *testing.T) { Lifecycle: &fakeLCM{store: st}, LookPath: func(string) (string, error) { return "/bin/true", nil }, }) + m.afterStartPromptInitialDelay = 0 + m.afterStartPromptRetryDelay = 0 if _, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "fix the button"}); err != nil { t.Fatal(err) @@ -591,6 +613,83 @@ func TestSpawn_DeliversPromptAfterStartWhenAgentRequestsIt(t *testing.T) { } } +func TestSpawn_InCommandPromptIsNotSentAfterStart(t *testing.T) { + st := newFakeStore() + st.projects["mer"] = domain.ProjectRecord{ID: "mer", Config: testRoleAgents()} + agent := &recordingAgent{} + msg := &fakeMessenger{} + m := New(Deps{ + Runtime: &fakeRuntime{}, Agents: singleAgent{agent: agent}, Workspace: &fakeWorkspace{}, Store: st, + Messenger: msg, Lifecycle: &fakeLCM{store: st}, LookPath: func(string) (string, error) { return "/bin/true", nil }, + }) + + if _, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "do it"}); err != nil { + t.Fatal(err) + } + if agent.lastLaunch.Prompt != "do it" { + t.Fatalf("launch prompt = %q, want in-command prompt", agent.lastLaunch.Prompt) + } + if len(msg.msgs) != 0 { + t.Fatalf("messenger sends = %v, want none for in-command delivery", msg.msgs) + } +} + +func TestSpawn_AfterStartPromptSentAfterRuntimeHandleIsRecorded(t *testing.T) { + st := newFakeStore() + st.projects["mer"] = domain.ProjectRecord{ID: "mer", Config: testRoleAgents()} + agent := &recordingAgent{} + rt := &fakeRuntime{} + ws := &fakeWorkspace{} + msg := &fakeMessenger{} + msg.onSend = func(id domain.SessionID, msg string) error { + rec := st.sessions[id] + if rec.Metadata.RuntimeHandleID == "" { + return errors.New("prompt sent before runtime handle landed") + } + return nil + } + m := New(Deps{ + Runtime: rt, Agents: singleAgent{agent: afterStartAgent{recordingAgent: agent}}, Workspace: ws, Store: st, + Messenger: msg, Lifecycle: &fakeLCM{store: st}, LookPath: func(string) (string, error) { return "/bin/true", nil }, + }) + m.afterStartPromptInitialDelay = 0 + m.afterStartPromptRetryDelay = 0 + + rec, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "do it"}) + if err != nil { + t.Fatal(err) + } + if agent.lastLaunch.Prompt != "" { + t.Fatalf("launch prompt = %q, want empty for after-start delivery", agent.lastLaunch.Prompt) + } + if len(msg.ids) != 1 || msg.ids[0] != rec.ID || len(msg.msgs) != 1 || msg.msgs[0] != "do it" { + t.Fatalf("sent ids/messages = %v/%v, want one prompt to %s", msg.ids, msg.msgs, rec.ID) + } + if rt.destroyed != 0 || ws.destroyed != 0 || st.sessions[rec.ID].IsTerminated { + t.Fatalf("successful after-start delivery should keep session live; destroyed runtime/workspace=%d/%d terminated=%v", rt.destroyed, ws.destroyed, st.sessions[rec.ID].IsTerminated) + } +} + +func TestSpawn_AfterStartPromptRetriesTransportFailures(t *testing.T) { + st := newFakeStore() + st.projects["mer"] = domain.ProjectRecord{ID: "mer", Config: testRoleAgents()} + agent := &recordingAgent{} + msg := &fakeMessenger{failErr: errors.New("pane not ready"), failCount: 2} + m := New(Deps{ + Runtime: &fakeRuntime{}, Agents: singleAgent{agent: afterStartAgent{recordingAgent: agent}}, Workspace: &fakeWorkspace{}, Store: st, + Messenger: msg, Lifecycle: &fakeLCM{store: st}, LookPath: func(string) (string, error) { return "/bin/true", nil }, + }) + m.afterStartPromptInitialDelay = 0 + m.afterStartPromptRetryDelay = 0 + + if _, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "do it"}); err != nil { + t.Fatal(err) + } + if len(msg.msgs) != 1 || msg.msgs[0] != "do it" { + t.Fatalf("messages after retries = %v, want one successful prompt", msg.msgs) + } +} + func TestSpawn_AfterStartPromptFailureCleansUpSpawn(t *testing.T) { st := newFakeStore() st.projects["mer"] = domain.ProjectRecord{ID: "mer", Config: testRoleAgents()} @@ -608,6 +707,9 @@ func TestSpawn_AfterStartPromptFailureCleansUpSpawn(t *testing.T) { Lifecycle: lcm, LookPath: func(string) (string, error) { return "/bin/true", nil }, }) + m.afterStartPromptInitialDelay = 0 + m.afterStartPromptRetryDelay = 0 + m.afterStartPromptAttempts = 3 _, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "fix the button"}) if err == nil { @@ -656,6 +758,9 @@ func TestSpawn_AfterStartPromptFailureCleansUpWorkspaceProjectRows(t *testing.T) Lifecycle: lcm, LookPath: func(string) (string, error) { return "/bin/true", nil }, }) + m.afterStartPromptInitialDelay = 0 + m.afterStartPromptRetryDelay = 0 + m.afterStartPromptAttempts = 3 _, err := m.Spawn(ctx, ports.SpawnConfig{ProjectID: "mer", Kind: domain.KindWorker, Prompt: "fix the button"}) if err == nil || !strings.Contains(err.Error(), "deliver prompt") {