Skip to content

fix: harden after-start prompt delivery #2520

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
nikhilachale wants to merge 1 commit into from
+170 −11
Closed

fix: harden after-start prompt delivery #2520

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 60 additions & 6 deletions backend/internal/session_manager/manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <agent> <event>`. The session
// PATH pin (hookPATH) only works when the daemon's own executable carries this
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
115 changes: 110 additions & 5 deletions backend/internal/session_manager/manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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)
Expand All @@ -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()}
Expand All @@ -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 {
Expand Down Expand Up @@ -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") {
Expand Down
Loading