diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 8923585..fec341d 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -72,6 +72,7 @@ func main() { // sandbox, it re-runs THIS binary, which self-restricts and execs the real // command. Must run first; a normal launch falls straight through. sandbox.MaybeExecConfined() + llm.UserAgent = userAgent() var ( workspace string doInit bool diff --git a/internal/llm/client.go b/internal/llm/client.go index 72b0156..c341168 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -183,6 +183,18 @@ func (c *deadlineConn) Read(b []byte) (int, error) { // transport DOES bound the connect/handshake/response-header phases so a dead // connection (e.g. after the laptop sleeps) can't wedge a request forever — // short of that, a stuck socket read ignores context cancellation. +// UserAgent names this program on every request to the model server, so a +// proxy in front of it (LLMTray, a gateway) can tell it from any other Go +// client — they all said "Go-http-client/1.1". Set once at startup; empty +// leaves Go's default. +var UserAgent string + +func setUserAgent(req *http.Request) { + if UserAgent != "" { + req.Header.Set("User-Agent", UserAgent) + } +} + func NewOpenAIClient(c config.LLM) *OpenAIClient { idle := 90 * time.Second if c.IdleTimeoutSeconds > 0 { @@ -567,6 +579,7 @@ func (c *OpenAIClient) send(ctx context.Context, buf []byte, reqCompl *int) (Mes } req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "text/event-stream") + setUserAgent(req) if strings.Contains(c.baseURL, "openrouter.ai") { // OpenRouter uses these to attribute traffic to the app (rankings, some // free-tier access). Harmless elsewhere, so only set for OpenRouter. diff --git a/internal/llm/detect.go b/internal/llm/detect.go index 91dcd39..8fc5035 100644 --- a/internal/llm/detect.go +++ b/internal/llm/detect.go @@ -25,6 +25,7 @@ func ListModels(ctx context.Context, baseURL, apiKey string, hc *http.Client) ([ if err != nil { return nil, err } + setUserAgent(req) if apiKey != "" { req.Header.Set("Authorization", "Bearer "+apiKey) } @@ -70,6 +71,7 @@ func DetectModelContext(ctx context.Context, baseURL, apiKey, model string, hc * if err != nil { return 0 } + setUserAgent(req) if apiKey != "" { req.Header.Set("Authorization", "Bearer "+apiKey) } @@ -123,6 +125,7 @@ func ListLMStudioModels(ctx context.Context, baseURL string, hc *http.Client) ([ if err != nil { return nil, err } + setUserAgent(req) resp, err := hc.Do(req) if err != nil { return nil, err @@ -154,6 +157,7 @@ func DetectContextWindow(ctx context.Context, baseURL, model string, hc *http.Cl if err != nil { return 0 } + setUserAgent(req) resp, err := hc.Do(req) if err != nil { return 0 diff --git a/internal/llm/useragent_test.go b/internal/llm/useragent_test.go new file mode 100644 index 0000000..0036621 --- /dev/null +++ b/internal/llm/useragent_test.go @@ -0,0 +1,63 @@ +package llm + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/ipsupport-llc/ipsupport-code/internal/config" +) + +// Every request to the model server names this program: a proxy in front of +// it saw only "Go-http-client/1.1", the same as any other Go tool. +func TestEveryModelRequestCarriesTheUserAgent(t *testing.T) { + defer func(ua string) { UserAgent = ua }(UserAgent) + UserAgent = "ipsupport-code/9.9.9" + + var mu sync.Mutex + var seen []string // "path user-agent", one per request + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + seen = append(seen, r.URL.Path+" "+r.Header.Get("User-Agent")) + mu.Unlock() + switch r.URL.Path { + case "/v1/chat/completions": + io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}`) + default: + io.WriteString(w, `{"data":[{"id":"m","state":"loaded","loaded_context_length":8192}]}`) + } + })) + defer srv.Close() + + ctx := context.Background() + base := srv.URL + "/v1" + calls := map[string]func(){ + "ListModels": func() { _, _ = ListModels(ctx, base, "", srv.Client()) }, + "DetectModelContext": func() { _ = DetectModelContext(ctx, base, "", "m", srv.Client()) }, + "ListLMStudioModels": func() { _, _ = ListLMStudioModels(ctx, base, srv.Client()) }, + "DetectContextWindow": func() { _ = DetectContextWindow(ctx, base, "m", srv.Client()) }, + "Chat": func() { + _, _ = NewOpenAIClient(config.LLM{BaseURL: base, Model: "m"}).Chat(ctx, []Message{User("hi")}, nil) + }, + } + for name, call := range calls { + mu.Lock() + seen = nil + mu.Unlock() + call() + mu.Lock() + if len(seen) == 0 { + t.Errorf("%s made no request", name) + } + for _, s := range seen { + if !strings.HasSuffix(s, " "+UserAgent) { + t.Errorf("%s: request %q, want User-Agent %q", name, s, UserAgent) + } + } + mu.Unlock() + } +}