Add support for importing PRs into work units

cmd/work.go

@@ -12,12 +12,8 @@ import (
 	"github.com/newhook/co/internal/control"
 	"github.com/newhook/co/internal/db"
 	"github.com/newhook/co/internal/git"
-	"github.com/newhook/co/internal/mise"
-	"github.com/newhook/co/internal/names"
 	"github.com/newhook/co/internal/project"
-	cosignal "github.com/newhook/co/internal/signal"
 	"github.com/newhook/co/internal/work"
-	"github.com/newhook/co/internal/worktree"
 	"github.com/spf13/cobra"
 )
 
@@ -214,7 +210,6 @@ func runWorkCreate(cmd *cobra.Command, args []string) error {
 
 	mainRepoPath := proj.MainRepoPath()
 	gitOps := git.NewOperations()
-	wtOps := worktree.NewOperations()
 	beadID := args[0]
 
 	// Expand the bead (handles epics and transitive deps)
@@ -245,7 +240,6 @@ func runWorkCreate(cmd *cobra.Command, args []string) error {
 	// Determine branch name
 	var branchName string
 	var useExistingBranch bool
-	var branchExistsOnRemote bool
 
 	if flagFromBranch != "" {
 		// Use an existing branch
@@ -260,7 +254,6 @@ func runWorkCreate(cmd *cobra.Command, args []string) error {
 		if !existsLocal && !existsRemote {
 			return fmt.Errorf("branch %s does not exist locally or on remote", branchName)
 		}
-		branchExistsOnRemote = existsRemote
 	} else if flagBranchName != "" {
 		// Use provided branch name
 		branchName = flagBranchName
@@ -281,101 +274,25 @@ func runWorkCreate(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	// Generate work ID
-	workID, err := proj.DB.GenerateWorkID(ctx, branchName, proj.Config.Project.Name)
+	// Create work asynchronously (control plane handles worktree creation, git push, orchestrator spawn)
+	result, err := work.CreateWorkAsyncWithOptions(ctx, proj, work.CreateWorkAsyncOptions{
+		BranchName:        branchName,
+		BaseBranch:        baseBranch,
+		RootIssueID:       beadID,
+		Auto:              flagAutoRun,
+		UseExistingBranch: useExistingBranch,
+		BeadIDs:           expandedIssueIDs,
+	})
 	if err != nil {
-		return fmt.Errorf("failed to generate work ID: %w", err)
+		return fmt.Errorf("failed to create work: %w", err)
 	}
-	fmt.Printf("Work ID: %s\n", workID)
 
-	// Block signals during critical worktree creation
-	cosignal.BlockSignals()
-	defer cosignal.UnblockSignals()
-
-	// Create work subdirectory
-	workDir := filepath.Join(proj.Root, workID)
-	if err := os.Mkdir(workDir, 0755); err != nil {
-		return fmt.Errorf("failed to create work directory: %w", err)
+	fmt.Printf("\nCreated work: %s\n", result.WorkID)
+	if result.WorkerName != "" {
+		fmt.Printf("Worker: %s\n", result.WorkerName)
 	}
-
-	// Create git worktree inside work directory
-	worktreePath := filepath.Join(workDir, "tree")
-
-	if useExistingBranch {
-		// For existing branches, fetch the branch first
-		if err := gitOps.FetchBranch(ctx, mainRepoPath, branchName); err != nil {
-			// Ignore fetch errors - branch might only exist locally
-			fmt.Printf("Note: Could not fetch branch %s from origin (may only exist locally)\n", branchName)
-		}
-
-		// Create worktree from existing branch
-		if err := wtOps.CreateFromExisting(ctx, mainRepoPath, worktreePath, branchName); err != nil {
-			os.RemoveAll(workDir)
-			return err
-		}
-
-		// Only push if branch doesn't exist on remote yet
-		if !branchExistsOnRemote {
-			if err := gitOps.PushSetUpstream(ctx, branchName, worktreePath); err != nil {
-				_ = wtOps.RemoveForce(ctx, mainRepoPath, worktreePath)
-				os.RemoveAll(workDir)
-				return err
-			}
-		}
-	} else {
-		// Fetch latest from origin for the base branch
-		if err := gitOps.FetchBranch(ctx, mainRepoPath, baseBranch); err != nil {
-			os.RemoveAll(workDir)
-			return fmt.Errorf("failed to fetch base branch: %w", err)
-		}
-
-		// Create worktree with new branch based on origin/<baseBranch>
-		if err := wtOps.Create(ctx, mainRepoPath, worktreePath, branchName, "origin/"+baseBranch); err != nil {
-			os.RemoveAll(workDir)
-			return err
-		}
-
-		// Push branch and set upstream
-		if err := gitOps.PushSetUpstream(ctx, branchName, worktreePath); err != nil {
-			_ = wtOps.RemoveForce(ctx, mainRepoPath, worktreePath)
-			os.RemoveAll(workDir)
-			return err
-		}
-	}
-
-	// Initialize mise in worktree if needed
-	if err := mise.Initialize(worktreePath); err != nil {
-		fmt.Printf("Warning: mise initialization failed: %v\n", err)
-	}
-
-	// Get a human-readable name for this worker
-	workerName, err := names.GetNextAvailableName(ctx, proj.DB.DB)
-	if err != nil {
-		fmt.Printf("Warning: failed to get worker name: %v\n", err)
-	}
-
-	// Create work record in database with the root issue ID (the original bead that was expanded)
-	if err := proj.DB.CreateWork(ctx, workID, workerName, worktreePath, branchName, baseBranch, beadID, flagAutoRun); err != nil {
-		_ = wtOps.RemoveForce(ctx, mainRepoPath, worktreePath)
-		os.RemoveAll(workDir)
-		return fmt.Errorf("failed to create work record: %w", err)
-	}
-
-	// Add beads to work_beads
-	if err := work.AddBeadsToWorkInternal(ctx, proj, workID, expandedIssueIDs); err != nil {
-		_ = wtOps.RemoveForce(ctx, mainRepoPath, worktreePath)
-		os.RemoveAll(workDir)
-		return fmt.Errorf("failed to add beads to work: %w", err)
-	}
-
-	fmt.Printf("\nCreated work: %s\n", workID)
-	if workerName != "" {
-		fmt.Printf("Worker: %s\n", workerName)
-	}
-	fmt.Printf("Directory: %s\n", workDir)
-	fmt.Printf("Worktree: %s\n", worktreePath)
-	fmt.Printf("Branch: %s\n", branchName)
-	fmt.Printf("Base Branch: %s\n", baseBranch)
+	fmt.Printf("Branch: %s\n", result.BranchName)
+	fmt.Printf("Base Branch: %s\n", result.BaseBranch)
 
 	// Display beads
 	fmt.Printf("\nBeads (%d):\n", len(groupIssues))
@@ -392,45 +309,17 @@ func runWorkCreate(cmd *cobra.Command, args []string) error {
 		printSessionCreatedNotification(sessionResult.SessionName)
 	}
 
-	// If --auto, run the full automated workflow
-	if flagAutoRun {
-		fmt.Println("\nRunning automated workflow...")
-		result, err := work.RunWorkAuto(ctx, proj, workID, os.Stdout)
-		if err != nil {
-			return fmt.Errorf("failed to run automated workflow: %w", err)
-		}
-		if result.OrchestratorSpawned {
-			fmt.Println("Orchestrator spawned in zellij tab.")
-		}
-		// Ensure control plane is running (handles scheduled tasks like PR feedback polling)
-		// Note: InitializeSession spawns control plane for new sessions, but we call
-		// EnsureControlPlane for existing sessions that might have a dead control plane
-		if err := control.EnsureControlPlane(ctx, proj); err != nil {
-			fmt.Printf("Warning: failed to ensure control plane: %v\n", err)
-		}
-		fmt.Println("Switch to the zellij session to monitor progress.")
-		return nil
-	}
-
-	// Spawn the orchestrator for this work
-	fmt.Println("\nSpawning orchestrator...")
-	if err := claude.SpawnWorkOrchestrator(ctx, workID, proj.Config.Project.Name, worktreePath, workerName, os.Stdout); err != nil {
-		fmt.Printf("Warning: failed to spawn orchestrator: %v\n", err)
-		fmt.Println("You can start it manually with: co run")
-	} else {
-		fmt.Println("Orchestrator is running in zellij tab.")
-	}
-
-	// Ensure control plane is running (handles scheduled tasks like PR feedback polling)
-	// Note: InitializeSession spawns control plane for new sessions, but we call
-	// EnsureControlPlane for existing sessions that might have a dead control plane
+	// Ensure control plane is running (handles worktree creation, orchestrator spawning, etc.)
 	if err := control.EnsureControlPlane(ctx, proj); err != nil {
 		fmt.Printf("Warning: failed to ensure control plane: %v\n", err)
 	}
 
-	fmt.Printf("\nNext steps:\n")
-	fmt.Printf("  cd %s\n", workID)
-	fmt.Printf("  co run               # Execute tasks\n")
+	if flagAutoRun {
+		fmt.Println("\nAutomated workflow will start once the control plane creates the worktree.")
+	}
+
+	fmt.Printf("\nThe control plane will create the worktree and start the orchestrator.\n")
+	fmt.Printf("Switch to the zellij session to monitor progress.\n")
 
 	return nil
 }

cmd/work_import_pr.go

@@ -0,0 +1,130 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/newhook/co/internal/control"
+	"github.com/newhook/co/internal/github"
+	"github.com/newhook/co/internal/project"
+	"github.com/newhook/co/internal/work"
+	"github.com/spf13/cobra"
+)
+
+var workImportPRCmd = &cobra.Command{
+	Use:   "import-pr <pr-url>",
+	Short: "Import a PR into a work unit",
+	Long: `Create a work unit from an existing GitHub pull request.
+
+This command fetches the PR's branch, creates a worktree, and sets up the work
+for further development or review. The PR's branch becomes the work's feature branch.
+
+A bead is automatically created from the PR metadata to track the work in the beads system.
+
+Examples:
+  co work import-pr https://github.com/owner/repo/pull/123
+  co work import-pr https://github.com/owner/repo/pull/123 --branch custom-branch-name`,
+	Args: cobra.ExactArgs(1),
+	RunE: runWorkImportPR,
+}
+
+var flagImportPRBranch string
+
+func init() {
+	workImportPRCmd.Flags().StringVar(&flagImportPRBranch, "branch", "", "override the local branch name (default: use PR's branch name)")
+	workCmd.AddCommand(workImportPRCmd)
+}
+
+func runWorkImportPR(cmd *cobra.Command, args []string) error {
+	ctx := GetContext()
+
+	// Find project
+	proj, err := project.Find(ctx, "")
+	if err != nil {
+		return err
+	}
+	defer proj.Close()
+
+	prURL := args[0]
+
+	// Create GitHub client and PR importer
+	ghClient := github.NewClient()
+	importer := work.NewPRImporter(ghClient)
+
+	// Fetch PR metadata first (user needs to see PR info)
+	fmt.Printf("Fetching PR metadata from %s...\n", prURL)
+	metadata, err := importer.FetchPRMetadata(ctx, prURL, "")
+	if err != nil {
+		return fmt.Errorf("failed to fetch PR metadata: %w", err)
+	}
+
+	fmt.Printf("PR #%d: %s\n", metadata.Number, metadata.Title)
+	fmt.Printf("Author: %s\n", metadata.Author)
+	fmt.Printf("State: %s\n", metadata.State)
+	fmt.Printf("Branch: %s -> %s\n", metadata.HeadRefName, metadata.BaseRefName)
+
+	// Check if PR is still open
+	if metadata.State != "OPEN" {
+		fmt.Printf("Warning: PR is %s\n", metadata.State)
+	}
+
+	// Determine branch name
+	branchName := flagImportPRBranch
+	if branchName == "" {
+		branchName = metadata.HeadRefName
+	}
+
+	// Create a bead from PR metadata (user needs feedback on bead creation)
+	fmt.Printf("\nCreating bead from PR metadata...\n")
+	beadResult, err := importer.CreateBeadFromPR(ctx, metadata, &work.CreateBeadOptions{
+		BeadsDir:     proj.BeadsPath(),
+		SkipIfExists: true,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create bead: %w", err)
+	}
+	if beadResult.Created {
+		fmt.Printf("Created bead: %s\n", beadResult.BeadID)
+	} else {
+		fmt.Printf("Bead already exists: %s (%s)\n", beadResult.BeadID, beadResult.SkipReason)
+	}
+	rootIssueID := beadResult.BeadID
+
+	// Schedule PR import via control plane (handles worktree, git, mise)
+	fmt.Printf("\nScheduling PR import...\n")
+	result, err := work.ImportPRAsync(ctx, proj, work.ImportPRAsyncOptions{
+		PRURL:       prURL,
+		BranchName:  branchName,
+		RootIssueID: rootIssueID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to schedule PR import: %w", err)
+	}
+
+	fmt.Printf("\nCreated work: %s\n", result.WorkID)
+	if result.WorkerName != "" {
+		fmt.Printf("Worker: %s\n", result.WorkerName)
+	}
+	fmt.Printf("Branch: %s\n", result.BranchName)
+	fmt.Printf("PR URL: %s\n", prURL)
+	fmt.Printf("\nWorktree setup is in progress via control plane.\n")
+
+	// Initialize zellij session and spawn control plane if new session
+	sessionResult, err := control.InitializeSession(ctx, proj)
+	if err != nil {
+		fmt.Printf("Warning: failed to initialize zellij session: %v\n", err)
+	} else if sessionResult.SessionCreated {
+		// Display notification for new session
+		printSessionCreatedNotification(sessionResult.SessionName)
+	}
+
+	// Ensure control plane is running to process the import task
+	if err := control.EnsureControlPlane(ctx, proj); err != nil {
+		fmt.Printf("Warning: failed to ensure control plane: %v\n", err)
+	}
+
+	fmt.Printf("\nNext steps:\n")
+	fmt.Printf("  cd %s\n", result.WorkID)
+	fmt.Printf("  co run               # Execute tasks (after worktree is ready)\n")
+
+	return nil
+}