fix(ytmusic): batch playlists and cancel searches

external/ytmusic/cookie_provider.go

@@ -31,12 +31,15 @@ const (
 type cookieBase struct {
 	browser    string
 	fetchFn    func(browser string) ([]playlist.PlaylistInfo, error)
+	resolveFn  func(pageURL string, start, count int, browser ...string) ([]playlist.Track, error)
 	mu         sync.Mutex
 	playlists  []playlist.PlaylistInfo
 	trackCache map[string][]playlist.Track
 	disk       *ytCache
 }
 
+const cookiePlaylistBatchSize = 100
+
 func newCookieBase(browser string) *cookieBase {
 	return &cookieBase{
 		browser:    browser,
@@ -123,9 +126,21 @@ func (b *cookieBase) fetchTracks(target string) ([]playlist.Track, error) {
 	}
 	b.mu.Unlock()
 
-	tracks, err := resolve.ResolveYTDLBatch(target, 0, 0, b.browser)
-	if err != nil {
-		return nil, fmt.Errorf("ytmusic: resolve playlist tracks: %w", err)
+	resolveBatch := b.resolveFn
+	if resolveBatch == nil {
+		resolveBatch = resolve.ResolveYTDLBatch
+	}
+	var tracks []playlist.Track
+	for start := 0; ; {
+		batch, err := resolveBatch(target, start, cookiePlaylistBatchSize, b.browser)
+		if err != nil {
+			return nil, fmt.Errorf("ytmusic: resolve playlist tracks: %w", err)
+		}
+		tracks = append(tracks, batch...)
+		if len(batch) < cookiePlaylistBatchSize {
+			break
+		}
+		start += len(batch)
 	}
 
 	b.mu.Lock()
@@ -283,7 +298,7 @@ func (p *CookieProvider) Tracks(playlistID string) ([]playlist.Track, error) {
 }
 
 // SearchTracks performs a search query using yt-dlp's ytsearch: protocol.
-func (p *CookieProvider) SearchTracks(_ context.Context, query string, limit int) ([]playlist.Track, error) {
+func (p *CookieProvider) SearchTracks(ctx context.Context, query string, limit int) ([]playlist.Track, error) {
 	q := strings.TrimSpace(query)
 	if q == "" {
 		return nil, nil
@@ -291,7 +306,7 @@ func (p *CookieProvider) SearchTracks(_ context.Context, query string, limit int
 	if limit <= 0 {
 		limit = 10
 	}
-	tracks, err := resolve.ResolveYTDLBatch(fmt.Sprintf("ytsearch%d:%s", limit, q), 0, 0, p.base.browser)
+	tracks, err := resolve.ResolveYTDLBatchContext(ctx, fmt.Sprintf("ytsearch%d:%s", limit, q), 0, 0, p.base.browser)
 	if err != nil {
 		return nil, fmt.Errorf("ytmusic: search tracks: %w", err)
 	}

external/ytmusic/cookie_provider_test.go

@@ -3,9 +3,11 @@ package ytmusic
 import (
 	"context"
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
+	"slices"
 	"strings"
 	"testing"
 
@@ -253,6 +255,60 @@ func TestCookieProviderTracksCaching(t *testing.T) {
 	}
 }
 
+func TestCookieProviderTracksLoadsInBatches(t *testing.T) {
+	t.Setenv("HOME", t.TempDir())
+	base := newCookieBase("firefox")
+	var starts []int
+	base.resolveFn = func(_ string, start, count int, browser ...string) ([]playlist.Track, error) {
+		starts = append(starts, start)
+		if count != cookiePlaylistBatchSize {
+			t.Fatalf("count = %d, want %d", count, cookiePlaylistBatchSize)
+		}
+		if len(browser) != 1 || browser[0] != "firefox" {
+			t.Fatalf("browser = %v, want firefox", browser)
+		}
+		n := cookiePlaylistBatchSize
+		if start > 0 {
+			n = 20
+		}
+		tracks := make([]playlist.Track, n)
+		for i := range tracks {
+			tracks[i].Title = fmt.Sprintf("Track %d", start+i)
+		}
+		return tracks, nil
+	}
+
+	tracks, err := (&CookieProvider{base: base, kind: KindMusic}).Tracks("PL123")
+	if err != nil {
+		t.Fatalf("Tracks() error: %v", err)
+	}
+	if len(tracks) != 120 {
+		t.Fatalf("tracks = %d, want 120", len(tracks))
+	}
+	if !slices.Equal(starts, []int{0, cookiePlaylistBatchSize}) {
+		t.Fatalf("batch starts = %v, want [0 %d]", starts, cookiePlaylistBatchSize)
+	}
+}
+
+func TestCookieProviderSearchTracksHonorsCancellation(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("skipping Unix shell script test on Windows")
+	}
+	tmpDir := t.TempDir()
+	fakeYTDL := filepath.Join(tmpDir, "yt-dlp")
+	if err := os.WriteFile(fakeYTDL, []byte("#!/bin/sh\nexec sleep 10\n"), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err := NewCookieProvider("chrome", KindMusic).SearchTracks(ctx, "query", 10)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("SearchTracks() error = %v, want context.Canceled", err)
+	}
+}
+
 func TestNewCookieProvidersDoesNotMutateGlobalCookies(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("skipping Unix shell script test on Windows")

resolve/resolve.go

@@ -611,11 +611,19 @@ func resolveYouTube(pageURL string) ([]playlist.Track, error) {
 // If an optional browser is provided, cookies from that browser are used;
 // otherwise, the globally configured yt-dlp cookies browser is used.
 func ResolveYTDLBatch(pageURL string, start, count int, browser ...string) ([]playlist.Track, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	return ResolveYTDLBatchContext(ctx, pageURL, start, count, browser...)
+}
+
+// ResolveYTDLBatchContext is ResolveYTDLBatch with caller-controlled
+// cancellation and timeout.
+func ResolveYTDLBatchContext(ctx context.Context, pageURL string, start, count int, browser ...string) ([]playlist.Track, error) {
 	end := 0
 	if count > 0 {
 		end = start + count
 	}
-	return resolveYTDLRange(pageURL, start, end, browser...)
+	return resolveYTDLRangeContext(ctx, pageURL, start, end, browser...)
 }
 
 // resolveYTDL uses yt-dlp --flat-playlist to quickly enumerate tracks.
@@ -630,13 +638,16 @@ func resolveYTDL(pageURL string, maxItems ...int) ([]playlist.Track, error) {
 }
 
 func resolveYTDLRange(pageURL string, start, end int, browser ...string) ([]playlist.Track, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	return resolveYTDLRangeContext(ctx, pageURL, start, end, browser...)
+}
+
+func resolveYTDLRangeContext(ctx context.Context, pageURL string, start, end int, browser ...string) ([]playlist.Track, error) {
 	if _, err := exec.LookPath("yt-dlp"); err != nil {
 		return nil, fmt.Errorf("yt-dlp not found in PATH — see https://github.com/yt-dlp/yt-dlp#installation")
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-	defer cancel()
-
 	args := []string{"--flat-playlist", "-j", "--socket-timeout", "15"}
 	b := ""
 	if len(browser) > 0 {
@@ -656,12 +667,13 @@ func resolveYTDLRange(pageURL string, start, end int, browser ...string) ([]play
 	}
 	args = append(args, pageURL)
 	cmd := exec.CommandContext(ctx, "yt-dlp", args...)
+	cmd.WaitDelay = 3 * time.Second
 	var stderr strings.Builder
 	cmd.Stderr = &stderr
 	stdout, err := cmd.Output()
 	if err != nil {
-		if ctx.Err() == context.DeadlineExceeded {
-			return nil, fmt.Errorf("yt-dlp: timed out resolving %s (30s)", pageURL)
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, fmt.Errorf("yt-dlp: resolve %s: %w", pageURL, ctxErr)
 		}
 		msg := strings.TrimSpace(stderr.String())
 		if msg != "" {