byre/internal/agents/summary.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/store"
)

const summaryRequest = "Call set_task_summary with a summary of the whole task, not just this turn, for whoever searches closed tasks later: what it set out to do, what was built or tried, the key decisions and why, and what was deliberately left out."

const discardPrompt = "This task is being discarded: when this turn ends, byre stops this session and deletes the worktree and branch without merging. Don't change any files. " +
	summaryRequest + " Say why it is being dropped, as far as you know. Then end your turn."

const discardRestartPrompt = "byre restarted while this turn was running. " + discardPrompt

type setTaskSummaryArgs struct {
	Summary string `json:"summary"`
}

func (m *Manager) setTaskSummaryTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "set_task_summary",
		Description: "Set the summary of this whole task that byre keeps after it is merged or discarded, where agents find it with search_closed_tasks. " +
			"Cover what the task set out to do, what was built or tried, the key decisions and why, and what was deliberately left out. It replaces any earlier summary.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"summary": {"type": "string", "description": "Markdown, a few short paragraphs or a bulleted list."}
			},
			"required": ["summary"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := setTaskSummaryArgs{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed set_task_summary call: %w", err)
			}
			summary := strings.TrimSpace(args.Summary)
			if summary == "" {
				return "", fmt.Errorf("%w: summary must not be empty", ErrInvalid)
			}
			err = m.store.SetTaskSummary(ctx, l.taskID, summary, agentActor)
			if errors.Is(err, store.ErrNotFound) {
				return "", fmt.Errorf("%w: this task is closed", ErrConflict)
			}
			if err != nil {
				return "", err
			}
			return "Task summary saved.", nil
		},
	}
}

const searchClosedDescription = "Search tasks that were merged or discarded, newest closed first. A task matches when query appears, ignoring case, in its title, merge commit title, TODO item, summary or merge commit; without query every closed task matches. " +
	"Each result has the task's id, outcome, merge commit, close time, titles, TODO item and the summary its agent wrote when it was merged or discarded (tasks closed before byre kept summaries have none). " +
	"Use it to check whether something was already done or tried, and why."

const searchClosedProperties = `
	"query": {"type": "string"},
	"limit": {"type": "integer", "description": "At most this many tasks, default 20, at most 100."}`

type searchClosedArgs struct {
	ProjectID string `json:"project_id"`
	Query     string `json:"query"`
	Limit     int    `json:"limit"`
}

func (m *Manager) searchClosedTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name:        "search_closed_tasks",
		Description: searchClosedDescription + " Searches this task's project.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {` + searchClosedProperties + `}}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := searchClosedArgs{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed search_closed_tasks call: %w", err)
			}
			return m.searchClosedText(ctx, l.projectID, args.Query, args.Limit)
		},
	}
}

func (m *Manager) searchClosedText(ctx context.Context, projectID, query string, limit int) (string, error) {
	if limit <= 0 {
		limit = 20
	}
	limit = min(limit, 100)

	found, err := m.store.SearchClosedTasks(ctx, projectID, strings.TrimSpace(query), limit)
	if err != nil {
		return "", err
	}
	if len(found) == 0 {
		return "No closed tasks match.", nil
	}

	b := &strings.Builder{}
	for _, c := range found {
		t := c.Task
		outcome := string(t.Status)
		if t.MergeCommit != nil {
			outcome += " as " + *t.MergeCommit
		}
		closed := ""
		if t.ClosedAt != nil {
			closed = t.ClosedAt.Format(time.RFC3339)
		}
		fmt.Fprintf(b, "## %s %q\n\nproject %s, %s %s, created by %s\n", t.Id, t.Title, c.Project, outcome, closed, t.CreatedBy)
		if t.MergeTitle != t.Title {
			fmt.Fprintf(b, "merge commit title: %s\n", t.MergeTitle)
		}
		if t.TodoItem != nil {
			fmt.Fprintf(b, "TODO item: %s\n", *t.TodoItem)
		}
		if t.ForkedFrom != nil {
			fmt.Fprintf(b, "forked from %s\n", *t.ForkedFrom)
		}
		summary := strings.TrimSpace(c.Summary)
		if summary == "" {
			summary = "(no summary)"
		}
		fmt.Fprintf(b, "\n%s\n\n", summary)
	}
	return strings.TrimSpace(b.String()), nil
}

func (m *Manager) Discard(ctx context.Context, task api.Task, actor string) (api.Task, error) {
	state, err := m.store.RequestDiscard(ctx, task.Id, actor)
	switch {
	case errors.Is(err, store.ErrNotFound):
		return api.Task{}, fmt.Errorf("%w: task %s is closed", ErrConflict, task.Id)
	case errors.Is(err, store.ErrHasSubtasks):
		serr := m.checkNoSubtasks(ctx, task)
		if serr == nil {
			serr = fmt.Errorf("%w: task %s: %w", ErrConflict, task.Id, err)
		}
		return api.Task{}, serr
	case errors.Is(err, store.ErrDiscardPending), errors.Is(err, store.ErrMerging), errors.Is(err, store.ErrMergingIntoParent):
		return api.Task{}, fmt.Errorf("%w: task %s: %w", ErrConflict, task.Id, err)
	case err != nil:
		return api.Task{}, err
	}
	if state == store.DiscardClosing {
		go m.finishDiscard(task.Id)
	} else {
		go m.startDiscard(task.Id)
	}
	return m.store.GetTask(ctx, task.Id)
}

func (m *Manager) startDiscard(taskID string) {
	ctx := context.Background()
	started, err := m.launchDiscard(ctx, taskID)
	if err == nil {
		if started {
			slog.Info("discard turn started", "task", taskID)
		}
		return
	}
	slog.Error("start discard turn", "task", taskID, "err", err)
	err = m.store.CloseDiscard(ctx, taskID)
	if err != nil {
		slog.Error("close discard", "task", taskID, "err", err)
		return
	}
	m.finishDiscard(taskID)
}

func (m *Manager) launchDiscard(ctx context.Context, taskID string) (bool, error) {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return false, err
	}
	sess, found, err := m.liveSession(ctx, taskID)
	if err != nil {
		return false, err
	}
	if !found {
		creator, err := m.store.GetUser(ctx, task.CreatedBy)
		if err != nil {
			return false, err
		}
		sess, err = m.newSession(ctx, task, creator, sess.Member, "")
		if err != nil {
			return false, err
		}
	}

	err = m.store.StartDiscardTurn(ctx, taskID, sess.Id, "byre")
	if errors.Is(err, store.ErrStateChanged) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	_, err = m.launch(ctx, sess, task, "byre", discardPrompt, nil)
	return err == nil, err
}

func (m *Manager) finishDiscard(taskID string) {
	m.mu.Lock()
	if m.finishing[taskID] {
		m.mu.Unlock()
		return
	}
	m.finishing[taskID] = true
	m.mu.Unlock()
	defer func() {
		m.mu.Lock()
		delete(m.finishing, taskID)
		m.mu.Unlock()
	}()

	ctx := context.Background()
	err := m.discardNow(ctx, taskID)
	if err == nil {
		return
	}
	slog.Error("finish discard", "task", taskID, "err", err)
	sessions, serr := m.store.ListSessions(ctx, taskID)
	if serr != nil {
		slog.Error("list sessions for discard error", "task", taskID, "err", serr)
		return
	}
	task, terr := m.store.GetTask(ctx, taskID)
	if terr != nil {
		slog.Error("load task for discard error", "task", taskID, "err", terr)
		return
	}
	m.fileTaskError(ctx, task, sessions, "Discarding the task failed: "+err.Error()+". Discard again to retry; byre also retries when it next starts.")
}

func (m *Manager) discardNow(ctx context.Context, taskID string) error {
	unlock := m.lockTask(taskID)
	defer unlock()
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return err
	}
	if task.Status != api.TaskStatusActive {
		return nil
	}
	discard, err := m.store.GetDiscard(ctx, taskID)
	if err != nil {
		return err
	}
	project, err := m.store.GetProject(ctx, task.ProjectId)
	if err != nil {
		return err
	}
	err = m.RemoveTask(ctx, task, project, discard.By)
	if err != nil {
		return err
	}
	_, err = m.store.DiscardTask(ctx, task.Id, discard.By)
	if err != nil {
		return err
	}
	go m.releaseMerge(project.Id, task.Id)
	return nil
}

func (m *Manager) recoverDiscards(ctx context.Context) error {
	tasks, err := m.store.PendingDiscards(ctx)
	if err != nil {
		return err
	}
	for _, t := range tasks {
		discard, err := m.store.GetDiscard(ctx, t.Id)
		if err != nil {
			return err
		}
		switch discard.State {
		case store.DiscardHeld:
			go m.startDiscard(t.Id)
		case store.DiscardClosing:
			go m.finishDiscard(t.Id)
		case store.DiscardRunning:
			sess, err := m.store.GetSession(ctx, discard.Session)
			if err != nil {
				return err
			}
			if sess.State == api.AgentSessionStateRunning {
				continue
			}
			err = m.store.CloseDiscard(ctx, t.Id)
			if err != nil {
				return err
			}
			go m.finishDiscard(t.Id)
		}
	}
	return nil
}