ai/byre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 9 Wiki Activity

Wait for the agent's turn instead of dropping a held merge

Browse Source
This commit is contained in:
Ian Gulliver
2026-10-21 11:03:25 -07:00
parent 39fe50eff2
commit 4ba742aee0
4 changed files with 129 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/store/discard.go
+1 -1
View File
@@ -184,7 +184,7 @@ func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload ma
case mergeAfterTurn && (othersRunning || continues):
end.MergeWaiting = true
case mergeAfterTurn:
err = tx.QueryRowContext(ctx, `UPDATE tasks SET merge_after_turn = 0 WHERE id = ? RETURNING merging_by`, taskID).Scan(&end.MergeBy)
err = tx.QueryRowContext(ctx, `SELECT merging_by FROM tasks WHERE id = ?`, taskID).Scan(&end.MergeBy)
case state == DiscardNone && finalizeMerge && lead:
err = tx.QueryRowContext(ctx, `UPDATE tasks SET finalizing = merging_by, merging_by = '' WHERE id = ? AND merging_by != '' RETURNING finalizing`, taskID).Scan(&end.FinalizeBy)
if errors.Is(err, sql.ErrNoRows) {
internal/store/store.go
+87 -58
View File
@@ -712,13 +712,11 @@ func (s *Store) TakeQueuedMerge(ctx context.Context, id string) (string, error)
return login, err
}
func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, bool, error) {
func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, error) {
login := ""
claimed := false
err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
projectID := ""
err := tx.QueryRowContext(ctx, `SELECT project_id, merge_queued_by, EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?) FROM tasks WHERE id = ?`,
string(api.Implementer), string(api.AgentSessionStateRunning), id).Scan(&projectID, &login, &claimed)
err := tx.QueryRowContext(ctx, `SELECT project_id, merge_queued_by FROM tasks WHERE id = ?`, id).Scan(&projectID, &login)
if errors.Is(err, sql.ErrNoRows) {
return ErrNotFound
}
@@ -733,32 +731,59 @@ func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, bool,
return err
}
err = emit(api.Event{Kind: "task.merge_unqueued", ProjectId: &projectID, TaskId: &id, Actor: "byre", Payload: map[string]any{"id": id}})
if err != nil || !claimed {
return err
}
_, err = tx.ExecContext(ctx, `UPDATE tasks SET merging_by = ?, merge_after_turn = 1 WHERE id = ?`, login, id)
if err != nil {
return err
}
return emit(api.Event{Kind: "task.merge_started", ProjectId: &projectID, TaskId: &id, Actor: login, Payload: map[string]any{"id": id}})
})
return login, claimed, err
}
func (s *Store) TakeIdleMergeClaim(ctx context.Context, id string) (string, error) {
login := ""
err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
err := tx.QueryRowContext(ctx, `UPDATE tasks SET merge_after_turn = 0 WHERE id = ? AND merge_after_turn = 1
AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?) RETURNING merging_by`,
id, string(api.Implementer), string(api.AgentSessionStateRunning)).Scan(&login)
if errors.Is(err, sql.ErrNoRows) {
res, err := tx.ExecContext(ctx, `UPDATE tasks SET merging_by = ?, merge_after_turn = 1 WHERE id = ? AND discard_state = '' AND finalizing = ''`, login, id)
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil {
return err
}
if n == 0 {
login = ""
return nil
}
return err
return emit(api.Event{Kind: "task.merge_started", ProjectId: &projectID, TaskId: &id, Actor: login, Payload: map[string]any{"id": id}})
})
return login, err
}
func (s *Store) MergeClaim(ctx context.Context, id string) (string, error) {
login := ""
err := s.db.QueryRowContext(ctx, `SELECT merging_by FROM tasks WHERE id = ? AND merge_after_turn = 1`, id).Scan(&login)
if errors.Is(err, sql.ErrNoRows) {
return "", nil
}
return login, err
}
func (s *Store) ReserveMergeTurn(ctx context.Context, id, sessionID string) (bool, error) {
reserved := false
err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
res, err := tx.ExecContext(ctx, `UPDATE tasks SET merge_after_turn = 0 WHERE id = ? AND merge_after_turn = 1
AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?)`,
id, string(api.Implementer), string(api.AgentSessionStateRunning))
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil || n == 0 {
return err
}
idle := api.AgentSessionStateIdle
err = transitionSessionTx(ctx, tx, emit, sessionID, &idle, api.AgentSessionStateRunning, "byre", nil)
if err != nil {
return err
}
reserved = true
return nil
})
return reserved, err
}
func (s *Store) MergeClaims(ctx context.Context) ([]string, error) {
rows, err := s.db.QueryContext(ctx, `SELECT id FROM tasks WHERE status = ? AND merge_after_turn = 1`, string(api.TaskStatusActive))
if err != nil {
@@ -1416,46 +1441,50 @@ func (s *Store) TransitionSession(ctx context.Context, id string, from, to api.A
func (s *Store) transitionSession(ctx context.Context, id string, from *api.AgentSessionState, state api.AgentSessionState, actor string, payload map[string]any) error {
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
taskID := ""
projectID := ""
owner := ""
current := ""
role := ""
discard := ""
finalizing := ""
err := tx.QueryRowContext(ctx, `
SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.state, s.role, coalesce(t.discard_state, ''), coalesce(t.finalizing, '') FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?
`, id).Scan(&taskID, &projectID, &owner, &current, &role, &discard, &finalizing)
if errors.Is(err, sql.ErrNoRows) {
return ErrNotFound
}
if err != nil {
return err
}
if from != nil && current != string(*from) {
return fmt.Errorf("%w: session %s is %s, not %s", ErrStateChanged, id, current, *from)
}
if state == api.AgentSessionStateRunning && role == string(api.Implementer) && discard != "" {
return fmt.Errorf("%w: task %s is being discarded", ErrStateChanged, taskID)
}
if state == api.AgentSessionStateRunning && role == string(api.Implementer) && finalizing != "" {
return fmt.Errorf("%w: task %s is finishing its merge", ErrStateChanged, taskID)
}
_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(state), id)
if err != nil {
return err
}
if payload == nil {
payload = map[string]any{}
}
payload["state"] = string(state)
return emit(sessionEvent("session.state", projectID, taskID, owner, id, actor, payload))
return transitionSessionTx(ctx, tx, emit, id, from, state, actor, payload)
})
}
func transitionSessionTx(ctx context.Context, tx *sql.Tx, emit emitFunc, id string, from *api.AgentSessionState, state api.AgentSessionState, actor string, payload map[string]any) error {
taskID := ""
projectID := ""
owner := ""
current := ""
role := ""
discard := ""
finalizing := ""
err := tx.QueryRowContext(ctx, `
SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.state, s.role, coalesce(t.discard_state, ''), coalesce(t.finalizing, '') FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?
`, id).Scan(&taskID, &projectID, &owner, &current, &role, &discard, &finalizing)
if errors.Is(err, sql.ErrNoRows) {
return ErrNotFound
}
if err != nil {
return err
}
if from != nil && current != string(*from) {
return fmt.Errorf("%w: session %s is %s, not %s", ErrStateChanged, id, current, *from)
}
if state == api.AgentSessionStateRunning && role == string(api.Implementer) && discard != "" {
return fmt.Errorf("%w: task %s is being discarded", ErrStateChanged, taskID)
}
if state == api.AgentSessionStateRunning && role == string(api.Implementer) && finalizing != "" {
return fmt.Errorf("%w: task %s is finishing its merge", ErrStateChanged, taskID)
}
_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(state), id)
if err != nil {
return err
}
if payload == nil {
payload = map[string]any{}
}
payload["state"] = string(state)
return emit(sessionEvent("session.state", projectID, taskID, owner, id, actor, payload))
}
func (s *Store) SessionsWithSubagents(ctx context.Context) ([]api.AgentSession, error) {
return s.querySessions(ctx, `SELECT `+sessionColumns+` FROM agent_sessions WHERE subagents > 0`)
}