byre/internal/agents/fork.go

package agents

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"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)
	return t, err
}

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)
	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) (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:         fmt.Sprintf("%s (fork of #%d)", src.Title, n),
		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 := ""
	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 == "") {
			break
		}
		switch ev.Kind {
		case "agent.tool_call":
			open[id] = true
		case "agent.tool_result":
			delete(open, id)
		}
		at = uuid
		if after && len(open) == 0 {
			break
		}
	}
	return at
}