From d7c3c0aaaaf404de811ebcbf969b2bd3823058fc Mon Sep 17 00:00:00 2001 From: sususu Date: Fri, 2 Oct 2026 22:12:20 +0800 Subject: [PATCH] feat(antigravity): trace upstream connection reuse in debug logs --- internal/logging/global_logger.go | 6 +- .../antigravity_connection_trace_test.go | 74 ++++++++ .../runtime/executor/antigravity_executor.go | 2 +- .../executor/antigravity_executor_auth.go | 5 +- .../executor/antigravity_executor_credits.go | 2 +- .../executor/antigravity_executor_execute.go | 4 +- .../executor/antigravity_executor_stream.go | 2 +- .../executor/antigravity_executor_tokens.go | 2 +- .../helps/antigravity_connection_trace.go | 77 ++++++++ ...igravity_connection_trace_redirect_test.go | 123 +++++++++++++ .../antigravity_connection_trace_test.go | 173 ++++++++++++++++++ 11 files changed, 461 insertions(+), 9 deletions(-) create mode 100644 internal/runtime/executor/antigravity_connection_trace_test.go create mode 100644 internal/runtime/executor/helps/antigravity_connection_trace.go create mode 100644 internal/runtime/executor/helps/antigravity_connection_trace_redirect_test.go create mode 100644 internal/runtime/executor/helps/antigravity_connection_trace_test.go diff --git a/internal/logging/global_logger.go b/internal/logging/global_logger.go index 524fc4247..928fc2501 100644 --- a/internal/logging/global_logger.go +++ b/internal/logging/global_logger.go @@ -56,13 +56,17 @@ var logFieldOrder = []string{ "plugin_id", "plugin_name", "source_id", "version", "active_version", "retired_version", "overwritten", "mode", "budget", "level", "original_mode", "original_value", "min", "max", "clamped_to", "error", - "credential", "auth_id", "connection", "proxy_scheme", "remote_transport", + "credential", "auth_id", "auth_index", "connection", "proxy_scheme", "remote_transport", + "operation", "upstream_host", "reused", "was_idle", "idle_time", "media_session_id", "call_id", "peer", "state", "reason", } var quotedLogFields = map[string]struct{}{ "credential": {}, "auth_id": {}, + "auth_index": {}, + "upstream_host": {}, + "operation": {}, "connection": {}, "proxy_scheme": {}, "remote_transport": {}, diff --git a/internal/runtime/executor/antigravity_connection_trace_test.go b/internal/runtime/executor/antigravity_connection_trace_test.go new file mode 100644 index 000000000..370101b15 --- /dev/null +++ b/internal/runtime/executor/antigravity_connection_trace_test.go @@ -0,0 +1,74 @@ +package executor + +import ( + "context" + "io" + "net/http" + "net/http/httptrace" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +func TestAntigravityProjectDiscoveryConnectionTrace(t *testing.T) { + logger := log.StandardLogger() + oldLevel := logger.GetLevel() + oldHooks := logger.ReplaceHooks(make(log.LevelHooks)) + oldOutput := logger.Out + logger.SetLevel(log.DebugLevel) + logger.SetOutput(io.Discard) + hook := logtest.NewLocal(logger) + t.Cleanup(func() { + logger.SetLevel(oldLevel) + logger.ReplaceHooks(oldHooks) + logger.SetOutput(oldOutput) + }) + auth := &cliproxyauth.Auth{Index: "discovery-index"} + requests := 0 + ctx := logging.WithRequestID(context.Background(), "discovery-request-id") + ctx = context.WithValue(ctx, "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) { + requests++ + trace := httptrace.ContextClientTrace(req.Context()) + if trace == nil || trace.GotConn == nil { + t.Fatal("discovery request missing connection trace") + } + trace.GotConn(httptrace.GotConnInfo{}) + var body string + switch req.URL.Path { + case "/v1internal:loadCodeAssist": + body = `{"allowedTiers":[{"id":"free-tier","isDefault":true}]}` + case "/v1internal:onboardUser": + body = `{"done":true,"response":{"cloudaicompanionProject":{"id":"trace-project"}}}` + default: + t.Fatalf("unexpected discovery request: %s", req.URL.Path) + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + })) + executor := &AntigravityExecutor{} + projectID, errFetch := executor.fetchAntigravityProjectID(ctx, auth, "secret-token") + if errFetch != nil { + t.Fatal(errFetch) + } + if projectID != "trace-project" || requests != 2 { + t.Fatalf("project=%q requests=%d", projectID, requests) + } + var entries []*log.Entry + for _, entry := range hook.AllEntries() { + if entry.Message == "antigravity upstream connection acquired" { + entries = append(entries, entry) + } + } + if len(entries) != 2 { + t.Fatalf("connection entries=%d, want 2", len(entries)) + } + for i, host := range []string{"cloudcode-pa.googleapis.com", "daily-cloudcode-pa.googleapis.com"} { + fields := entries[i].Data + if fields["operation"] != "project_discovery" || fields["request_id"] != "discovery-request-id" || fields["auth_index"] != "discovery-index" || fields["upstream_host"] != host { + t.Fatalf("unexpected trace fields: %v", fields) + } + } +} diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 7a28df22a..9a53090c4 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -969,7 +969,7 @@ func (e *AntigravityExecutor) HttpRequest(ctx context.Context, auth *cliproxyaut } httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - return httpClient.Do(httpReq) + return helps.WithAntigravityHTTPClientTrace(httpClient, auth, "http_request").Do(httpReq) } // SupportsApplyPatch reports the actual executor contract, independent of its provider name. diff --git a/internal/runtime/executor/antigravity_executor_auth.go b/internal/runtime/executor/antigravity_executor_auth.go index cdba4a072..6571930f3 100644 --- a/internal/runtime/executor/antigravity_executor_auth.go +++ b/internal/runtime/executor/antigravity_executor_auth.go @@ -167,7 +167,7 @@ func (e *AntigravityExecutor) refreshTokenSingleFlight(ctx context.Context, auth httpReq.Header.Set("User-Agent", "Go-http-client/2.0") httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "oauth_refresh").Do(httpReq) if errDo != nil { return nil, errDo } @@ -240,7 +240,8 @@ func (e *AntigravityExecutor) fetchAntigravityProjectID(ctx context.Context, aut acquisitionCtx, cancelAcquisition := context.WithTimeout(ctx, antigravityCredentialAcquisitionTimeout) defer cancelAcquisition() httpClient := newAntigravityHTTPClient(acquisitionCtx, e.cfg, auth, 0) - projectID, errFetch := sdkAuth.FetchAntigravityProjectID(acquisitionCtx, token, httpClient) + tracedClient := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "project_discovery") + projectID, errFetch := sdkAuth.FetchAntigravityProjectID(acquisitionCtx, token, tracedClient) if errFetch != nil { return "", errFetch } diff --git a/internal/runtime/executor/antigravity_executor_credits.go b/internal/runtime/executor/antigravity_executor_credits.go index 0d931bee7..516453971 100644 --- a/internal/runtime/executor/antigravity_executor_credits.go +++ b/internal/runtime/executor/antigravity_executor_credits.go @@ -539,7 +539,7 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalanceForTask(ctx context httpReq.Header.Set("User-Agent", userAgent) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "credits_query").Do(httpReq) if errDo != nil { log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo) return diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go index 36401c3fc..3c7d552a1 100644 --- a/internal/runtime/executor/antigravity_executor_execute.go +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -130,7 +130,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au return resp, err } - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "generate").Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { @@ -340,7 +340,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * return resp, err } - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "generate").Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index 480835128..0d2d6243d 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -129,7 +129,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya err = errReq return nil, err } - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "generate_stream").Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index 3eb291624..9085fea5a 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -119,7 +119,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut }) cliproxyexecutor.MarkUpstreamAttempt(ctx) - httpResp, errDo := httpClient.Do(httpReq) + httpResp, errDo := helps.WithAntigravityHTTPClientTrace(httpClient, auth, "count_tokens").Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) return cliproxyexecutor.Response{}, errDo diff --git a/internal/runtime/executor/helps/antigravity_connection_trace.go b/internal/runtime/executor/helps/antigravity_connection_trace.go new file mode 100644 index 000000000..67144875c --- /dev/null +++ b/internal/runtime/executor/helps/antigravity_connection_trace.go @@ -0,0 +1,77 @@ +package helps + +import ( + "net/http" + "net/http/httptrace" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +// WithAntigravityHTTPClientTrace traces each actual request, including redirects. +// The client copy keeps the original pooled transport and client settings. +func WithAntigravityHTTPClientTrace(client *http.Client, auth *cliproxyauth.Auth, operation string) *http.Client { + if client == nil { + return nil + } + clone := *client + base := client.Transport + if base == nil { + base = http.DefaultTransport + } + authIndex := "" + if auth != nil { + authIndex = auth.Index + } + clone.Transport = &antigravityConnectionTraceTransport{ + base: base, auth: cliproxyauth.Auth{Index: authIndex}, operation: operation, + } + return &clone +} + +type antigravityConnectionTraceTransport struct { + base http.RoundTripper + auth cliproxyauth.Auth + operation string +} + +func (t *antigravityConnectionTraceTransport) RoundTrip(req *http.Request) (*http.Response, error) { + return t.base.RoundTrip(WithAntigravityConnectionTrace(req, &t.auth, t.operation)) +} + +// CloseIdleConnections preserves the underlying transport cleanup behavior. +func (t *antigravityConnectionTraceTransport) CloseIdleConnections() { + if closer, ok := t.base.(interface{ CloseIdleConnections() }); ok { + closer.CloseIdleConnections() + } +} + +// WithAntigravityConnectionTrace reports upstream connection reuse only in debug mode. +// It preserves existing trace hooks and never logs URLs, headers, or credential material. +func WithAntigravityConnectionTrace(req *http.Request, auth *cliproxyauth.Auth, operation string) *http.Request { + if req == nil || req.URL == nil || !log.IsLevelEnabled(log.DebugLevel) { + return req + } + authIndex := "" + if auth != nil { + authIndex = auth.Index + } + requestID := logging.GetRequestID(req.Context()) + host := req.URL.Hostname() + trace := &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + log.WithFields(log.Fields{ + "provider": "antigravity", + "operation": operation, + "upstream_host": host, + "auth_index": authIndex, + "request_id": requestID, + "reused": info.Reused, + "was_idle": info.WasIdle, + "idle_time": info.IdleTime.String(), + }).Debug("antigravity upstream connection acquired") + }, + } + return req.WithContext(httptrace.WithClientTrace(req.Context(), trace)) +} diff --git a/internal/runtime/executor/helps/antigravity_connection_trace_redirect_test.go b/internal/runtime/executor/helps/antigravity_connection_trace_redirect_test.go new file mode 100644 index 000000000..1bf796c66 --- /dev/null +++ b/internal/runtime/executor/helps/antigravity_connection_trace_redirect_test.go @@ -0,0 +1,123 @@ +package helps + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/http/httptrace" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +func TestAntigravityConnectionTraceRedirect(t *testing.T) { + logger := log.StandardLogger() + oldLevel := logger.GetLevel() + oldHooks := logger.ReplaceHooks(make(log.LevelHooks)) + oldOutput := logger.Out + logger.SetLevel(log.DebugLevel) + logger.SetOutput(io.Discard) + hook := logtest.NewLocal(logger) + t.Cleanup(func() { + logger.SetLevel(oldLevel) + logger.ReplaceHooks(oldHooks) + logger.SetOutput(oldOutput) + }) + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + })) + defer target.Close() + targetURL := strings.Replace(target.URL, "127.0.0.1", "localhost", 1) + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, targetURL+"/secret-path?token=secret-query", http.StatusFound) + })) + defer source.Close() + for _, operation := range []string{"generate_stream", "project_discovery"} { + t.Run(operation, func(t *testing.T) { + hook.Reset() + transport := &http.Transport{} + client := &http.Client{Transport: transport} + defer transport.CloseIdleConnections() + redirectCalls := 0 + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + redirectCalls++ + if req.URL.Hostname() != "localhost" { + t.Fatalf("unexpected redirect target: %s", req.URL.Hostname()) + } + return nil + } + tracedClient := WithAntigravityHTTPClientTrace(client, &cliproxyauth.Auth{Index: "redirect-index"}, operation) + existingCalls := 0 + ctx := logging.WithRequestID(context.Background(), "redirect-request-id") + ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{GotConn: func(httptrace.GotConnInfo) { existingCalls++ }}) + for attempt := 0; attempt < 2; attempt++ { + req, errReq := http.NewRequestWithContext(ctx, http.MethodGet, source.URL, nil) + if errReq != nil { + t.Fatal(errReq) + } + resp, errDo := tracedClient.Do(req) + if errDo != nil { + t.Fatal(errDo) + } + _, errRead := io.Copy(io.Discard, resp.Body) + errClose := resp.Body.Close() + if errRead != nil { + t.Fatal(errRead) + } + if errClose != nil { + t.Fatal(errClose) + } + if resp.Request.URL.Hostname() != "localhost" { + t.Fatal("redirect did not reach target") + } + } + entries := hook.AllEntries() + if len(entries) != 4 || existingCalls != 4 || redirectCalls != 2 { + t.Fatalf("entries=%d existing hooks=%d redirects=%d", len(entries), existingCalls, redirectCalls) + } + for i, entry := range entries { + host := "127.0.0.1" + if i%2 == 1 { + host = "localhost" + } + fields := entry.Data + if fields["upstream_host"] != host || fields["request_id"] != "redirect-request-id" || fields["auth_index"] != "redirect-index" || fields["operation"] != operation || fields["reused"] != (i >= 2) { + t.Fatalf("connection %d has incorrect fields: %v", i, fields) + } + formatted, errFormat := (&logging.LogFormatter{}).Format(entry) + if errFormat != nil { + t.Fatal(errFormat) + } + if strings.Contains(string(formatted), "secret") { + t.Fatalf("secret leaked: %s", formatted) + } + } + // The wrapper must not attach its trace to the original request context. + if trace := httptrace.ContextClientTrace(ctx); trace == nil { + t.Fatal("original trace lost") + } + }) + } +} + +func TestAntigravityTraceTransportCloseIdleConnections(t *testing.T) { + base := &antigravityTraceCleanupTransport{} + client := WithAntigravityHTTPClientTrace(&http.Client{Transport: base}, nil, "http_request") + client.CloseIdleConnections() + if !base.closed { + t.Fatal("idle connection cleanup was not forwarded") + } +} + +type antigravityTraceCleanupTransport struct{ closed bool } + +func (t *antigravityTraceCleanupTransport) RoundTrip(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody}, nil +} + +func (t *antigravityTraceCleanupTransport) CloseIdleConnections() { t.closed = true } diff --git a/internal/runtime/executor/helps/antigravity_connection_trace_test.go b/internal/runtime/executor/helps/antigravity_connection_trace_test.go new file mode 100644 index 000000000..58b8fdb27 --- /dev/null +++ b/internal/runtime/executor/helps/antigravity_connection_trace_test.go @@ -0,0 +1,173 @@ +package helps + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/http/httptrace" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +func TestAntigravityConnectionTrace(t *testing.T) { + logger := log.StandardLogger() + oldLevel := logger.GetLevel() + oldHooks := logger.ReplaceHooks(make(log.LevelHooks)) + oldOutput := logger.Out + logger.SetOutput(io.Discard) + hook := logtest.NewLocal(logger) + t.Cleanup(func() { + logger.SetLevel(oldLevel) + logger.ReplaceHooks(oldHooks) + logger.SetOutput(oldOutput) + }) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + })) + defer server.Close() + transport := &http.Transport{} + defer transport.CloseIdleConnections() + client := &http.Client{Transport: transport} + auth := &cliproxyauth.Auth{Index: "safe-index", ID: "secret-identity", Metadata: map[string]any{"access_token": "secret-token"}} + + logger.SetLevel(log.InfoLevel) + req, errReq := http.NewRequest(http.MethodGet, server.URL, nil) + if errReq != nil { + t.Fatal(errReq) + } + if got := WithAntigravityConnectionTrace(req, auth, "generate"); got != req { + t.Fatal("info mode must leave request unchanged") + } + if WithAntigravityConnectionTrace(nil, nil, "generate") != nil { + t.Fatal("nil request must remain nil") + } + + logger.SetLevel(log.DebugLevel) + invalidReq := &http.Request{Header: make(http.Header)} + if WithAntigravityConnectionTrace(invalidReq, auth, "generate") != invalidReq { + t.Fatal("nil URL request must remain unchanged") + } + if _, errDo := client.Do(WithAntigravityConnectionTrace(invalidReq, auth, "generate")); errDo == nil { + t.Fatal("nil URL must return an HTTP client error") + } + existingCalls := 0 + ctx := logging.WithRequestID(context.Background(), "test-request-id") + ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{GotConn: func(httptrace.GotConnInfo) { existingCalls++ }}) + for i := 0; i < 2; i++ { + req, errReq := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/secret-path?token=secret-query", nil) + if errReq != nil { + t.Fatal(errReq) + } + req.Header.Set("Authorization", "Bearer secret-header") + resp, errDo := client.Do(WithAntigravityConnectionTrace(req, auth, "generate")) + if errDo != nil { + t.Fatal(errDo) + } + _, errRead := io.Copy(io.Discard, resp.Body) + errClose := resp.Body.Close() + if errRead != nil { + t.Fatal(errRead) + } + if errClose != nil { + t.Fatal(errClose) + } + } + if existingCalls != 2 { + t.Fatalf("existing trace calls = %d, want 2", existingCalls) + } + entries := hook.AllEntries() + if len(entries) != 2 { + t.Fatalf("log entries = %d, want 2", len(entries)) + } + for i, entry := range entries { + if entry.Level != log.DebugLevel || entry.Message != "antigravity upstream connection acquired" { + t.Fatalf("unexpected entry: %+v", entry) + } + if entry.Data["reused"] != (i == 1) || entry.Data["was_idle"] != (i == 1) { + t.Fatalf("connection %d fields: %v", i, entry.Data) + } + if entry.Data["auth_index"] != "safe-index" || entry.Data["request_id"] != "test-request-id" || entry.Data["upstream_host"] != "127.0.0.1" { + t.Fatalf("unexpected fields: %v", entry.Data) + } + if _, ok := entry.Data["idle_time"].(string); !ok { + t.Fatal("missing idle_time") + } + formattedBytes, errString := (&logging.LogFormatter{}).Format(entry) + formatted := string(formattedBytes) + if errString != nil { + t.Fatal(errString) + } + for _, field := range []string{`auth_index="safe-index"`, `upstream_host="127.0.0.1"`, `operation="generate"`, "reused=", "was_idle=", "idle_time=", "test-request-id"[len("test-request-id")-8:]} { + if !strings.Contains(formatted, field) { + t.Fatalf("formatted log missing %q: %s", field, formatted) + } + } + if strings.Contains(formatted, "secret") { + t.Fatalf("secret leaked in log: %s", formatted) + } + } + + // Credential-discovery helpers create their own requests using the supplied client. + hook.Reset() + tracedClient := WithAntigravityHTTPClientTrace(client, auth, "project_discovery") + if client.Transport != transport || tracedClient == client { + t.Fatal("tracing must not modify the original client or pool") + } + for _, path := range []string{"/v1internal:loadCodeAssist", "/v1internal:onboardUser"} { + discoveryReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, server.URL+path, nil) + if errReq != nil { + t.Fatal(errReq) + } + resp, errDo := tracedClient.Do(discoveryReq) + if errDo != nil { + t.Fatal(errDo) + } + _, errRead := io.Copy(io.Discard, resp.Body) + errClose := resp.Body.Close() + if errRead != nil { + t.Fatal(errRead) + } + if errClose != nil { + t.Fatal(errClose) + } + } + entries = hook.AllEntries() + if len(entries) != 2 || existingCalls != 4 { + t.Fatalf("discovery entries=%d trace calls=%d", len(entries), existingCalls) + } + for _, entry := range entries { + if entry.Data["operation"] != "project_discovery" || entry.Data["request_id"] != "test-request-id" || entry.Data["reused"] != true { + t.Fatalf("discovery trace fields: %v", entry.Data) + } + } + if WithAntigravityHTTPClientTrace(nil, nil, "project_discovery") != nil { + t.Fatal("nil client must remain nil") + } + defaultClient := &http.Client{} + tracedDefault := WithAntigravityHTTPClientTrace(defaultClient, nil, "project_discovery") + if tracedDefault.Transport.(*antigravityConnectionTraceTransport).base != http.DefaultTransport || defaultClient.Transport != nil { + t.Fatal("default transport must be preserved") + } + hook.Reset() + req = WithAntigravityConnectionTrace(req, nil, "credits_query") + httptrace.ContextClientTrace(req.Context()).GotConn(httptrace.GotConnInfo{}) + if hook.LastEntry().Data["request_id"] != "" || hook.LastEntry().Data["operation"] != "credits_query" { + t.Fatal("background query must retain its operation without inventing a request ID") + } + if hook.LastEntry().Data["auth_index"] != "" { + t.Fatal("nil auth should have empty index") + } + logger.SetLevel(log.InfoLevel) + hook.Reset() + httptrace.ContextClientTrace(req.Context()).GotConn(httptrace.GotConnInfo{}) + if len(hook.AllEntries()) != 0 { + t.Fatal("trace should respect runtime log-level changes") + } +}