chore: refactoring, tests

Signed-off-by: Danny Kopping <[email protected]>

Author: dannykopping

bridge_integration_test.go

@@ -1152,8 +1152,8 @@ func createMockMCPSrv(t *testing.T) http.Handler {
 	return server.NewStreamableHTTPServer(s)
 }
 
-func cfg(url, key string) aibridge.ProviderConfig {
-	return aibridge.ProviderConfig{
+func cfg(url, key string) *aibridge.ProviderConfig {
+	return &aibridge.ProviderConfig{
 		BaseURL: url,
 		Key:     key,
 	}

config.go

@@ -4,7 +4,11 @@ import "sync/atomic"
 
 type ProviderConfig struct {
 	BaseURL, Key string
-	// EnableUpstreamLogging enables logging of upstream API requests and responses to /tmp/$provider.log
+	// UpstreamLoggingDir specifies the base directory for upstream logging.
+	// If empty, os.TempDir() will be used.
+	// Logs are written to $UpstreamLoggingDir/$provider/$model/$id.{req,res}.log
+	UpstreamLoggingDir string
+	// enableUpstreamLogging enables logging of upstream API requests and responses.
 	enableUpstreamLogging atomic.Bool
 }
 

provider_anthropic.go

@@ -104,7 +104,7 @@ func newAnthropicClient(logger slog.Logger, cfg *ProviderConfig, id, model strin
 	opts = append(opts, option.WithBaseURL(cfg.BaseURL))
 
 	if cfg.EnableUpstreamLogging() {
-		if middleware := createLoggingMiddleware(logger, "anthropic", id, model); middleware != nil {
+		if middleware := createLoggingMiddleware(logger, cfg, ProviderAnthropic, id, model); middleware != nil {
 			opts = append(opts, option.WithMiddleware(middleware))
 		}
 	}

provider_openai.go

@@ -107,7 +107,7 @@ func newOpenAIClient(logger slog.Logger, cfg *ProviderConfig, id, model string)
 	opts = append(opts, option.WithBaseURL(cfg.BaseURL))
 
 	if cfg.EnableUpstreamLogging() {
-		if middleware := createLoggingMiddleware(logger, "openai", id, model); middleware != nil {
+		if middleware := createLoggingMiddleware(logger, cfg, ProviderOpenAI, id, model); middleware != nil {
 			opts = append(opts, option.WithMiddleware(middleware))
 		}
 	}

request_logger.go

@@ -13,19 +13,20 @@ import (
 	"cdr.dev/slog"
 )
 
-// sanitizeModelName makes a model name safe for use as a directory name.
+// SanitizeModelName makes a model name safe for use as a directory name.
 // Replaces filesystem-unsafe characters with underscores.
-func sanitizeModelName(model string) string {
+func SanitizeModelName(model string) string {
+	repl := "_"
 	replacer := strings.NewReplacer(
-		"/", "_",
-		"\\", "_",
-		":", "_",
-		"*", "_",
-		"?", "_",
-		"\"", "_",
-		"<", "_",
-		">", "_",
-		"|", "_",
+		"/", repl,
+		"\\", repl,
+		":", repl,
+		"*", repl,
+		"?", repl,
+		"\"", repl,
+		"<", repl,
+		">", repl,
+		"|", repl,
 	)
 	return replacer.Replace(model)
 }
@@ -65,12 +66,19 @@ func logUpstreamError(logger *log.Logger, id, model string, err error) {
 }
 
 // createLoggingMiddleware creates a middleware function that logs requests and responses.
-// Logs are written to $TMPDIR/$provider/$model/$id.req.log and $TMPDIR/$provider/$model/$id.res.log
+// Logs are written to $baseDir/$provider/$model/$id.req.log and $baseDir/$provider/$model/$id.res.log
+// where baseDir is from cfg.UpstreamLoggingDir or os.TempDir() if not specified.
 // Returns nil if logging setup fails, logging errors via the provided logger.
-func createLoggingMiddleware(logger slog.Logger, provider, id, model string) func(*http.Request, func(*http.Request) (*http.Response, error)) (*http.Response, error) {
+func createLoggingMiddleware(logger slog.Logger, cfg *ProviderConfig, provider, id, model string) func(*http.Request, func(*http.Request) (*http.Response, error)) (*http.Response, error) {
 	ctx := context.Background()
-	safeModel := sanitizeModelName(model)
-	logDir := filepath.Join(os.TempDir(), provider, safeModel)
+	safeModel := SanitizeModelName(model)
+
+	baseDir := cfg.UpstreamLoggingDir
+	if baseDir == "" {
+		baseDir = os.TempDir()
+	}
+
+	logDir := filepath.Join(baseDir, provider, safeModel)
 
 	// Create the directory structure if it doesn't exist
 	if err := os.MkdirAll(logDir, 0755); err != nil {

request_logger_test.go

@@ -0,0 +1,175 @@
+package aibridge_test
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/aibridge"
+	"github.com/coder/aibridge/mcp"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/tools/txtar"
+)
+
+func TestRequestLogging(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		provider       string
+		fixture        []byte
+		route          string
+		createProvider func(*aibridge.ProviderConfig) aibridge.Provider
+	}{
+		{
+			provider: aibridge.ProviderAnthropic,
+			fixture:  antSimple,
+			route:    "/anthropic/v1/messages",
+			createProvider: func(cfg *aibridge.ProviderConfig) aibridge.Provider {
+				return aibridge.NewAnthropicProvider(cfg)
+			},
+		},
+		{
+			provider: aibridge.ProviderOpenAI,
+			fixture:  oaiSimple,
+			route:    "/openai/v1/chat/completions",
+			createProvider: func(cfg *aibridge.ProviderConfig) aibridge.Provider {
+				return aibridge.NewOpenAIProvider(cfg)
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.provider, func(t *testing.T) {
+			t.Parallel()
+
+			logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+
+			// Use a temp dir for this test
+			tmpDir := t.TempDir()
+
+			// Parse fixture
+			arc := txtar.Parse(tc.fixture)
+			files := filesMap(arc)
+			require.Contains(t, files, fixtureRequest)
+			require.Contains(t, files, fixtureNonStreamingResponse)
+
+			// Create mock server
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write(files[fixtureNonStreamingResponse])
+			}))
+			t.Cleanup(srv.Close)
+
+			cfg := aibridge.ProviderConfig{
+				BaseURL:            srv.URL,
+				Key:                apiKey,
+				UpstreamLoggingDir: tmpDir,
+			}
+			cfg.SetEnableUpstreamLogging(true)
+
+			provider := tc.createProvider(&cfg)
+			client := &mockRecorderClient{}
+			mcpProxy := mcp.NewServerProxyManager(nil)
+
+			bridge, err := aibridge.NewRequestBridge(context.Background(), []aibridge.Provider{provider}, logger, client, mcpProxy)
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				_ = bridge.Shutdown(context.Background())
+			})
+
+			// Make a request
+			req, err := http.NewRequestWithContext(t.Context(), "POST", tc.route, strings.NewReader(string(files[fixtureRequest])))
+			require.NoError(t, err)
+			req.Header.Set("Content-Type", "application/json")
+			req = req.WithContext(aibridge.AsActor(req.Context(), userID, nil))
+			rec := httptest.NewRecorder()
+			bridge.ServeHTTP(rec, req)
+			require.Equal(t, 200, rec.Code)
+
+			// Check that log files were created
+			// Parse the request to get the model name
+			var reqData map[string]any
+			require.NoError(t, json.Unmarshal(files[fixtureRequest], &reqData))
+			model := reqData["model"].(string)
+
+			logDir := filepath.Join(tmpDir, tc.provider, model)
+			entries, err := os.ReadDir(logDir)
+			require.NoError(t, err, "log directory should exist")
+			require.NotEmpty(t, entries, "log directory should contain files")
+
+			// Should have at least one .req.log and one .res.log file
+			var hasReq, hasRes bool
+			for _, entry := range entries {
+				name := entry.Name()
+				if strings.HasSuffix(name, ".req.log") {
+					hasReq = true
+					// Verify the file has content
+					content, err := os.ReadFile(filepath.Join(logDir, name))
+					require.NoError(t, err)
+					require.NotEmpty(t, content, "request log should have content")
+					require.Contains(t, string(content), "POST")
+				} else if strings.HasSuffix(name, ".res.log") {
+					hasRes = true
+					// Verify the file has content
+					content, err := os.ReadFile(filepath.Join(logDir, name))
+					require.NoError(t, err)
+					require.NotEmpty(t, content, "response log should have content")
+					require.Contains(t, string(content), "200")
+				}
+			}
+			require.True(t, hasReq, "should have request log file")
+			require.True(t, hasRes, "should have response log file")
+		})
+	}
+}
+
+func TestSanitizeModelName(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "simple model",
+			input:    "gpt-4o",
+			expected: "gpt-4o",
+		},
+		{
+			name:     "model with slash",
+			input:    "gpt-4o/mini",
+			expected: "gpt-4o_mini",
+		},
+		{
+			name:     "model with colon",
+			input:    "o1:2024-12-17",
+			expected: "o1_2024-12-17",
+		},
+		{
+			name:     "model with backslash",
+			input:    "model\\name",
+			expected: "model_name",
+		},
+		{
+			name:     "model with multiple special chars",
+			input:    "model:name/version?",
+			expected: "model_name_version_",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := aibridge.SanitizeModelName(tt.input)
+			require.Equal(t, tt.expected, result)
+		})
+	}
+}