local_client.go

// Custom code. Not generated by Stainless.
package stagehand

import (
	"context"
	"fmt"
	"os"
	"sync"

	"github.com/browserbase/stagehand-go/v3/internal/requestconfig"
	"github.com/browserbase/stagehand-go/v3/lib/local"
	"github.com/browserbase/stagehand-go/v3/option"
)

type localServerOption struct {
	mu      sync.Mutex
	manager *local.ServerManager
	initErr error
}

func newLocalServerOption() *localServerOption {
	return &localServerOption{}
}

func (o *localServerOption) Apply(cfg *requestconfig.RequestConfig) error {
	var modelAPIKey string
	var browserbaseAPIKey string
	var browserbaseProjectID string
	if cfg != nil {
		modelAPIKey = cfg.ModelAPIKey
		browserbaseAPIKey = cfg.BrowserbaseAPIKey
		browserbaseProjectID = cfg.BrowserbaseProjectID
	}
	if modelAPIKey == "" {
		if key := os.Getenv("MODEL_API_KEY"); key != "" {
			modelAPIKey = key
		}
		if modelAPIKey != "" && cfg != nil {
			if err := option.WithModelAPIKey(modelAPIKey).Apply(cfg); err != nil {
				return err
			}
		}
	}
	if browserbaseAPIKey == "" {
		if key := os.Getenv("BROWSERBASE_API_KEY"); key != "" {
			browserbaseAPIKey = key
			if cfg != nil {
				if err := option.WithBrowserbaseAPIKey(browserbaseAPIKey).Apply(cfg); err != nil {
					return err
				}
			}
		}
	}
	if browserbaseProjectID == "" {
		if key := os.Getenv("BROWSERBASE_PROJECT_ID"); key != "" {
			browserbaseProjectID = key
			if cfg != nil {
				if err := option.WithBrowserbaseProjectID(browserbaseProjectID).Apply(cfg); err != nil {
					return err
				}
			}
		}
	}
	if modelAPIKey == "" {
		return fmt.Errorf("MODEL_API_KEY is required for local mode")
	}

	manager, err := o.ensureManager()
	if err != nil {
		return err
	}
	if modelAPIKey != "" {
		manager.SetModelAPIKey(modelAPIKey)
	}
	if browserbaseAPIKey != "" {
		manager.SetBrowserbaseAPIKey(browserbaseAPIKey)
	}
	if browserbaseProjectID != "" {
		manager.SetBrowserbaseProjectID(browserbaseProjectID)
	}

	ctx := cfg.Context
	if ctx == nil {
		ctx = context.Background()
	}

	baseURL, err := manager.EnsureRunning(ctx)
	if err != nil {
		return err
	}

	return option.WithBaseURL(baseURL).Apply(cfg)
}

func (o *localServerOption) Close() error {
	o.mu.Lock()
	manager := o.manager
	o.mu.Unlock()

	if manager == nil {
		return nil
	}

	return manager.Close()
}

func (o *localServerOption) ensureManager() (*local.ServerManager, error) {
	o.mu.Lock()
	defer o.mu.Unlock()

	if o.manager != nil || o.initErr != nil {
		return o.manager, o.initErr
	}

	manager, err := local.NewServerManager()
	if err != nil {
		o.initErr = err
		return nil, err
	}

	o.manager = manager
	return manager, nil
}

func serverModeFromOptions(opts []option.RequestOption) string {
	mode := "remote"
	for _, opt := range opts {
		if serverOpt, ok := opt.(option.ServerOption); ok && serverOpt.Server() != "" {
			mode = serverOpt.Server()
		}
	}
	return mode
}

// Close shuts down any local mode processes associated with this client.
func (c Client) Close() error {
	var firstErr error
	for _, opt := range c.Options {
		if closer, ok := opt.(interface{ Close() error }); ok {
			if err := closer.Close(); err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	return firstErr
}