byre/internal/agents/pause.go

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 hookToolUseID(input json.RawMessage) string {
	in := struct {
		ToolUseID string `json:"tool_use_id"`
	}{}
	json.Unmarshal(input, &in)
	return in.ToolUseID
}

func (m *Manager) holdWhilePaused(ctx context.Context, l *live, toolUseID string) 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, "tool_use_id", toolUseID)
			if toolUseID != "" {
				err := m.emit(context.WithoutCancel(ctx), l, agentActor, "agent.held", map[string]any{"tool_use_id": toolUseID})
				if err != nil {
					slog.Error("record held tool call", "session", l.id, "err", err)
				}
			}
		}
		held = true
		select {
		case <-ch:
		case <-ctx.Done():
			return held
		}
	}
}

func (m *Manager) heldBeforeTool(ctx context.Context, l *live, toolUseID string) bool {
	return m.holdWhilePaused(ctx, l, toolUseID) && 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, hookToolUseID(input)) {
			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
	}
}