kandev/apps/backend/internal/backendapp/branch_materializer.go

package backendapp

import (
	"context"
	"fmt"
	"path/filepath"
	"time"

	"go.uber.org/zap"

	"github.com/kandev/kandev/internal/agent/runtime/lifecycle"
	"github.com/kandev/kandev/internal/common/logger"
	"github.com/kandev/kandev/internal/task/models"
	taskservice "github.com/kandev/kandev/internal/task/service"
	"github.com/kandev/kandev/internal/worktree"
)

// branchMaterializerRepo is the slim repo surface the mid-session worktree
// materializer needs. Defined as an interface so the wiring stays decoupled
// from the concrete sqlite repo.
type branchMaterializerRepo interface {
	GetTask(ctx context.Context, id string) (*models.Task, error)
	GetTaskRepository(ctx context.Context, id string) (*models.TaskRepository, error)
	ListTaskRepositories(ctx context.Context, taskID string) ([]*models.TaskRepository, error)
	GetRepository(ctx context.Context, id string) (*models.Repository, error)
	ListTaskSessions(ctx context.Context, taskID string) ([]*models.TaskSession, error)
	GetTaskEnvironmentByTaskID(ctx context.Context, taskID string) (*models.TaskEnvironment, error)
	UpdateTaskEnvironment(ctx context.Context, env *models.TaskEnvironment) error
}

// agentctlRescanner is the lifecycle-tier surface the materializer uses to
// reach a running agentctl by session_id. The lifecycle manager owns the
// authenticated agentctl client; executors_running may carry endpoint metadata
// for diagnostics/recovery, but callers should not bypass lifecycle ownership.
// Defined as an interface so the materializer stays decoupled from the
// lifecycle package.
type agentctlRescanner interface {
	RescanWorkspaceForSession(ctx context.Context, sessionID, workDir string, sourceRoots ...[]string) error
	NotifyWorktreeMaterialized(ctx context.Context, wt lifecycle.MaterializedWorktree)
}

// branchMaterializer creates a worktree on disk for a newly added
// task_repositories row so MCP-driven "add branch" surfaces the worktree in
// the UI without waiting for a session relaunch. Implements
// taskservice.BranchMaterializer.
//
// It is best-effort: when there's no active session (no tasksession row yet,
// or the task has never launched), it no-ops and lets the standard
// multi-repo prepare path materialize on the next launch.
type branchMaterializer struct {
	repo        branchMaterializerRepo
	worktreeMgr *worktree.Manager
	rescanner   agentctlRescanner
	logger      *logger.Logger
}

type branchMaterialization struct {
	environment  *models.TaskEnvironment
	session      *models.TaskSession
	worktree     *worktree.Worktree
	repositoryID string
	slug         string
	taskID       string
}

func newBranchMaterializer(repo branchMaterializerRepo, mgr *worktree.Manager, lc *lifecycle.Manager, log *logger.Logger) *branchMaterializer {
	var rescanner agentctlRescanner
	if lc != nil {
		rescanner = lc
	}
	return &branchMaterializer{
		repo:        repo,
		worktreeMgr: mgr,
		rescanner:   rescanner,
		logger:      log.WithFields(zap.String("component", "branch_materializer")),
	}
}

// MaterializeBranch creates the worktree dir + persists the
// task_session_worktrees row for the just-inserted task_repositories row.
// Designed to be idempotent: the worktree manager's reuse path catches a
// rerun for the same (session, repo, branch_slug) triple and returns the
// existing worktree.
func (b *branchMaterializer) MaterializeBranch(ctx context.Context, taskID, taskRepositoryID string) error {
	materialization, err := b.materializeUnfinalized(ctx, taskID, taskRepositoryID)
	if err != nil || materialization == nil {
		return err
	}
	b.finalize(materialization, ctx)
	return nil
}

// materializeUnfinalized creates the worktree but deliberately does not
// promote, rescan, or notify. Workspace-source batches call this phase for
// every repository and publish only after their folders and live sessions have
// all adopted the new root.
func (b *branchMaterializer) materializeUnfinalized(ctx context.Context, taskID, taskRepositoryID string) (*branchMaterialization, error) {
	if b == nil || b.worktreeMgr == nil {
		return nil, nil
	}
	req, env, session, slug, ok, err := b.prepareMaterializeRequest(ctx, taskID, taskRepositoryID)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, nil
	}
	wt, err := b.worktreeMgr.Create(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("create worktree: %w", err)
	}
	b.logger.Info("materialized branch worktree",
		zap.String("task_id", taskID),
		zap.String("task_repository_id", taskRepositoryID),
		zap.String("worktree_id", wt.ID),
		zap.String("path", wt.Path),
		zap.String("branch", wt.Branch))
	return &branchMaterialization{environment: env, session: session, worktree: wt, repositoryID: req.RepositoryID, slug: slug, taskID: taskID}, nil
}

func (b *branchMaterializer) finalize(materialization *branchMaterialization, ctx context.Context) {
	b.finalizeMaterialize(ctx, materialization.environment, materialization.session, materialization.worktree, materialization.repositoryID, materialization.slug, materialization.taskID)
}

// prepareMaterializeRequest builds the worktree.CreateRequest plus the
// (env, session, slug) context the finalize step needs. Returns ok=false
// for the legitimate "skip materialize" cases (no local path, no active
// session, env not provisioned); ok=true means CreateRequest is ready.
func (b *branchMaterializer) prepareMaterializeRequest(
	ctx context.Context, taskID, taskRepositoryID string,
) (worktree.CreateRequest, *models.TaskEnvironment, *models.TaskSession, string, bool, error) {
	task, tr, repo, err := b.loadContext(ctx, taskID, taskRepositoryID)
	if err != nil {
		return worktree.CreateRequest{}, nil, nil, "", false, err
	}
	if repo.LocalPath == "" {
		// Provider-backed repos that haven't been cloned yet need the
		// orchestrator's repoCloner path. Defer to next launch.
		b.logger.Info("skipping materialize: repository has no local path",
			zap.String("repository_id", repo.ID), zap.String("task_id", taskID))
		return worktree.CreateRequest{}, nil, nil, "", false, nil
	}
	session, err := b.pickActiveSession(ctx, taskID)
	if err != nil {
		return worktree.CreateRequest{}, nil, nil, "", false, err
	}
	if session == nil {
		b.logger.Info("skipping materialize: no active session for task",
			zap.String("task_id", taskID))
		return worktree.CreateRequest{}, nil, nil, "", false, nil
	}
	env, err := b.repo.GetTaskEnvironmentByTaskID(ctx, taskID)
	if err != nil {
		return worktree.CreateRequest{}, nil, nil, "", false, fmt.Errorf("lookup task environment: %w", err)
	}
	if env == nil || env.TaskDirName == "" {
		b.logger.Info("skipping materialize: task environment not provisioned yet",
			zap.String("task_id", taskID))
		return worktree.CreateRequest{}, nil, nil, "", false, nil
	}
	slug := deriveBranchSlugForRow(tr)
	req := worktree.CreateRequest{
		TaskID:                 taskID,
		SessionID:              session.ID,
		TaskTitle:              task.Title,
		RepositoryID:           repo.ID,
		RepositoryPath:         repo.LocalPath,
		BaseBranch:             tr.BaseBranch,
		FallbackBaseBranch:     repo.DefaultBranch,
		CheckoutBranch:         tr.CheckoutBranch,
		WorktreeBranchPrefix:   repo.WorktreeBranchPrefix,
		WorktreeBranchTemplate: repo.WorktreeBranchTemplate,
		WorktreeBranchTicket:   worktree.TicketForBranchName(task.Identifier, task.Metadata),
		PullBeforeWorktree:     repo.PullBeforeWorktree,
		TaskDirName:            env.TaskDirName,
		RepoName:               repo.Name,
		BranchSlug:             slug,
	}
	return req, env, session, slug, true, nil
}

// finalizeMaterialize handles the post-Create plumbing: promote the task
// environment's workspace_path to the task root when the new worktree is a
// sibling, ping the running agentctl to add per-repo trackers, and emit a
// frontend WS event so the session's worktree tabs render immediately.
func (b *branchMaterializer) finalizeMaterialize(
	ctx context.Context,
	env *models.TaskEnvironment,
	session *models.TaskSession,
	wt *worktree.Worktree,
	repositoryID, slug, taskID string,
) {
	taskRoot := b.promoteWorkspacePathIfNeeded(ctx, env, wt.Path)
	b.notifyAgentctlRescan(ctx, session, taskRoot)
	if b.rescanner != nil {
		b.rescanner.NotifyWorktreeMaterialized(ctx, lifecycle.MaterializedWorktree{
			TaskID:            taskID,
			SessionID:         session.ID,
			WorktreeID:        wt.ID,
			WorktreePath:      wt.Path,
			WorktreeBranch:    wt.Branch,
			RepositoryID:      repositoryID,
			BranchSlug:        slug,
			TaskWorkspacePath: taskRoot,
		})
	}
}

// promoteWorkspacePathIfNeeded switches task_environments.workspace_path
// to the task root that contains the newly-created worktree, so subsequent
// session relaunches pick up multi-repo mode via prepareMultiRepo. The
// agent process's CWD does not change — that's set at launch time.
// Returns the resolved task-root path so the caller can pass it to the
// agentctl rescan.
//
// Anchors the task root in env.TaskDirName rather than env.WorkspacePath:
// a few historical envs landed with workspace_path pointing at the source
// repo's local_path (the orchestrator's computeWorkspacePath fallback)
// instead of the primary worktree. Comparing against env.WorkspacePath
// alone would refuse to promote those envs because the source repo isn't
// a sibling of the new worktree — but the task root we want IS still
// `filepath.Dir(newWorktreePath)` and its basename matches env.TaskDirName,
// which is a reliable cross-check that the worktree landed under the right
// task root and not somewhere unrelated.
func (b *branchMaterializer) promoteWorkspacePathIfNeeded(ctx context.Context, env *models.TaskEnvironment, newWorktreePath string) string {
	taskRoot := filepath.Dir(newWorktreePath)
	if env.WorkspacePath == taskRoot {
		return taskRoot
	}
	if env.TaskDirName == "" || filepath.Base(taskRoot) != env.TaskDirName {
		b.logger.Warn("skip workspace_path promotion: new worktree parent does not match env.task_dir_name",
			zap.String("env_workspace_path", env.WorkspacePath),
			zap.String("env_task_dir_name", env.TaskDirName),
			zap.String("new_worktree_path", newWorktreePath),
			zap.String("derived_task_root", taskRoot))
		return env.WorkspacePath
	}
	prev := env.WorkspacePath
	env.WorkspacePath = taskRoot
	env.UpdatedAt = time.Now().UTC()
	if err := b.repo.UpdateTaskEnvironment(ctx, env); err != nil {
		b.logger.Warn("failed to promote workspace_path to task root",
			zap.String("task_environment_id", env.ID),
			zap.String("from", prev),
			zap.String("to", taskRoot),
			zap.Error(err))
		return prev
	}
	b.logger.Info("promoted workspace_path to task root",
		zap.String("task_environment_id", env.ID),
		zap.String("from", prev),
		zap.String("to", taskRoot))
	return taskRoot
}

// notifyAgentctlRescan tells the running agentctl instance attached to
// this session to re-discover repo subdirs and add per-repo trackers for
// new sibling worktrees. Routes through the lifecycle manager because the
// lifecycle manager is the single owner of authenticated per-execution
// agentctl clients.
//
// No-op when no rescanner is wired (test stubs, agentless test envs) or
// when the session has no active execution (agent stopped). The next
// session launch picks up the new layout via prepareMultiRepo regardless.
func (b *branchMaterializer) notifyAgentctlRescan(ctx context.Context, session *models.TaskSession, taskRoot string) {
	if b.rescanner == nil {
		return
	}
	rescanCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	if err := b.rescanner.RescanWorkspaceForSession(rescanCtx, session.ID, taskRoot); err != nil {
		b.logger.Warn("agentctl rescan failed; new worktree will appear after next session restart",
			zap.String("session_id", session.ID),
			zap.String("task_root", taskRoot),
			zap.Error(err))
	}
}

// loadContext fetches the task + task_repository row + repository entity in
// one place so MaterializeBranch can fail-fast on missing data without
// scattering nil checks across the body.
func (b *branchMaterializer) loadContext(ctx context.Context, taskID, taskRepositoryID string) (*models.Task, *models.TaskRepository, *models.Repository, error) {
	task, err := b.repo.GetTask(ctx, taskID)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("get task: %w", err)
	}
	if task == nil {
		return nil, nil, nil, fmt.Errorf("task not found: %s", taskID)
	}
	tr, err := b.repo.GetTaskRepository(ctx, taskRepositoryID)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("get task repository: %w", err)
	}
	if tr == nil {
		return nil, nil, nil, fmt.Errorf("task repository not found: %s", taskRepositoryID)
	}
	repo, err := b.repo.GetRepository(ctx, tr.RepositoryID)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("get repository: %w", err)
	}
	if repo == nil {
		return nil, nil, nil, fmt.Errorf("repository not found: %s", tr.RepositoryID)
	}
	return task, tr, repo, nil
}

// pickActiveSession returns the most recently active session for the task,
// or nil if none qualifies. Worktree persistence requires a session_id, so
// when there's no candidate we no-op rather than orphaning a worktree
// without an owning session.
func (b *branchMaterializer) pickActiveSession(ctx context.Context, taskID string) (*models.TaskSession, error) {
	sessions, err := b.repo.ListTaskSessions(ctx, taskID)
	if err != nil {
		return nil, fmt.Errorf("list sessions: %w", err)
	}
	var best *models.TaskSession
	for _, s := range sessions {
		if !sessionEligibleForMaterialize(s) {
			continue
		}
		if best == nil || s.UpdatedAt.After(best.UpdatedAt) {
			best = s
		}
	}
	return best, nil
}

func sessionEligibleForMaterialize(s *models.TaskSession) bool {
	switch s.State {
	case models.TaskSessionStateRunning,
		models.TaskSessionStateStarting,
		models.TaskSessionStateWaitingForInput,
		models.TaskSessionStateCreated:
		return true
	}
	return false
}

// deriveBranchSlugForRow chooses the slug input that uniquely identifies the
// row's worktree on disk. CheckoutBranch wins when set (local-executor and
// PR-style flows put the branch there); BaseBranch is the fallback for the
// worktree-executor flow where the branch lives in BaseBranch and
// CheckoutBranch is empty.
func deriveBranchSlugForRow(tr *models.TaskRepository) string {
	slug := worktree.SanitizeBranchSlug(tr.CheckoutBranch)
	if slug == "" {
		slug = worktree.SanitizeBranchSlug(tr.BaseBranch)
	}
	return slug
}

// ensureBranchMaterializerSatisfiesInterface is a compile-time guard so the
// interface contract drift surfaces at build, not at runtime.
var _ taskservice.BranchMaterializer = (*branchMaterializer)(nil)