byre/internal/agents/fork.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"
)

func (m *Manager) Fork(ctx context.Context, src api.Task, project api.Project, n int, creator api.User) (api.Task, error) {
	t, _, err := m.fork(ctx, src, project, n, creator, forkTitle(src, n))
	return t, err
}

func forkTitle(src api.Task, n int) string {
	return fmt.Sprintf("%s (fork of #%d)", src.Title, n)
}

func (m *Manager) forkTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "fork",
		Description: "Fork this task here, for when two ways forward are both worth finishing and the human should compare the results rather than you picking one. " +
			"byre snapshots your worktree now and starts a new task from that snapshot, with its own branch and worktree, whose agent continues this conversation up to just before this call " +
			"and then gets prompt as its next message. You carry on here with the other way. " +
			"Each fork doubles the work and gives the human another result to compare, so fork only when you strongly believe both paths are worth following.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {
			"title": {"type": "string", "description": "The new task's title, naming the path it takes."},
			"prompt": {"type": "string", "description": "The forked agent's next message: tell it that it is the other side of a fork and which path to take."}
		}, "required": ["title", "prompt"]}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := struct {
				Title  string `json:"title"`
				Prompt string `json:"prompt"`
			}{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed fork call: %w", err)
			}
			if strings.TrimSpace(args.Title) == "" || strings.TrimSpace(args.Prompt) == "" {
				return "", errors.New("title and prompt are required")
			}
			sess, err := m.store.GetSession(ctx, l.id)
			if err != nil {
				return "", err
			}
			creator, err := m.store.GetUser(ctx, sess.CreatedBy)
			if err != nil {
				return "", err
			}
			src, err := m.store.GetTask(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			project, err := m.store.GetProject(ctx, src.ProjectId)
			if err != nil {
				return "", err
			}
			snap, err := m.cutSnapshot(ctx, liveTarget(l), m.ownsTask(l.id), api.SnapshotTriggerFork)
			if err != nil {
				return "", err
			}
			if snap == nil {
				return "", fmt.Errorf("%w: task %s is closing", ErrConflict, src.Id)
			}
			t, sessionID, err := m.fork(ctx, src, project, snap.N, creator, args.Title)
			if err != nil {
				return "", err
			}
			_, _, err = m.Message(ctx, sessionID, agentActor, args.Prompt, "")
			if err != nil {
				return "", fmt.Errorf("task %s forked, but its agent did not start: %w", t.Id, err)
			}
			slog.Info("task forked by its agent", "task", src.Id, "fork", t.Id, "snapshot", snap.N)
			return fmt.Sprintf("Forked at snapshot #%d as task %s %q, branch %s, worktree %s; its agent continues from just before this call with your prompt. Carry on here with your own path.",
				snap.N, t.Id, t.Title, t.Branch, t.WorktreePath), nil
		},
	}
}

func (m *Manager) AnswerAgain(ctx context.Context, creator api.User, create api.InboxResponseCreate) (api.Task, error) {
	item, err := m.store.GetInboxItem(ctx, create.ItemId)
	if err != nil {
		return api.Task{}, err
	}
	if item.State != api.Answered {
		return api.Task{}, fmt.Errorf("%w: question %s is %s, not answered", ErrConflict, item.Id, item.State)
	}
	n := 0
	for _, r := range item.Responses {
		if r.SnapshotN != nil {
			n = *r.SnapshotN
		}
	}
	if n == 0 {
		return api.Task{}, fmt.Errorf("%w: question %s's answer reached the agent without a snapshot to fork from", ErrConflict, item.Id)
	}
	open := item
	open.State = api.Open
	_, err = CheckResponse(open, create)
	if err != nil {
		return api.Task{}, err
	}

	src, err := m.store.GetTask(ctx, item.TaskId)
	if err != nil {
		return api.Task{}, err
	}
	project, err := m.store.GetProject(ctx, src.ProjectId)
	if err != nil {
		return api.Task{}, err
	}
	t, sessionID, err := m.fork(ctx, src, project, n, creator, forkTitle(src, n))
	if err != nil {
		return api.Task{}, err
	}

	id := store.NewID()
	err = m.store.CreateAnsweredItem(ctx, api.InboxItem{
		Id:             id,
		ProjectId:      t.ProjectId,
		TaskId:         t.Id,
		AgentSessionId: sessionID,
		ThreadId:       id,
		Payload:        item.Payload,
		CreatedAt:      time.Now().UTC(),
	}, newResponse(id, creator.Login, create))
	if err != nil {
		return api.Task{}, fmt.Errorf("task %s forked, but its answer was not recorded: %w", t.Id, err)
	}
	slog.Info("question answered again", "item", item.Id, "fork", t.Id, "snapshot", n, "user", creator.Login)

	go func() {
		err := m.Deliver(context.Background(), sessionID)
		if err != nil {
			slog.Error("deliver inbox responses", "session", sessionID, "err", err)
		}
	}()
	return t, nil
}

func (m *Manager) fork(ctx context.Context, src api.Task, project api.Project, n int, creator api.User, title string) (api.Task, string, error) {
	if creator.GitName == "" || creator.GitEmail == "" {
		return api.Task{}, "", ErrNoGitIdentity
	}
	if src.Status != api.TaskStatusActive {
		return api.Task{}, "", fmt.Errorf("%w: task %s is %s", ErrConflict, src.Id, src.Status)
	}

	snap, err := m.store.GetSnapshot(ctx, src.Id, n)
	if err != nil {
		return api.Task{}, "", err
	}
	srcSess, err := m.store.GetSession(ctx, snap.AgentSessionId)
	if err != nil {
		return api.Task{}, "", err
	}
	if srcSess.Vendor == api.Antigravity {
		return api.Task{}, "", fmt.Errorf("%w: snapshot %d of task %s is from an Antigravity session, and agy can't fork a conversation", ErrConflict, n, src.Id)
	}
	events, err := m.store.TaskEvents(ctx, src.Id, 0, -1)
	if err != nil {
		return api.Task{}, "", err
	}
	at := forkPoint(events, snap.AgentSessionId, snap.EventSeq)
	if at == "" {
		return api.Task{}, "", fmt.Errorf("%w: snapshot %d of task %s has no recorded agent message at or before it to fork the conversation from", ErrConflict, n, src.Id)
	}

	var frozen *store.FrozenWorkflow
	if src.Workflow != nil {
		f, err := m.store.TaskWorkflow(ctx, src.Id)
		if err != nil {
			return api.Task{}, "", err
		}
		frozen = &f
	}

	t, err := m.createTask(ctx, project, api.Task{
		Title:         title,
		MergeTitle:    src.MergeTitle,
		CreatedBy:     creator.Login,
		TodoItem:      src.TodoItem,
		ForkedFrom:    &src.Id,
		ForkedAt:      &n,
		WorkflowRun:   src.WorkflowRun,
		WorkflowStage: src.WorkflowStage,
		ParentTaskId:  src.ParentTaskId,
	}, frozen, snap.CommitSha, creator.Login)
	if err != nil {
		return api.Task{}, "", err
	}

	id := newUUID()
	err = m.store.CreateSession(ctx, api.AgentSession{
		Id:              id,
		TaskId:          t.Id,
		Vendor:          srcSess.Vendor,
		Model:           srcSess.Model,
		VendorSessionId: initialVendorSessionID(srcSess.Vendor, id),
		State:           api.AgentSessionStateIdle,
		Role:            api.Implementer,
		CreatedBy:       creator.Login,
		CreatedAt:       time.Now().UTC(),
		ForkedFrom:      &srcSess.VendorSessionId,
		ForkedAt:        &at,
	}, t.ProjectId)
	if err != nil {
		return api.Task{}, "", fmt.Errorf("task %s created, but its agent session was not: %w", t.Id, err)
	}
	return t, id, nil
}

func forkPoint(events []api.Event, sessionID string, seq int64) string {
	at := ""
	prev := ""
	cut := ""
	open := map[string]bool{}
	for _, ev := range events {
		if ev.AgentSessionId == nil || *ev.AgentSessionId != sessionID {
			continue
		}
		uuid, _ := ev.Payload["vendor_uuid"].(string)
		if uuid == "" {
			continue
		}
		id := ""
		switch ev.Kind {
		case "agent.tool_call":
			id, _ = ev.Payload["id"].(string)
		case "agent.tool_result":
			id, _ = ev.Payload["tool_use_id"].(string)
		}
		after := ev.Seq > seq
		if after && (len(open) == 0 || id == "") {
			return at
		}
		switch ev.Kind {
		case "agent.tool_call":
			if len(open) == 0 {
				cut = at
				if at == uuid {
					cut = prev
				}
			}
			open[id] = true
		case "agent.tool_result":
			delete(open, id)
		}
		if uuid != at {
			prev = at
		}
		at = uuid
		if after && len(open) == 0 {
			return at
		}
	}
	// Calls still open have no results yet, as when the agent forks from inside a tool call: cut before them.
	if len(open) > 0 {
		return cut
	}
	return at
}