chore: improvements

Signed-off-by: Danny Kopping <danny@coder.com>

Author: dannykopping

example/main.go

@@ -5,7 +5,6 @@ package main
 import (
 	"context"
 	"database/sql"
-	"log"
 	"net/http"
 	"os"
 	"regexp"
@@ -27,15 +26,21 @@ func main() {
 	// Initialize SQLite database with WAL mode for better concurrency.
 	db, err := sql.Open("sqlite", "aibridge.db?_journal_mode=WAL&_busy_timeout=5000")
 	if err != nil {
-		log.Fatalf("open database: %v", err)
+		logger.Fatal(ctx, "open database", slog.Error(err))
 	}
 	defer db.Close()
 	db.SetMaxOpenConns(1) // SQLite only supports one writer at a time.
 
 	if err := initSchema(db); err != nil {
-		log.Fatalf("init schema: %v", err)
+		logger.Fatal(ctx, "init schema", slog.Error(err))
 	}
 
+	recorder, err := NewSQLiteRecorder(db, logger)
+	if err != nil {
+		logger.Fatal(ctx, "create recorder", slog.Error(err))
+	}
+	defer recorder.Close()
+
 	// Configure providers.
 	providers := []aibridge.Provider{
 		aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
@@ -63,25 +68,25 @@ func main() {
 		nil,                                  // denylist
 	)
 	if err != nil {
-		log.Fatalf("create deepwiki mcp proxy: %v", err)
+		logger.Fatal(ctx, "create deepwiki mcp proxy", slog.Error(err))
 	}
 
 	mcpProxy = mcp.NewServerProxyManager(map[string]mcp.ServerProxier{"deepwiki": deepwikiProxy})
 	if err := mcpProxy.Init(ctx); err != nil {
-		log.Printf("mcp init warning: %v", err)
+		logger.Warn(ctx, "mcp init warning", slog.Error(err))
 	}
 
 	// Create the bridge with SQLite recorder.
 	bridge, err := aibridge.NewRequestBridge(
 		ctx,
 		providers,
-		&SQLiteRecorder{db: db, logger: logger},
+		recorder,
 		mcpProxy,
 		metrics,
 		logger,
 	)
 	if err != nil {
-		log.Fatalf("create bridge: %v", err)
+		logger.Fatal(ctx, "create bridge", slog.Error(err))
 	}
 	defer bridge.Shutdown(ctx)
 
@@ -90,9 +95,9 @@ func main() {
 	mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
 	mux.Handle("/", actorMiddleware(bridge))
 
-	log.Println("listening on :8080")
+	logger.Info(ctx, "listening on :8080")
 	if err := http.ListenAndServe(":8080", mux); err != nil {
-		log.Fatal(err)
+		logger.Fatal(ctx, "http server error", slog.Error(err))
 	}
 }
 

example/recorder.go

@@ -14,13 +14,67 @@ import (
 type SQLiteRecorder struct {
 	db     *sql.DB
 	logger slog.Logger
+
+	stmtInsertInterception    *sql.Stmt
+	stmtUpdateInterception    *sql.Stmt
+	stmtInsertTokenUsage      *sql.Stmt
+	stmtInsertPromptUsage     *sql.Stmt
+	stmtInsertToolUsage       *sql.Stmt
+}
+
+func NewSQLiteRecorder(db *sql.DB, logger slog.Logger) (*SQLiteRecorder, error) {
+	r := &SQLiteRecorder{db: db, logger: logger}
+
+	var err error
+	r.stmtInsertInterception, err = db.Prepare(`
+		INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, started_at, metadata)
+		VALUES (?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return nil, err
+	}
+
+	r.stmtUpdateInterception, err = db.Prepare(`
+		UPDATE aibridge_interceptions SET ended_at = ? WHERE id = ?`)
+	if err != nil {
+		return nil, err
+	}
+
+	r.stmtInsertTokenUsage, err = db.Prepare(`
+		INSERT INTO aibridge_token_usages (id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at)
+		VALUES (?, ?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return nil, err
+	}
+
+	r.stmtInsertPromptUsage, err = db.Prepare(`
+		INSERT INTO aibridge_user_prompts (id, interception_id, provider_response_id, prompt, metadata, created_at)
+		VALUES (?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return nil, err
+	}
+
+	r.stmtInsertToolUsage, err = db.Prepare(`
+		INSERT INTO aibridge_tool_usages (id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at)
+		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return nil, err
+	}
+
+	return r, nil
+}
+
+func (r *SQLiteRecorder) Close() error {
+	r.stmtInsertInterception.Close()
+	r.stmtUpdateInterception.Close()
+	r.stmtInsertTokenUsage.Close()
+	r.stmtInsertPromptUsage.Close()
+	r.stmtInsertToolUsage.Close()
+	return nil
 }
 
 func (r *SQLiteRecorder) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error {
 	metadata, _ := json.Marshal(req.Metadata)
-	_, err := r.db.ExecContext(ctx, `
-		INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, started_at, metadata)
-		VALUES (?, ?, ?, ?, ?, ?)`,
+	_, err := r.stmtInsertInterception.ExecContext(ctx,
 		req.ID, req.InitiatorID, req.Provider, req.Model, req.StartedAt, string(metadata),
 	)
 	if err != nil {
@@ -30,10 +84,7 @@ func (r *SQLiteRecorder) RecordInterception(ctx context.Context, req *aibridge.I
 }
 
 func (r *SQLiteRecorder) RecordInterceptionEnded(ctx context.Context, req *aibridge.InterceptionRecordEnded) error {
-	_, err := r.db.ExecContext(ctx, `
-		UPDATE aibridge_interceptions SET ended_at = ? WHERE id = ?`,
-		req.EndedAt, req.ID,
-	)
+	_, err := r.stmtUpdateInterception.ExecContext(ctx, req.EndedAt, req.ID)
 	if err != nil {
 		r.logger.Warn(ctx, "failed to record interception end", slog.Error(err))
 	}
@@ -51,9 +102,7 @@ func (r *SQLiteRecorder) RecordTokenUsage(ctx context.Context, req *aibridge.Tok
 	}
 	metadata, _ := json.Marshal(merged)
 
-	_, err := r.db.ExecContext(ctx, `
-		INSERT INTO aibridge_token_usages (id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at)
-		VALUES (?, ?, ?, ?, ?, ?, ?)`,
+	_, err := r.stmtInsertTokenUsage.ExecContext(ctx,
 		uuid.NewString(), req.InterceptionID, req.MsgID, req.Input, req.Output, string(metadata), req.CreatedAt,
 	)
 	if err != nil {
@@ -64,9 +113,7 @@ func (r *SQLiteRecorder) RecordTokenUsage(ctx context.Context, req *aibridge.Tok
 
 func (r *SQLiteRecorder) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error {
 	metadata, _ := json.Marshal(req.Metadata)
-	_, err := r.db.ExecContext(ctx, `
-		INSERT INTO aibridge_user_prompts (id, interception_id, provider_response_id, prompt, metadata, created_at)
-		VALUES (?, ?, ?, ?, ?, ?)`,
+	_, err := r.stmtInsertPromptUsage.ExecContext(ctx,
 		uuid.NewString(), req.InterceptionID, req.MsgID, req.Prompt, string(metadata), req.CreatedAt,
 	)
 	if err != nil {
@@ -90,9 +137,7 @@ func (r *SQLiteRecorder) RecordToolUsage(ctx context.Context, req *aibridge.Tool
 		invocationError = &errStr
 	}
 
-	_, err := r.db.ExecContext(ctx, `
-		INSERT INTO aibridge_tool_usages (id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at)
-		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+	_, err := r.stmtInsertToolUsage.ExecContext(ctx,
 		uuid.NewString(), req.InterceptionID, req.MsgID, serverURL, req.Tool, string(input), req.Injected, invocationError, string(metadata), req.CreatedAt,
 	)
 	if err != nil {