-
Notifications
You must be signed in to change notification settings - Fork 5
feat: add upstream request/response logging for passthrough routes #186
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
32e23b1
991d4a9
1a03d34
7b47019
153f439
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -32,26 +32,21 @@ type MiddlewareNext = func(*http.Request) (*http.Response, error) | |
| // Middleware is an HTTP middleware function compatible with SDK WithMiddleware options. | ||
| type Middleware = func(*http.Request, MiddlewareNext) (*http.Response, error) | ||
|
|
||
| // NewMiddleware returns a middleware function that dumps requests and responses to files. | ||
| // Files are written to the path returned by DumpPath. | ||
| // NewBridgeMiddleware returns a middleware function that dumps requests and responses to files. | ||
| // If baseDir is empty, returns nil (no middleware). | ||
| func NewMiddleware(baseDir, provider, model string, interceptionID uuid.UUID, logger slog.Logger, clk quartz.Clock) Middleware { | ||
| func NewBridgeMiddleware(baseDir string, provider string, model string, interceptionID uuid.UUID, logger slog.Logger, clk quartz.Clock) Middleware { | ||
| if baseDir == "" { | ||
| return nil | ||
| } | ||
|
|
||
| d := &dumper{ | ||
| baseDir: baseDir, | ||
| provider: provider, | ||
| model: model, | ||
| interceptionID: interceptionID, | ||
| clk: clk, | ||
| logger: logger, | ||
| dumpPath: interceptDumpPath(baseDir, provider, model, interceptionID, clk), | ||
| logger: logger, | ||
| } | ||
|
|
||
| return func(req *http.Request, next MiddlewareNext) (*http.Response, error) { | ||
| if err := d.dumpRequest(req); err != nil { | ||
| logger.Named("apidump").Warn(context.Background(), "failed to dump request", slog.Error(err)) | ||
| logger.Named("apidump").Warn(req.Context(), "failed to dump request", slog.Error(err)) | ||
| } | ||
|
|
||
| // TODO: https://github.com/coder/aibridge/issues/129 | ||
|
|
@@ -61,24 +56,20 @@ func NewMiddleware(baseDir, provider, model string, interceptionID uuid.UUID, lo | |
| } | ||
|
|
||
| if err := d.dumpResponse(resp); err != nil { | ||
| logger.Named("apidump").Warn(context.Background(), "failed to dump response", slog.Error(err)) | ||
| logger.Named("apidump").Warn(req.Context(), "failed to dump response", slog.Error(err)) | ||
| } | ||
|
|
||
| return resp, nil | ||
| } | ||
| } | ||
|
|
||
| type dumper struct { | ||
| baseDir string | ||
| provider string | ||
| model string | ||
| interceptionID uuid.UUID | ||
| clk quartz.Clock | ||
| logger slog.Logger | ||
| dumpPath string | ||
| logger slog.Logger | ||
| } | ||
|
|
||
| func (d *dumper) dumpRequest(req *http.Request) error { | ||
| dumpPath := d.path(SuffixRequest) | ||
| dumpPath := d.dumpPath + SuffixRequest | ||
| if err := os.MkdirAll(filepath.Dir(dumpPath), 0o755); err != nil { | ||
| return fmt.Errorf("create dump dir: %w", err) | ||
| } | ||
|
|
@@ -98,25 +89,44 @@ func (d *dumper) dumpRequest(req *http.Request) error { | |
|
|
||
| // Build raw HTTP request format | ||
| var buf bytes.Buffer | ||
| fmt.Fprintf(&buf, "%s %s %s\r\n", req.Method, req.URL.RequestURI(), req.Proto) | ||
| d.writeRedactedHeaders(&buf, req.Header, sensitiveRequestHeaders, map[string]string{ | ||
| _, err := fmt.Fprintf(&buf, "%s %s %s\r\n", req.Method, req.URL.RequestURI(), req.Proto) | ||
| if err != nil { | ||
| return fmt.Errorf("write request uri: %w", err) | ||
| } | ||
| err = d.writeRedactedHeaders(&buf, req.Header, sensitiveRequestHeaders, map[string]string{ | ||
| "Content-Length": fmt.Sprintf("%d", len(prettyBody)), | ||
| }) | ||
| if err != nil { | ||
| return fmt.Errorf("write request headers: %w", err) | ||
| } | ||
|
|
||
| fmt.Fprintf(&buf, "\r\n") | ||
| _, err = fmt.Fprintf(&buf, "\r\n") | ||
| if err != nil { | ||
| return fmt.Errorf("write request header terminator: %w", err) | ||
| } | ||
| buf.Write(prettyBody) | ||
| buf.WriteByte('\n') | ||
|
|
||
| return os.WriteFile(dumpPath, buf.Bytes(), 0o644) | ||
| } | ||
|
|
||
| func (d *dumper) dumpResponse(resp *http.Response) error { | ||
| dumpPath := d.path(SuffixResponse) | ||
| dumpPath := d.dumpPath + SuffixResponse | ||
dannykopping marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| // Build raw HTTP response headers | ||
| var headerBuf bytes.Buffer | ||
| fmt.Fprintf(&headerBuf, "%s %s\r\n", resp.Proto, resp.Status) | ||
| d.writeRedactedHeaders(&headerBuf, resp.Header, sensitiveResponseHeaders, nil) | ||
| fmt.Fprintf(&headerBuf, "\r\n") | ||
| _, err := fmt.Fprintf(&headerBuf, "%s %s\r\n", resp.Proto, resp.Status) | ||
| if err != nil { | ||
| return fmt.Errorf("write response status: %w", err) | ||
| } | ||
| err = d.writeRedactedHeaders(&headerBuf, resp.Header, sensitiveResponseHeaders, nil) | ||
| if err != nil { | ||
| return fmt.Errorf("write response headers: %w", err) | ||
| } | ||
| _, err = fmt.Fprintf(&headerBuf, "\r\n") | ||
| if err != nil { | ||
| return fmt.Errorf("write response header terminator: %w", err) | ||
| } | ||
|
|
||
| // Wrap the response body to capture it as it streams | ||
| if resp.Body != nil { | ||
|
|
@@ -141,7 +151,7 @@ func (d *dumper) dumpResponse(resp *http.Response) error { | |
| // for deterministic output. | ||
| // `sensitive` and `overrides` must both supply keys in canoncialized form. | ||
| // See [textproto.MIMEHeader]. | ||
| func (d *dumper) writeRedactedHeaders(w io.Writer, headers http.Header, sensitive map[string]struct{}, overrides map[string]string) { | ||
| func (d *dumper) writeRedactedHeaders(w io.Writer, headers http.Header, sensitive map[string]struct{}, overrides map[string]string) error { | ||
| // Collect all header keys including overrides. | ||
| headerKeys := make([]string, 0, len(headers)+len(overrides)) | ||
| seen := make(map[string]struct{}, len(headers)+len(overrides)) | ||
|
|
@@ -163,7 +173,10 @@ func (d *dumper) writeRedactedHeaders(w io.Writer, headers http.Header, sensitiv | |
| // If no values exist but we have an override, use that. | ||
| if len(values) == 0 { | ||
| if override, ok := overrides[key]; ok { | ||
| fmt.Fprintf(w, "%s: %s\r\n", key, override) | ||
| _, err := fmt.Fprintf(w, "%s: %s\r\n", key, override) | ||
| if err != nil { | ||
| return fmt.Errorf("write response header override: %w", err) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This function is used for both request/responses, right? Consider removing |
||
| } | ||
| } | ||
| continue | ||
| } | ||
|
|
@@ -175,16 +188,71 @@ func (d *dumper) writeRedactedHeaders(w io.Writer, headers http.Header, sensitiv | |
| if isSensitive { | ||
| value = redactHeaderValue(value) | ||
| } | ||
| fmt.Fprintf(w, "%s: %s\r\n", key, value) | ||
| _, err := fmt.Fprintf(w, "%s: %s\r\n", key, value) | ||
| if err != nil { | ||
| return fmt.Errorf("write response headers: %w", err) | ||
| } | ||
| } | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // interceptDumpPath returns the base file path (without req/resp suffix) for an interception dump. | ||
| func interceptDumpPath(baseDir string, provider string, model string, interceptionID uuid.UUID, clk quartz.Clock) string { | ||
| safeModel := strings.ReplaceAll(model, "/", "-") | ||
| return filepath.Join(baseDir, provider, safeModel, fmt.Sprintf("%d-%s", clk.Now().UTC().UnixMilli(), interceptionID)) | ||
| } | ||
|
|
||
| // passthroughDumpPath returns the base file path (without req/resp suffix) for a passthrough dump. | ||
| func passthroughDumpPath(baseDir string, provider string, urlPath string, clk quartz.Clock) string { | ||
| safeURLPath := strings.ReplaceAll(strings.TrimPrefix(urlPath, "/"), "/", "-") | ||
| return filepath.Join(baseDir, provider, "passthrough", fmt.Sprintf("%d-%s-%s", clk.Now().UTC().UnixMilli(), safeURLPath, uuid.NewString()[:4])) | ||
| } | ||
|
|
||
| // NewPassthroughMiddleware returns http.RoundTripper that dumps requests and responses to files. | ||
| // If baseDir is empty, returns the original transport unchanged. | ||
| // Used for logging in pass through routes. | ||
| func NewPassthroughMiddleware(transport http.RoundTripper, baseDir string, provider string, logger slog.Logger, clk quartz.Clock) http.RoundTripper { | ||
| if baseDir == "" { | ||
| return transport | ||
| } | ||
| return &dumpRoundTripper{ | ||
| inner: transport, | ||
| baseDir: baseDir, | ||
| provider: provider, | ||
| clk: clk, | ||
| logger: logger, | ||
| } | ||
| } | ||
|
|
||
| // path returns the path to a request/response dump file for a given interception. | ||
| // suffix should be SuffixRequest or SuffixResponse. | ||
| func (d *dumper) path(suffix string) string { | ||
| safeModel := strings.ReplaceAll(d.model, "/", "-") | ||
| return filepath.Join(d.baseDir, d.provider, safeModel, fmt.Sprintf("%d-%s%s", d.clk.Now().UTC().UnixMilli(), d.interceptionID, suffix)) | ||
| type dumpRoundTripper struct { | ||
| inner http.RoundTripper | ||
| baseDir string | ||
| provider string | ||
| clk quartz.Clock | ||
| logger slog.Logger | ||
| } | ||
|
|
||
| func (rt *dumpRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { | ||
| dumper := dumper{ | ||
| dumpPath: passthroughDumpPath(rt.baseDir, rt.provider, req.URL.Path, rt.clk), | ||
| logger: rt.logger, | ||
| } | ||
|
|
||
| if err := dumper.dumpRequest(req); err != nil { | ||
| dumper.logger.Named("apidump").Warn(req.Context(), "failed to dump passthrough request", slog.Error(err)) | ||
| } | ||
|
|
||
| resp, err := rt.inner.RoundTrip(req) | ||
| if err != nil { | ||
| return resp, err | ||
| } | ||
|
|
||
| if err := dumper.dumpResponse(resp); err != nil { | ||
| dumper.logger.Named("apidump").Warn(req.Context(), "failed to dump passthrough response", slog.Error(err)) | ||
| } | ||
|
|
||
| return resp, nil | ||
| } | ||
|
|
||
| // prettyPrintJSON returns indented JSON if body is valid JSON, otherwise returns body as-is. | ||
|
|
@@ -194,12 +262,11 @@ func prettyPrintJSON(body []byte) []byte { | |
| if len(body) == 0 { | ||
| return body | ||
| } | ||
| result := pretty.Pretty(body) | ||
| // pretty.Pretty returns a truncated/modified result for invalid JSON, | ||
| // so check if the result is valid JSON; if not, return the original. | ||
| if !json.Valid(result) { | ||
| return body | ||
|
|
||
| result := body | ||
| if json.Valid(body) { | ||
| result = pretty.Pretty(body) | ||
| } | ||
|
Comment on lines
+266
to
+269
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why did this change? Since |
||
| // Trim trailing newline added by pretty.Pretty. | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Was there a reson for removing trailing new line? |
||
| return bytes.TrimSuffix(result, []byte("\n")) | ||
|
|
||
| return result | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.