Skip to content

Commit 6925994

Browse files
luispaterluispater
committed
Support multiple auth files per instance.
Refactor the architecture to handle multiple auth files for each instance, enabling better scalability and customization. Updates include changes to configurations, handlers, and the initialization logic while maintaining backward compatibility. Removed unused `AuthDownload` endpoint and improved error handling throughout.
1 parent f7e2bda commit 6925994

File tree

8 files changed

+218
-295
lines changed

8 files changed

+218
-295
lines changed

‎internal/api/handlers.go‎

Lines changed: 71 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,13 @@ var ScreenshotMutex sync.Mutex
2727
// APIHandlers contains the handlers for API endpoints
2828
type APIHandlers struct {
2929
queue *RequestQueue
30-
pages map[string]*chrome.Page
30+
pages map[string][]*chrome.Page
3131
debug bool
3232
appConfig *config.AppConfig
3333
}
3434

3535
// NewAPIHandlers creates a new API handlers instance
36-
func NewAPIHandlers(appConfig *config.AppConfig, queue *RequestQueue, pages map[string]*chrome.Page, debug bool) *APIHandlers {
36+
func NewAPIHandlers(appConfig *config.AppConfig, queue *RequestQueue, pages map[string][]*chrome.Page, debug bool) *APIHandlers {
3737
return &APIHandlers{
3838
queue: queue,
3939
pages: pages,
@@ -120,10 +120,21 @@ func (h *APIHandlers) TakeScreenshot(c *gin.Context) {
120120
c.Status(http.StatusNotFound)
121121
return
122122
}
123+
instanceIndex, ok := c.GetQuery("index")
124+
if !ok {
125+
c.Status(http.StatusNotFound)
126+
return
127+
}
128+
index, err := strconv.ParseUint(instanceIndex, 10, 64)
129+
if err != nil {
130+
c.Status(http.StatusNotFound)
131+
return
132+
}
133+
123134
if page, hasKey := h.pages[instanceName]; hasKey {
124135
var buf []byte
125-
pageCtx := page.GetContext()
126-
if err := chromedp.Run(pageCtx, chromedp.CaptureScreenshot(&buf)); err != nil {
136+
pageCtx := page[index].GetContext()
137+
if err = chromedp.Run(pageCtx, chromedp.CaptureScreenshot(&buf)); err != nil {
127138
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to take screenshot: %v", err), "code": 500})
128139
return
129140
}
@@ -140,20 +151,16 @@ func (h *APIHandlers) TakeScreenshot(c *gin.Context) {
140151
func (h *APIHandlers) ProxyPac(c *gin.Context) {
141152
c.Status(http.StatusOK)
142153
c.Header("Content-Type", "application/x-ns-proxy-autoconfig")
143-
if h.appConfig.InstanceAlone {
144-
index, ok := c.GetQuery("index")
145-
if ok {
146-
i, err := strconv.ParseUint(index, 10, 32)
147-
if err != nil {
148-
_, _ = c.Writer.Write([]byte(`function FindProxyForURL(url, host) {return "PROXY 127.0.0.1:3120";}`))
149-
return
150-
}
151-
152-
port := 3120 + i
153-
_, _ = c.Writer.Write([]byte(fmt.Sprintf(`function FindProxyForURL(url, host) {return "PROXY 127.0.0.1:%d";}`, port)))
154+
index, ok := c.GetQuery("index")
155+
if ok {
156+
i, err := strconv.ParseUint(index, 10, 32)
157+
if err != nil {
158+
_, _ = c.Writer.Write([]byte(`function FindProxyForURL(url, host) {return "PROXY 127.0.0.1:3120";}`))
159+
return
154160
}
155-
} else {
156-
_, _ = c.Writer.Write([]byte(`function FindProxyForURL(url, host) {return "PROXY 127.0.0.1:3120";}`))
161+
162+
port := 3120 + i
163+
_, _ = c.Writer.Write([]byte(fmt.Sprintf(`function FindProxyForURL(url, host) {return "PROXY 127.0.0.1:%d";}`, port)))
157164
}
158165
}
159166

@@ -200,9 +207,26 @@ func (h *APIHandlers) ChatCompletions(c *gin.Context) {
200207
instanceIndex = i
201208
}
202209
}
203-
if page, ok := h.pages[instanceName]; ok {
204-
defer page.RequestMutex.Unlock()
205-
page.RequestMutex.Lock()
210+
var page *chrome.Page
211+
defer func() {
212+
if page != nil {
213+
page.RequestMutex.Unlock()
214+
}
215+
}()
216+
217+
if pages, ok := h.pages[instanceName]; ok {
218+
locked := false
219+
for i := 0; i < len(pages); i++ {
220+
page = pages[i]
221+
if page.RequestMutex.TryLock() {
222+
locked = true
223+
break
224+
}
225+
}
226+
if !locked {
227+
page = pages[0]
228+
page.RequestMutex.Lock()
229+
}
206230
} else {
207231
c.JSON(http.StatusNotFound, ErrorResponse{
208232
Error: ErrorDetail{
@@ -224,6 +248,7 @@ func (h *APIHandlers) ChatCompletions(c *gin.Context) {
224248
CreatedAt: time.Now(),
225249
Context: c,
226250
InstanceName: instanceName,
251+
Page: page,
227252
}
228253

229254
// Add a task to queue
@@ -268,8 +293,7 @@ func (h *APIHandlers) ChatCompletions(c *gin.Context) {
268293
}
269294
}
270295

271-
func (h *APIHandlers) handleContextCanceled(instanceIndex int) {
272-
page := h.pages[h.appConfig.Instance[instanceIndex].Name]
296+
func (h *APIHandlers) handleContextCanceled(instanceIndex int, page *chrome.Page) {
273297
r, err := runner.NewRunnerManager(h.appConfig.Instance[instanceIndex], page, h.debug, false)
274298
if err != nil {
275299
log.Error(err)
@@ -305,7 +329,7 @@ func (h *APIHandlers) handleNonStreamingResponse(instanceIndex int, c *gin.Conte
305329
case <-c.Request.Context().Done():
306330
if c.Request.Context().Err().Error() == "context canceled" {
307331
log.Debugf("Client disconnected: %v", c.Request.Context().Err())
308-
h.handleContextCanceled(instanceIndex)
332+
h.handleContextCanceled(instanceIndex, response.Page)
309333
response.Runner.Abort()
310334
}
311335
return
@@ -356,7 +380,7 @@ func (h *APIHandlers) handleStreamingResponse(instanceIndex int, c *gin.Context,
356380
case <-c.Request.Context().Done():
357381
if c.Request.Context().Err().Error() == "context canceled" {
358382
log.Debugf("Client disconnected: %v", c.Request.Context().Err())
359-
h.handleContextCanceled(instanceIndex)
383+
h.handleContextCanceled(instanceIndex, response.Page)
360384
response.Runner.Abort()
361385
}
362386
return
@@ -416,10 +440,20 @@ func (h *APIHandlers) BrowserReload(c *gin.Context) {
416440
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("Page for instance '%s' not found", instanceName), "code": 404})
417441
return
418442
}
443+
instanceIndex, ok := c.GetQuery("index")
444+
if !ok {
445+
c.Status(http.StatusNotFound)
446+
return
447+
}
448+
index, err := strconv.ParseUint(instanceIndex, 10, 64)
449+
if err != nil {
450+
c.Status(http.StatusNotFound)
451+
return
452+
}
419453

420454
// Navigate to the instance URL
421-
pageCtx := page.GetContext()
422-
if err := chromedp.Run(pageCtx, chromedp.Navigate(instanceConfig.URL)); err != nil {
455+
pageCtx := page[index].GetContext()
456+
if err = chromedp.Run(pageCtx, chromedp.Navigate(instanceConfig.URL)); err != nil {
423457
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to reload page: %v", err), "code": 500})
424458
return
425459
}
@@ -432,6 +466,7 @@ func (h *APIHandlers) BrowserReload(c *gin.Context) {
432466
func (h *APIHandlers) AuthUpload(c *gin.Context) {
433467
var requestData struct {
434468
Name string `json:"name" binding:"required"`
469+
Index *int `json:"index" binding:"required"`
435470
Token string `json:"token" binding:"required"`
436471
Auth string `json:"auth" binding:"required"`
437472
}
@@ -469,7 +504,7 @@ func (h *APIHandlers) AuthUpload(c *gin.Context) {
469504
}
470505

471506
// Ensure the auth directory exists
472-
authAbsPath, err := filepath.Abs(instanceConfig.Auth.File)
507+
authAbsPath, err := filepath.Abs(instanceConfig.Auth.Files[*requestData.Index])
473508
if err != nil {
474509
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to get absolute path: %v", err), "code": 500})
475510
return
@@ -484,29 +519,29 @@ func (h *APIHandlers) AuthUpload(c *gin.Context) {
484519
}
485520

486521
// Write the auth info to file
487-
if err = os.WriteFile(instanceConfig.Auth.File, []byte(requestData.Auth), 0644); err != nil {
522+
if err = os.WriteFile(instanceConfig.Auth.Files[*requestData.Index], []byte(requestData.Auth), 0644); err != nil {
488523
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to write auth file: %v", err), "code": 500})
489524
return
490525
}
491526

492527
// Get the page for this instance
493-
page, exists := h.pages[instanceConfig.Name]
528+
pages, exists := h.pages[instanceConfig.Name]
494529
if !exists {
495530
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("Page for instance '%s' not found", instanceConfig.Name), "code": 404})
496531
return
497532
}
498533

499534
// Navigate to the instance URL
500-
pageCtx := page.GetContext()
535+
pageCtx := pages[*requestData.Index].GetContext()
501536
err = chrome.ClearCookies(pageCtx)
502537
if err != nil {
503538
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to load auth info: %v", err), "code": 500})
504539
return
505540
}
506-
page.Close()
541+
pages[*requestData.Index].Close()
507542

508543
pageLoaded := func() {
509-
r, errNewRunnerManager := runner.NewRunnerManager(*instanceConfig, page, h.appConfig.Debug, false) // Pass pageCtx
544+
r, errNewRunnerManager := runner.NewRunnerManager(*instanceConfig, pages[*requestData.Index], h.appConfig.Debug, false) // Pass pageCtx
510545
if errNewRunnerManager != nil {
511546
log.Error(errNewRunnerManager)
512547
}
@@ -517,72 +552,23 @@ func (h *APIHandlers) AuthUpload(c *gin.Context) {
517552
log.Debugf("all of the init system rules are executed.")
518553
}
519554

520-
page, err = h.pages[instanceConfig.Name].GetBrowserManager().NewPage(instanceConfig.URL, instanceConfig.Adapter, instanceConfig.Auth.File, pageLoaded)
555+
p, err := h.pages[instanceConfig.Name][*requestData.Index].GetBrowserManager().NewPage(instanceConfig.URL, instanceConfig.Adapter, instanceConfig.Auth.Files[*requestData.Index], pageLoaded)
521556
if err != nil {
522557
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to load auth info: %v", err), "code": 500})
523558
return
524559
}
525-
h.pages[instanceConfig.Name] = page
526-
pageCtx = page.GetContext()
560+
h.pages[instanceConfig.Name][*requestData.Index] = p
561+
pageCtx = p.GetContext()
527562

528563
if err = chromedp.Run(pageCtx, chromedp.Navigate(instanceConfig.URL)); err != nil {
529564
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to reload page: %v", err), "code": 500})
530565
return
531566
}
532567

533-
log.Debugf("Successfully uploaded auth info for instance '%s' to file: %s", requestData.Name, instanceConfig.Auth.File)
568+
log.Debugf("Successfully uploaded auth info for instance '%s' to file: %s", requestData.Name, instanceConfig.Auth.Files[*requestData.Index])
534569
c.JSON(http.StatusOK, gin.H{"message": "Auth info uploaded successfully"})
535570
}
536571

537-
// AuthDownload handles the /v1/auth/download endpoint
538-
func (h *APIHandlers) AuthDownload(c *gin.Context) {
539-
instanceName, ok := c.GetQuery("name")
540-
if !ok || instanceName == "" {
541-
c.JSON(http.StatusBadRequest, gin.H{"error": "Missing required parameter: name", "code": 400})
542-
return
543-
}
544-
545-
// Find the instance configuration
546-
var instanceConfig *config.AppConfigInstance
547-
for i := range h.appConfig.Instance {
548-
if h.appConfig.Instance[i].Name == instanceName {
549-
instanceConfig = &h.appConfig.Instance[i]
550-
break
551-
}
552-
}
553-
554-
if instanceConfig == nil {
555-
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("Instance '%s' not found", instanceName), "code": 404})
556-
return
557-
}
558-
559-
// Check if auth file exists
560-
if _, err := os.Stat(instanceConfig.Auth.File); os.IsNotExist(err) {
561-
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("Auth file not found for instance '%s'", instanceName), "code": 404})
562-
return
563-
}
564-
565-
// Read the auth file
566-
authData, err := os.ReadFile(instanceConfig.Auth.File)
567-
if err != nil {
568-
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to read auth file: %v", err), "code": 500})
569-
return
570-
}
571-
572-
// Validate the auth JSON format
573-
var authInfo chrome.AuthInfo
574-
if err = json.Unmarshal(authData, &authInfo); err != nil {
575-
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Invalid auth file format: %v", err), "code": 500})
576-
return
577-
}
578-
579-
log.Debugf("Successfully downloaded auth info for instance '%s' from file: %s", instanceName, instanceConfig.Auth.File)
580-
581-
// Return the auth info as JSON
582-
c.Header("Content-Type", "application/json")
583-
c.JSON(http.StatusOK, authInfo)
584-
}
585-
586572
// AuthUploadPage handles the GET /v1/auth/upload endpoint to serve the upload page
587573
func (h *APIHandlers) AuthUploadPage(c *gin.Context) {
588574
// Serve the HTML content

‎internal/api/models.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package api
22

33
import (
44
"github.com/gin-gonic/gin"
5+
"github.com/luispater/anyAIProxyAPI/internal/browser/chrome"
56
"github.com/luispater/anyAIProxyAPI/internal/runner"
67
"time"
78
)
@@ -32,6 +33,7 @@ type RequestTask struct {
3233
CreatedAt time.Time `json:"created_at"`
3334
Context *gin.Context `json:"context"`
3435
InstanceName string `json:"instance_name"`
36+
Page *chrome.Page `json:"page"`
3537
}
3638

3739
// TaskResponse represents the response from processing a task
@@ -41,4 +43,5 @@ type TaskResponse struct {
4143
Stream chan string `json:"-"`
4244
Error error `json:"error,omitempty"`
4345
Runner *runner.RunnerManager
46+
Page *chrome.Page `json:"page"`
4447
}

‎internal/api/processor.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,13 @@ import (
1717

1818
// ChatProcessor implements TaskProcessor interface
1919
type ChatProcessor struct {
20-
pages map[string]*chrome.Page
20+
pages map[string][]*chrome.Page
2121
debug bool
2222
appConfig *config.AppConfig
2323
}
2424

2525
// NewChatProcessor creates a new chat processor
26-
func NewChatProcessor(appConfig *config.AppConfig, pages map[string]*chrome.Page, debug bool) *ChatProcessor {
26+
func NewChatProcessor(appConfig *config.AppConfig, pages map[string][]*chrome.Page, debug bool) *ChatProcessor {
2727
return &ChatProcessor{
2828
pages: pages,
2929
debug: debug,
@@ -70,15 +70,14 @@ func (cp *ChatProcessor) processNonStreamingTask(appConfigInstance config.AppCon
7070
errChannel := make(chan error)
7171
streamChan := make(chan string, 100)
7272

73-
page := cp.pages[appConfigInstance.Name]
74-
r, errNewRunnerManager := runner.NewRunnerManager(appConfigInstance, page, cp.debug, false)
73+
r, errNewRunnerManager := runner.NewRunnerManager(appConfigInstance, task.Page, cp.debug, false)
7574
go func() {
7675
if errNewRunnerManager != nil {
7776
log.Debug(errNewRunnerManager)
7877
return
7978
}
8079
r.SetVariable("REQUEST", task.Request, "string")
81-
r.SetVariable("PAGE", page, "ptr")
80+
r.SetVariable("PAGE", task.Page, "ptr")
8281
r.SetVariable("PAGE-DATA-CHANNEL", channel, "ptr")
8382
err := r.Run("chat_completions")
8483
if err != nil {
@@ -93,7 +92,7 @@ func (cp *ChatProcessor) processNonStreamingTask(appConfigInstance config.AppCon
9392
defer close(streamChan)
9493
for !done {
9594
select {
96-
case pageError := <-page.Error:
95+
case pageError := <-task.Page.Error:
9796
streamChan <- fmt.Sprintf(`{"error": "%v"}`, pageError)
9897
return
9998
case err := <-errChannel:
@@ -153,6 +152,7 @@ func (cp *ChatProcessor) processNonStreamingTask(appConfigInstance config.AppCon
153152
Success: true,
154153
Stream: streamChan,
155154
Runner: r,
155+
Page: task.Page,
156156
}
157157
}
158158

@@ -163,15 +163,14 @@ func (cp *ChatProcessor) processStreamingTask(appConfigInstance config.AppConfig
163163
channel := make(chan *adapter.AdapterResponse)
164164
errChannel := make(chan error)
165165

166-
page := cp.pages[appConfigInstance.Name]
167-
r, errNewRunnerManager := runner.NewRunnerManager(appConfigInstance, page, cp.debug, false)
166+
r, errNewRunnerManager := runner.NewRunnerManager(appConfigInstance, task.Page, cp.debug, false)
168167
go func() {
169168
if errNewRunnerManager != nil {
170169
log.Debug(errNewRunnerManager)
171170
return
172171
}
173172
r.SetVariable("REQUEST", task.Request, "string")
174-
r.SetVariable("PAGE", page, "ptr")
173+
r.SetVariable("PAGE", task.Page, "ptr")
175174
r.SetVariable("PAGE-DATA-CHANNEL", channel, "ptr")
176175
err := r.Run("chat_completions")
177176
if err != nil {
@@ -200,7 +199,7 @@ func (cp *ChatProcessor) processStreamingTask(appConfigInstance config.AppConfig
200199
var done bool
201200
for !done {
202201
select {
203-
case pageError := <-page.Error:
202+
case pageError := <-task.Page.Error:
204203
streamChan <- fmt.Sprintf(`{"error": "%v"}`, pageError)
205204
return
206205
case err := <-errChannel:
@@ -248,6 +247,7 @@ func (cp *ChatProcessor) processStreamingTask(appConfigInstance config.AppConfig
248247
Success: true,
249248
Stream: streamChan,
250249
Runner: r,
250+
Page: task.Page,
251251
}
252252
}
253253

0 commit comments

Comments
 (0)