refactor: remove container selection, extract helpers, add logging and tests

- Remove container picker UI, keyboard, settings, and ContainerOptions entirely — yt-dlp handles container format automatically - Remove Container field from DownloadJob struct and related references - Extract buildFormatList helper to deduplicate format list building between handleURLInput and handleBackStep - Break runDownload into trackDownloadProgress + handleDownloadCompletion - Add HTTP timeout (30s) to downloadFile helper - Log applyQuota and AddHistory errors instead of discarding - Log nil stepState warnings in all handler entry points - Include job.StderrLog in download failure messages - Add tests for buildFormatList, findFormatByID, determineMediaType, and buildFormatString - Update README with search, DRM fallback, and quota feature docs - Remove unused strings import from keyboard.go
2026-06-25 23:41:26 +03:30
parent dfe946a41b
commit 37e7f918d0
10 changed files with 266 additions and 296 deletions
--- a/internal/handler/download.go
+++ b/internal/handler/download.go
@@ -28,7 +28,6 @@ type stepState struct {
 	SecondaryID   string
 	SecondaryFmt  *domain.Format
 	Language      string
-	Container     string
 	FormatIDs     []string
 	PlaylistState *domain.PlaylistState
 	SearchState   *domain.SearchState
@@ -138,30 +137,21 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string,

 	h.sendMediaPreview(ctx, chatID, info)

-	// Build unified format list: video formats first, then audio
-	deduped := dl.DeduplicateResolutions(info.Formats)
-	var videoFmts, audioFmts []domain.Format
-	for _, f := range deduped {
-		if f.HasVideo {
-			videoFmts = append(videoFmts, f)
+	ss.FormatIDs, _ = buildFormatList(info.Formats)
+	h.setUserStepState(chatID, ss)
+	labels := make([]string, len(ss.FormatIDs))
+	for i, id := range ss.FormatIDs {
+		if id == "best" {
+			labels[i] = "Best"
 		} else {
-			audioFmts = append(audioFmts, f)
+			f := findFormatByID(info.Formats, id)
+			if f != nil {
+				labels[i] = formatLabelForDisplay(*f)
+			} else {
+				labels[i] = id
+			}
 		}
 	}
-	ss.FormatIDs = make([]string, 0, len(deduped)+1)
-	labels := make([]string, 0, len(deduped)+1)
-	ss.FormatIDs = append(ss.FormatIDs, "best")
-	labels = append(labels, "Best")
-	for _, f := range videoFmts {
-		ss.FormatIDs = append(ss.FormatIDs, f.ID)
-		labels = append(labels, formatLabelForDisplay(f))
-	}
-	for _, f := range audioFmts {
-		ss.FormatIDs = append(ss.FormatIDs, f.ID)
-		labels = append(labels, formatLabelForDisplay(f))
-	}
-
-	h.setUserStepState(chatID, ss)
 	kb := formatKeyboard("format_", labels, ss.FormatIDs)
 	h.sendWithKB(ctx, chatID, "Select format:", kb)
 }
@@ -170,6 +160,7 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string,
 func (h *Handler) handleFormatSelection(ctx context.Context, chatID int64, msgID int, userID int64, formatID string) {
 	ss := h.getUserStepState(chatID)
 	if ss == nil {
+		slog.Warn("nil stepState in handleFormatSelection", "chat_id", chatID, "user_id", userID)
 		return
 	}
 	ss.MainFormatID = formatID
@@ -183,7 +174,6 @@ func (h *Handler) handleFormatSelection(ctx context.Context, chatID int64, msgID
 		h.advanceToLanguage(ctx, chatID, msgID, userID, ss)
 		return
 	}
-
 	for i := range ss.MediaInfo.Formats {
 		f := &ss.MediaInfo.Formats[i]
 		if f.ID == formatID {
@@ -255,6 +245,7 @@ func (h *Handler) showSecondaryFormats(ctx context.Context, chatID int64, msgID
 func (h *Handler) handleSecondarySelection(ctx context.Context, chatID int64, msgID int, userID int64, formatID string) {
 	ss := h.getUserStepState(chatID)
 	if ss == nil {
+		slog.Warn("nil stepState in handleSecondarySelection", "chat_id", chatID, "user_id", userID)
 		return
 	}

@@ -280,7 +271,7 @@ func (h *Handler) handleSecondarySelection(ctx context.Context, chatID int64, ms
 	h.advanceToLanguage(ctx, chatID, msgID, userID, ss)
 }

-// advanceToLanguage moves to language selection, or skips to container if N/A.
+// advanceToLanguage moves to language selection, or starts the download if only one language.
 func (h *Handler) advanceToLanguage(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) {
 	ss.Step = 2
 	h.setUserStepState(chatID, ss)
@@ -290,17 +281,7 @@ func (h *Handler) advanceToLanguage(ctx context.Context, chatID int64, msgID int
 		h.editText(ctx, chatID, msgID, "Select language:", &kb)
 		return
 	}
-	h.advanceToContainer(ctx, chatID, msgID, userID, ss)
-}
-
-// advanceToContainer shows the container picker.
-func (h *Handler) advanceToContainer(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) {
-	ss.Step = 3
-	h.setUserStepState(chatID, ss)
-
-	containers := dl.ContainerOptions(ss.determineMediaType())
-	kb := containerKeyboard(containers)
-	h.editText(ctx, chatID, msgID, "Select container format:", &kb)
+	h.startDownload(ctx, chatID, msgID, userID, ss)
 }

 func (ss *stepState) determineMediaType() domain.MediaType {
@@ -327,20 +308,15 @@ func (ss *stepState) buildFormatString() string {
 func (h *Handler) handleLanguageSelection(ctx context.Context, chatID int64, msgID int, userID int64, language string) {
 	ss := h.getUserStepState(chatID)
 	if ss == nil {
+		slog.Warn("nil stepState in handleLanguageSelection", "chat_id", chatID, "user_id", userID)
 		return
 	}
 	ss.Language = language
-	h.advanceToContainer(ctx, chatID, msgID, userID, ss)
+	h.startDownload(ctx, chatID, msgID, userID, ss)
 }

-// handleContainerSelection processes the container choice and starts the download.
-func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, msgID int, userID int64, container string) {
-	ss := h.getUserStepState(chatID)
-	if ss == nil {
-		return
-	}
-	ss.Container = container
-
+// startDownload creates the download job and begins the download process.
+func (h *Handler) startDownload(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) {
 	mediaType := ss.determineMediaType()
 	formatString := ss.buildFormatString()

@@ -368,7 +344,6 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
 		MediaType:      mediaType,
 		SelectedFormat: format,
 		FormatString:   formatString,
-		Container:      container,
 		Language:       ss.Language,
 		Status:         domain.StatusDownloading,
 		Title:          ss.MediaInfo.Title,
@@ -392,7 +367,6 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
 		}
 	}()

-	// Acquire semaphore slot or wait for cancellation
 	select {
 	case h.downloadSem <- struct{}{}:
 	case <-ctx.Done():
@@ -414,21 +388,33 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {

 	updates, err := h.YtDlp.StartDownload(job, h.downloadDir)
 	if err != nil {
-		h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil)
+		errMsg := err.Error()
+		if job.StderrLog != "" {
+			errMsg += ": " + job.StderrLog
+		}
+		h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", errMsg), nil)
 		job.Status = domain.StatusFailed
-		h.activeJobsMu.Lock()
-		delete(h.activeJobs, job.UserID)
-		h.activeJobsMu.Unlock()
+		h.clearActiveJob(job.UserID)
 		return
 	}

+	h.trackDownloadProgress(ctx, job, updates)
+
+	h.clearActiveJob(job.UserID)
+
+	if job.Status == domain.StatusCompleted {
+		h.handleDownloadCompletion(ctx, job)
+	}
+}
+
+// trackDownloadProgress reads progress updates from yt-dlp and updates the Telegram message.
+func (h *Handler) trackDownloadProgress(ctx context.Context, job *domain.DownloadJob, updates <-chan *domain.DownloadJob) {
 	lastProgress := 0.0
-loop:
 	for {
 		select {
 		case upd, ok := <-updates:
 			if !ok {
-				break loop
+				return
 			}
 			slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status)

@@ -441,23 +427,21 @@ loop:

 			if upd.Status == domain.StatusCancelled {
 				h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
-				h.activeJobsMu.Lock()
-				delete(h.activeJobs, job.UserID)
-				h.activeJobsMu.Unlock()
 				return
 			}

 			if upd.Status == domain.StatusFailed {
 				h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
-				h.activeJobsMu.Lock()
-				delete(h.activeJobs, job.UserID)
-				h.activeJobsMu.Unlock()
+				return
+			}
+
+			if upd.Status == domain.StatusCompleted {
 				return
 			}

 			if upd.Progress > 0 && lastProgress == 0 {
 				lastProgress = upd.Progress
-			} else if upd.Progress-lastProgress < domain.ProgressThreshold && upd.Status != domain.StatusCompleted {
+			} else if upd.Progress-lastProgress < domain.ProgressThreshold {
 				continue
 			} else {
 				lastProgress = upd.Progress
@@ -476,57 +460,54 @@ loop:
 			dl.CancelDownload(job)
 			for range updates {
 			}
-			break loop
+			return
 		}
 	}
+}

-	h.activeJobsMu.Lock()
-	delete(h.activeJobs, job.UserID)
-	h.activeJobsMu.Unlock()
-
-	h.cancelFuncsMu.Lock()
-	delete(h.cancelFuncs, job.UserID)
-	h.cancelFuncsMu.Unlock()
-
-	if job.Status == domain.StatusCompleted {
-		pattern := filepath.Join(h.downloadDir, job.ID+"_*")
-		files, err := filepath.Glob(pattern)
-		if err == nil && len(files) > 0 {
-			f := files[0]
-			fileInfo, err := os.Stat(f)
-			if err == nil {
-				if fileInfo.Size() > h.maxFileSize {
-					h.editText(ctx, job.ChatID, job.MessageID,
-						fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileInfo.Size()), formatBytes(h.maxFileSize)), nil)
-					os.Remove(f)
-					return
-				}
-
-				if !job.QuotaApplied && job.FileSize > 0 {
-					job.QuotaApplied = true
-					h.applyQuota(job.UserID, job.FileSize)
-				}
-
-				formatID := job.SelectedFormat.ID
-				h.sendDocument(ctx, job.ChatID, f, job.Title)
-				os.Remove(f)
-
-				_ = h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{
-					URL:       job.URL,
-					Title:     job.Title,
-					MediaType: job.MediaType.String(),
-					FormatID:  formatID,
-					Container: job.Container,
-					FileSize:  job.FileSize,
-					Status:    "completed",
-				})
-
-				h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", nil)
-				return
-			}
-		}
+// handleDownloadCompletion sends the downloaded file and records history.
+func (h *Handler) handleDownloadCompletion(ctx context.Context, job *domain.DownloadJob) {
+	pattern := filepath.Join(h.downloadDir, job.ID+"_*")
+	files, err := filepath.Glob(pattern)
+	if err != nil || len(files) == 0 {
 		h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil)
+		return
 	}
+
+	f := files[0]
+	fileInfo, err := os.Stat(f)
+	if err != nil {
+		h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil)
+		return
+	}
+
+	if fileInfo.Size() > h.maxFileSize {
+		h.editText(ctx, job.ChatID, job.MessageID,
+			fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileInfo.Size()), formatBytes(h.maxFileSize)), nil)
+		os.Remove(f)
+		return
+	}
+
+	if !job.QuotaApplied && job.FileSize > 0 {
+		job.QuotaApplied = true
+		h.applyQuota(job.UserID, job.FileSize)
+	}
+
+	h.sendDocument(ctx, job.ChatID, f, job.Title)
+	os.Remove(f)
+
+	if err := h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{
+		URL:       job.URL,
+		Title:     job.Title,
+		MediaType: job.MediaType.String(),
+		FormatID:  job.SelectedFormat.ID,
+		FileSize:  job.FileSize,
+		Status:    "completed",
+	}); err != nil {
+		slog.Error("add history failed", "user_id", job.UserID, "error", err)
+	}
+
+	h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", nil)
 }

 func (h *Handler) sendDocument(ctx context.Context, chatID int64, filePath, title string) {
@@ -579,6 +560,7 @@ func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, c
 	defer h.answerCb(ctx, cbID)
 	ss := h.getUserStepState(chatID)
 	if ss == nil {
+		slog.Warn("nil stepState in handleBackStep", "chat_id", chatID, "user_id", userID)
 		h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
 		return
 	}
@@ -589,26 +571,20 @@ func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, c
 		h.clearStepState(chatID)
 		h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
 	case 0:
-		// Re-show format list
-		deduped := dl.DeduplicateResolutions(ss.MediaInfo.Formats)
-		var videoFmts, audioFmts []domain.Format
-		for _, f := range deduped {
-			if f.HasVideo {
-				videoFmts = append(videoFmts, f)
+		ss.FormatIDs, _ = buildFormatList(ss.MediaInfo.Formats)
+		labels := make([]string, len(ss.FormatIDs))
+		for i, id := range ss.FormatIDs {
+			if id == "best" {
+				labels[i] = "Best"
 			} else {
-				audioFmts = append(audioFmts, f)
+				f := findFormatByID(ss.MediaInfo.Formats, id)
+				if f != nil {
+					labels[i] = formatLabelForDisplay(*f)
+				} else {
+					labels[i] = id
+				}
 			}
 		}
-		ss.FormatIDs = []string{"best"}
-		labels := []string{"Best"}
-		for _, f := range videoFmts {
-			ss.FormatIDs = append(ss.FormatIDs, f.ID)
-			labels = append(labels, formatLabelForDisplay(f))
-		}
-		for _, f := range audioFmts {
-			ss.FormatIDs = append(ss.FormatIDs, f.ID)
-			labels = append(labels, formatLabelForDisplay(f))
-		}
 		h.setUserStepState(chatID, ss)
 		kb := formatKeyboard("format_", labels, ss.FormatIDs)
 		h.editText(ctx, chatID, msgID, "Select format:", &kb)
@@ -644,6 +620,34 @@ func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID
 		"Send me the URL to download:", nil)
 }

+// buildFormatList returns a deduplicated format ID list with "best" first, then video, then audio.
+func buildFormatList(formats []domain.Format) ([]string, error) {
+	deduped := dl.DeduplicateResolutions(formats)
+	ids := make([]string, 0, len(deduped)+1)
+	ids = append(ids, "best")
+	for _, f := range deduped {
+		if f.HasVideo {
+			ids = append(ids, f.ID)
+		}
+	}
+	for _, f := range deduped {
+		if !f.HasVideo {
+			ids = append(ids, f.ID)
+		}
+	}
+	return ids, nil
+}
+
+// findFormatByID returns the format with the given ID from the list, or nil.
+func findFormatByID(formats []domain.Format, id string) *domain.Format {
+	for i := range formats {
+		if formats[i].ID == id {
+			return &formats[i]
+		}
+	}
+	return nil
+}
+
 func formatLabelForDisplay(f domain.Format) string {
 	if f.Resolution != "" {
 		label := f.Resolution
@@ -721,7 +725,8 @@ func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *doma

 // downloadFile downloads a URL to a temp file in the given directory and returns the path.
 func downloadFile(url, dir string) (string, error) {
-	resp, err := http.Get(url)
+	client := &http.Client{Timeout: 30 * time.Second}
+	resp, err := client.Get(url)
 	if err != nil {
 		return "", fmt.Errorf("http get: %w", err)
 	}