free-Chatgpt-api/agent.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	activeConvMu sync.Mutex
	activeConvID string
)

func getActiveConversationID() string {
	activeConvMu.Lock()
	defer activeConvMu.Unlock()
	return activeConvID
}

func setActiveConversationID(id string) {
	activeConvMu.Lock()
	activeConvID = id
	activeConvMu.Unlock()
}

func clearActiveConversationID() {
	activeConvMu.Lock()
	activeConvID = ""
	activeConvMu.Unlock()
}

func Start() {
	loadConfig()
	LoadCookiesFromFile()

	promptFlag := flag.String("prompt", "", "Prompt to send to ChatGPT after the extension connects")
	onceFlag := flag.Bool("once", false, "Send the prompt, print the response, then exit")
	timeoutFlag := flag.Duration("timeout", cfg.Timeout(), "Maximum time to wait for extension/response")
	flag.Parse()

	http.HandleFunc("/", handleWS)
	http.HandleFunc("/health", handleHealth)
	http.HandleFunc("/api/ext/callback", handleCallback)
	http.HandleFunc("/api/ext/reload", func(w http.ResponseWriter, r *http.Request) {
		extMu.Lock()
		conn := extConn
		extMu.Unlock()
		if conn == nil {
			writeJSONStatus(w, http.StatusServiceUnavailable, map[string]any{"ok": false, "error": "extension not connected"})
			return
		}
		err := conn.WriteJSON(WSMessage{
			Method: "reload_extension",
		})
		if err != nil {
			writeJSONStatus(w, http.StatusInternalServerError, map[string]any{"ok": false, "error": err.Error()})
			return
		}
		writeJSON(w, map[string]any{"ok": true, "message": "reload request sent"})
	})
	http.HandleFunc("/api/conversation", func(w http.ResponseWriter, r *http.Request) {
		convID := r.URL.Query().Get("conversation_id")
		if convID == "" {
			writeJSONStatus(w, http.StatusBadRequest, map[string]any{"ok": false, "error": "conversation_id required"})
			return
		}
		token, err := getSessionToken()
		if err != nil {
			writeJSONStatus(w, http.StatusServiceUnavailable, map[string]any{"ok": false, "error": err.Error()})
			return
		}
		res, err := callChatGPTAPI(apiCallParams{
			URL:    "https://chatgpt.com/backend-api/conversation/" + convID,
			Method: "GET",
			Headers: map[string]string{
				"Authorization": "Bearer " + token,
			},
		}, cfg.APITimeout())
		if err != nil {
			writeJSONStatus(w, http.StatusInternalServerError, map[string]any{"ok": false, "error": err.Error()})
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(res.Status)
		w.Write([]byte(res.Body))
	})
	http.HandleFunc("/api/chat", handleChat)
	http.HandleFunc("/api/chat/bulk", handleChatBulk)
	http.HandleFunc("/api/chat/edit", handleChatEdit)
	http.HandleFunc("/api/chatgpt/test", handleChat)
	http.HandleFunc("/api/sniffs", handleSniffs)
	http.HandleFunc("/api/download", handleDownload)
	http.HandleFunc("/v1/chat/completions", handleOpenAIChat)
	http.HandleFunc("/api/cookies/status", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, map[string]any{
			"ok":              true,
			"has_cookies":     HasCookies(),
			"last_sync":       GetLastCookieSync().Format(time.RFC3339),
			"cookie_header_len": len(GetCachedCookieHeader()),
			"extension_connected": isExtensionConnected(),
		})
	})

	if *promptFlag != "" {
		go runPromptFromCLI(*promptFlag, *timeoutFlag, *onceFlag)
	}

	log.Printf("ChatGPT Agent listening on http://%s", cfg.ListenAddr)
	log.Println("Reload the ChatGPT Browser Bridge extension if it is not connected.")
	log.Println("Endpoints: GET/POST /api/chat, POST /api/chat/bulk, POST /api/chat/edit, GET /api/sniffs, GET /health")
	if err := http.ListenAndServe(cfg.ListenAddr, nil); err != nil {
		log.Fatal(err)
	}
}

func runPromptFromCLI(prompt string, timeout time.Duration, once bool) {
	log.Println("Waiting for extension API bridge connection...")
	if !waitForExtension(timeout) {
		log.Println("Extension did not connect before timeout")
		if once {
			os.Exit(1)
		}
		return
	}

	log.Printf("Sending prompt: %s", prompt)
	text, _, err := sendChat(prompt)
	if err != nil {
		log.Printf("ChatGPT request failed: %v", err)
		if once {
			os.Exit(1)
		}
		return
	}

	fmt.Println(text)
	if once {
		os.Exit(0)
	}
}