Add context support to Playwright browser initialization

Refactored Playwright initialization to ensure context propagation. Updated `NewPlayWrightBrowser` and related methods to accept `context.Context` for better cancellation and timeout handling. Improved error resilience and concurrency during browser setup.

cmd/browser/pkg/browser/flags.go

@@ -43,7 +43,7 @@ var Flags = BrowserFlags{
 	},
 }
 
-func FromCommand(_ context.Context, cmd *cli.Command) (extractor.Browser, error) {
+func FromCommand(ctx context.Context, cmd *cli.Command) (extractor.Browser, error) {
 	var opts extractor.PlayWrightBrowserOptions
 
 	if ua := cmd.String("user-agent"); ua != "" {
@@ -72,5 +72,5 @@ func FromCommand(_ context.Context, cmd *cli.Command) (extractor.Browser, error)
 
 	opts.ShowBrowser = cmd.Bool("visible")
 
-	return extractor.NewPlayWrightBrowser(opts)
+	return extractor.NewPlayWrightBrowser(ctx, opts)
 }

playwright.go

@@ -90,7 +90,7 @@ func playwrightCookieToCookie(cookie playwright.Cookie) Cookie {
 	}
 }
 
-func NewPlayWrightBrowser(opts ...PlayWrightBrowserOptions) (Browser, error) {
+func NewPlayWrightBrowser(ctx context.Context, opts ...PlayWrightBrowserOptions) (Browser, error) {
 	var thirtySeconds = 30 * time.Second
 	opt := PlayWrightBrowserOptions{
 		UserAgent:               "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:142.0) Gecko/20100101 Firefox/142.0",
@@ -131,126 +131,161 @@ func NewPlayWrightBrowser(opts ...PlayWrightBrowserOptions) (Browser, error) {
 		opt.ShowBrowser = o.ShowBrowser
 	}
 
-	pw, err := playwright.Run()
+	// Check if context is already done
-
+	if err := ctx.Err(); err != nil {
-	if err != nil {
-		err = playwright.Install()
-
-		if err != nil {
-			return nil, err
-		}
-
-		pw, err = playwright.Run()
-
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	var bt playwright.BrowserType
-
-	switch opt.Browser {
-	case PlayWrightBrowserSelectionChromium:
-		bt = pw.Chromium
-		if opt.PlayWrightServerAddress == "" {
-			opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_CHROMIUM")
-		}
-
-	case PlayWrightBrowserSelectionFirefox:
-		bt = pw.Firefox
-		if opt.PlayWrightServerAddress == "" {
-			opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_FIREFOX")
-		}
-
-	case PlayWrightBrowserSelectionWebKit:
-		bt = pw.WebKit
-		if opt.PlayWrightServerAddress == "" {
-			opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_WEBKIT")
-		}
-
-	default:
-		return nil, ErrInvalidBrowserSelection
-	}
-	var browser playwright.Browser
-
-	var launch = true
-	if opt.PlayWrightServerAddress != "" && !opt.UseLocalOnly {
-		launch = false
-		slog.Info("connecting to playwright server", "address", opt.PlayWrightServerAddress)
-		var timeout float64 = 30000
-		browser, err = bt.Connect(opt.PlayWrightServerAddress, playwright.BrowserTypeConnectOptions{Timeout: &timeout})
-
-		if err != nil {
-			if opt.DontLaunchOnConnectFailure {
-				return nil, err
-			}
-			slog.Warn("failed to connect to playwright server, launching local browser", "err", err)
-			launch = true
-		}
-	}
-
-	if launch {
-		browser, err = bt.Launch(playwright.BrowserTypeLaunchOptions{
-			Headless: playwright.Bool(!opt.ShowBrowser),
-		})
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	var viewport *playwright.Size
-	if opt.Dimensions.Width > 0 && opt.Dimensions.Height > 0 {
-		viewport = &playwright.Size{
-			Width:  opt.Dimensions.Width,
-			Height: opt.Dimensions.Height,
-		}
-	}
-
-	var scheme *playwright.ColorScheme
-
-	if opt.DarkMode {
-		scheme = playwright.ColorSchemeDark
-	} else {
-		scheme = playwright.ColorSchemeNoPreference
-	}
-
-	c, err := browser.NewContext(playwright.BrowserNewContextOptions{
-		UserAgent:   playwright.String(opt.UserAgent),
-		Viewport:    viewport,
-		ColorScheme: scheme,
-	})
-	if err != nil {
 		return nil, err
 	}
 
-	if opt.CookieJar != nil {
+	type browserResult struct {
-		cookies, err := opt.CookieJar.GetAll()
+		browser Browser
-		if err != nil {
+		err     error
-			return nil, fmt.Errorf("error getting cookies from cookie jar: %w", err)
-		}
-
-		pwCookies := make([]playwright.OptionalCookie, len(cookies))
-
-		for i, cookie := range cookies {
-			pwCookies[i] = cookieToPlaywrightOptionalCookie(cookie)
-		}
-
-		err = c.AddCookies(pwCookies)
-
-		if err != nil {
-			return nil, fmt.Errorf("error adding cookies to browser: %w", err)
-		}
 	}
 
-	return playWrightBrowser{
+	// Create a channel for the result
-		pw:         pw,
+	resultCh := make(chan browserResult, 1)
-		browser:    browser,
+
-		userAgent:  opt.UserAgent,
+	// Launch browser initialization in a separate goroutine
-		timeout:    *opt.Timeout,
+	go func() {
-		cookieJar:  opt.CookieJar,
+		pw, err := playwright.Run()
-		ctx:        c,
+
-		serverAddr: opt.PlayWrightServerAddress,
+		if err != nil {
-	}, nil
+			err = playwright.Install()
+
+			if err != nil {
+				resultCh <- browserResult{nil, err}
+				return
+			}
+
+			pw, err = playwright.Run()
+
+			if err != nil {
+				resultCh <- browserResult{nil, err}
+				return
+			}
+		}
+
+		var bt playwright.BrowserType
+
+		switch opt.Browser {
+		case PlayWrightBrowserSelectionChromium:
+			bt = pw.Chromium
+			if opt.PlayWrightServerAddress == "" {
+				opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_CHROMIUM")
+			}
+
+		case PlayWrightBrowserSelectionFirefox:
+			bt = pw.Firefox
+			if opt.PlayWrightServerAddress == "" {
+				opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_FIREFOX")
+			}
+
+		case PlayWrightBrowserSelectionWebKit:
+			bt = pw.WebKit
+			if opt.PlayWrightServerAddress == "" {
+				opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_WEBKIT")
+			}
+
+		default:
+			resultCh <- browserResult{nil, ErrInvalidBrowserSelection}
+			return
+		}
+		var browser playwright.Browser
+
+		var launch = true
+		if opt.PlayWrightServerAddress != "" && !opt.UseLocalOnly {
+			launch = false
+			slog.Info("connecting to playwright server", "address", opt.PlayWrightServerAddress)
+			var timeout float64 = 30000
+			browser, err = bt.Connect(opt.PlayWrightServerAddress, playwright.BrowserTypeConnectOptions{Timeout: &timeout})
+
+			if err != nil {
+				if opt.DontLaunchOnConnectFailure {
+					resultCh <- browserResult{nil, err}
+					return
+				}
+				slog.Warn("failed to connect to playwright server, launching local browser", "err", err)
+				launch = true
+			}
+		}
+
+		if launch {
+			browser, err = bt.Launch(playwright.BrowserTypeLaunchOptions{
+				Headless: playwright.Bool(!opt.ShowBrowser),
+			})
+			if err != nil {
+				resultCh <- browserResult{nil, err}
+				return
+			}
+		}
+
+		var viewport *playwright.Size
+		if opt.Dimensions.Width > 0 && opt.Dimensions.Height > 0 {
+			viewport = &playwright.Size{
+				Width:  opt.Dimensions.Width,
+				Height: opt.Dimensions.Height,
+			}
+		}
+
+		var scheme *playwright.ColorScheme
+
+		if opt.DarkMode {
+			scheme = playwright.ColorSchemeDark
+		} else {
+			scheme = playwright.ColorSchemeNoPreference
+		}
+
+		c, err := browser.NewContext(playwright.BrowserNewContextOptions{
+			UserAgent:   playwright.String(opt.UserAgent),
+			Viewport:    viewport,
+			ColorScheme: scheme,
+		})
+		if err != nil {
+			resultCh <- browserResult{nil, err}
+			return
+		}
+
+		if opt.CookieJar != nil {
+			cookies, err := opt.CookieJar.GetAll()
+			if err != nil {
+				resultCh <- browserResult{nil, fmt.Errorf("error getting cookies from cookie jar: %w", err)}
+				return
+			}
+
+			pwCookies := make([]playwright.OptionalCookie, len(cookies))
+
+			for i, cookie := range cookies {
+				pwCookies[i] = cookieToPlaywrightOptionalCookie(cookie)
+			}
+
+			err = c.AddCookies(pwCookies)
+
+			if err != nil {
+				resultCh <- browserResult{nil, fmt.Errorf("error adding cookies to browser: %w", err)}
+				return
+			}
+		}
+
+		resultCh <- browserResult{
+			browser: playWrightBrowser{
+				pw:         pw,
+				browser:    browser,
+				userAgent:  opt.UserAgent,
+				timeout:    *opt.Timeout,
+				cookieJar:  opt.CookieJar,
+				ctx:        c,
+				serverAddr: opt.PlayWrightServerAddress,
+			},
+			err: nil,
+		}
+	}()
+
+	// Wait for either context cancellation or browser initialization completion
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case result := <-resultCh:
+		return result.browser, result.err
+	}
 }
 
 func (b playWrightBrowser) updateCookies(_ context.Context, page playwright.Page) error {
@@ -337,7 +372,7 @@ func deferClose(cl io.Closer) {
 }
 
 func Screenshot(ctx context.Context, target string, timeout time.Duration) ([]byte, error) {
-	browser, err := NewPlayWrightBrowser(PlayWrightBrowserOptions{
+	browser, err := NewPlayWrightBrowser(ctx, PlayWrightBrowserOptions{
 		Timeout: &timeout,
 	})