Add a pause that holds an agent after its current step

internal/agents/pause.go

@@ -0,0 +1,177 @@
+package agents
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+
+	"g.fc.run/ai/byre/internal/api"
+	"g.fc.run/ai/byre/internal/hookcmd"
+)
+
+const pausedCallNotice = "Not run: the human paused you just before this call and sent you a message while you were paused. Read it, then carry on, making this call again only if it still fits."
+
+func (m *Manager) Pause(ctx context.Context, sessionID, actor string) (api.AgentSession, error) {
+	sess, err := m.store.GetSession(ctx, sessionID)
+	if err != nil {
+		return api.AgentSession{}, err
+	}
+	if sess.Role != api.Implementer || sess.State == api.AgentSessionStateCompleted {
+		return api.AgentSession{}, fmt.Errorf("%w: session %s is not an active task agent", ErrConflict, sess.Id)
+	}
+
+	unlock := m.lockDelivery(sessionID)
+	defer unlock()
+	if m.isPaused(sessionID) {
+		return sess, nil
+	}
+	m.mu.Lock()
+	m.paused[sessionID] = make(chan struct{})
+	m.mu.Unlock()
+	err = m.store.SetSessionPaused(context.WithoutCancel(ctx), sessionID, true, actor)
+	if err != nil {
+		m.mu.Lock()
+		delete(m.paused, sessionID)
+		m.mu.Unlock()
+		return api.AgentSession{}, err
+	}
+	slog.Info("session paused", "session", sessionID, "task", sess.TaskId, "actor", actor)
+	return m.store.GetSession(ctx, sessionID)
+}
+
+func (m *Manager) Resume(ctx context.Context, sessionID, actor string) (api.AgentSession, error) {
+	sess, err := m.store.GetSession(ctx, sessionID)
+	if err != nil {
+		return api.AgentSession{}, err
+	}
+
+	unlock := m.lockDelivery(sessionID)
+	if !m.isPaused(sessionID) {
+		unlock()
+		return sess, nil
+	}
+	err = m.store.SetSessionPaused(context.WithoutCancel(ctx), sessionID, false, actor)
+	if err != nil {
+		unlock()
+		return api.AgentSession{}, err
+	}
+	m.dropPause(sessionID)
+	unlock()
+	slog.Info("session resumed", "session", sessionID, "task", sess.TaskId, "actor", actor)
+
+	go func() {
+		err := m.Deliver(context.Background(), sessionID)
+		if err != nil {
+			slog.Error("deliver after resume", "session", sessionID, "err", err)
+		}
+	}()
+	return m.store.GetSession(ctx, sessionID)
+}
+
+func (m *Manager) loadPauses(ctx context.Context) error {
+	sessions, err := m.store.PausedSessions(ctx)
+	if err != nil {
+		return err
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for _, sess := range sessions {
+		m.paused[sess.Id] = make(chan struct{})
+	}
+	return nil
+}
+
+func (m *Manager) dropPause(sessionID string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	ch := m.paused[sessionID]
+	if ch != nil {
+		close(ch)
+		delete(m.paused, sessionID)
+	}
+}
+
+func (m *Manager) waitResumed(sessionID string) {
+	for {
+		m.mu.Lock()
+		ch := m.paused[sessionID]
+		m.mu.Unlock()
+		if ch == nil {
+			return
+		}
+		<-ch
+	}
+}
+
+func (m *Manager) recoverAfterPause(sessionID string) {
+	ctx := context.Background()
+	for {
+		m.waitResumed(sessionID)
+		unlock := m.lockDelivery(sessionID)
+		if m.isPaused(sessionID) {
+			unlock()
+			continue
+		}
+		sess, err := m.store.GetSession(ctx, sessionID)
+		if err == nil && sess.State == api.AgentSessionStateRunning {
+			err = m.recoverSession(ctx, sess)
+		}
+		unlock()
+		if err != nil {
+			slog.Error("resume session after pause", "session", sessionID, "err", err)
+		}
+		return
+	}
+}
+
+func (m *Manager) isPaused(sessionID string) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.paused[sessionID] != nil
+}
+
+func (m *Manager) holdWhilePaused(ctx context.Context, l *live) bool {
+	held := false
+	for {
+		m.mu.Lock()
+		ch := m.paused[l.id]
+		m.mu.Unlock()
+		if ch == nil {
+			return held
+		}
+		if !held {
+			slog.Info("agent held by pause", "session", l.id)
+		}
+		held = true
+		select {
+		case <-ch:
+		case <-ctx.Done():
+			return held
+		}
+	}
+}
+
+func (m *Manager) heldBeforeTool(ctx context.Context, l *live) bool {
+	return m.holdWhilePaused(ctx, l) && ctx.Err() == nil
+}
+
+func (m *Manager) beforeTool(l *live) hookcmd.Func {
+	return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
+		if !m.heldBeforeTool(ctx, l) {
+			return map[string]any{}, nil
+		}
+		mt := m.deliverMidTurn(ctx, l)
+		if mt.text == "" {
+			return map[string]any{}, nil
+		}
+		return map[string]any{
+			"hookSpecificOutput": map[string]any{
+				"hookEventName":            "PreToolUse",
+				"permissionDecision":       "deny",
+				"permissionDecisionReason": pausedCallNotice,
+				"additionalContext":        mt.text,
+			},
+		}, nil
+	}
+}