byre/internal/agents/fork.go

package agents

import (
	"context"
	"fmt"
	"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) {
	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
	}
	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, 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
}