Logging improvements

.golangci.yml

@@ -23,6 +23,7 @@ linters:
   - containedctx
   - contextcheck
   - copyloopvar
+  - depguard
   - dogsled
   - errcheck
   - errchkjson
@@ -58,6 +59,7 @@ linters:
       - errcheck
       - forcetypeassert
       - gosec
+      - depguard
     - path: 'cmd/(.+)/main\.go'
       linters:
       - forbidigo
@@ -77,6 +79,12 @@ linters:
       allow-no-explanation: []
       require-explanation: true
       require-specific: true
+    depguard:
+      rules:
+        test-only:
+          deny:
+            - pkg: github.com/replicate/cog-runtime/internal/loggingtest
+              desc: "loggingtest is test-only"
     revive:
       confidence: 0.8
       severity: warning

cmd/cog/main.go

@@ -9,11 +9,9 @@ import (
 	"time"
 
 	"github.com/alecthomas/kong"
-	"github.com/replicate/go/logging"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 
 	"github.com/replicate/cog-runtime/internal/config"
+	"github.com/replicate/cog-runtime/internal/logging"
 	"github.com/replicate/cog-runtime/internal/runner"
 	"github.com/replicate/cog-runtime/internal/service"
 	"github.com/replicate/cog-runtime/internal/version"
@@ -42,30 +40,15 @@ type CLI struct {
 	Test   TestCmd   `cmd:"" help:"Run model tests to verify functionality"`
 }
 
-// createBaseLogger creates a base logger with configurable level
-func createBaseLogger(name string) *zap.Logger {
-	logLevel := os.Getenv("COG_LOG_LEVEL")
-	if logLevel == "" {
-		logLevel = "info"
-	}
-	_, err := zapcore.ParseLevel(logLevel)
-	if err != nil {
-		fmt.Printf("Failed to parse log level \"%s\": %s\n", logLevel, err) //nolint:forbidigo // logger setup error reporting
-	}
-
-	return logging.New(name).WithOptions(zap.IncreaseLevel(zapcore.DebugLevel))
-}
-
 // buildServiceConfig converts CLI ServerCmd to service configuration
 func buildServiceConfig(s *ServerCmd) (config.Config, error) {
-	log := createBaseLogger("cog-config").Sugar()
+	log := logging.New("cog-config").Sugar()
 
 	logLevel := log.Level()
-	log.Infow("log level", "level", logLevel)
-	log.Infow("env log level", "level", os.Getenv("COG_LOG_LEVEL"))
+	log.Debugw("log level", "level", logLevel)
 	// One-shot mode requires procedure mode
 	if s.OneShot && !s.UseProcedureMode {
-		log.Error("one-shot mode requires procedure mode")
+		log.Fatal("one-shot mode requires procedure mode")
 		return config.Config{}, fmt.Errorf("one-shot mode requires procedure mode, use --use-procedure-mode")
 	}
 
@@ -115,7 +98,7 @@ func buildServiceConfig(s *ServerCmd) (config.Config, error) {
 
 func (s *ServerCmd) Run() error {
 	// Create base logger
-	baseLogger := createBaseLogger("cog")
+	baseLogger := logging.New("cog")
 	log := baseLogger.Sugar()
 
 	// Build service configuration
@@ -143,7 +126,7 @@ func (s *ServerCmd) Run() error {
 }
 
 func (s *SchemaCmd) Run() error {
-	log := createBaseLogger("cog-schema").Sugar()
+	log := logging.New("cog-schema").Sugar()
 
 	wd, err := os.Getwd()
 	if err != nil {
@@ -169,7 +152,7 @@ func (s *SchemaCmd) Run() error {
 }
 
 func (t *TestCmd) Run() error {
-	log := createBaseLogger("cog-test").Sugar()
+	log := logging.New("cog-test").Sugar()
 
 	wd, err := os.Getwd()
 	if err != nil {

internal/logging/logger.go

@@ -0,0 +1,171 @@
+package logging
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+// Custom log levels with Trace below Debug
+const (
+	TraceLevel = zapcore.Level(-8) // Below Debug (-4)
+)
+
+// customLowercaseLevelEncoder handles our custom Trace level display (lowercase)
+func customLowercaseLevelEncoder(level zapcore.Level, enc zapcore.PrimitiveArrayEncoder) {
+	switch level {
+	case TraceLevel:
+		enc.AppendString("trace")
+	default:
+		zapcore.LowercaseLevelEncoder(level, enc)
+	}
+}
+
+// customColorLevelEncoder handles our custom Trace level display (with colors)
+func customColorLevelEncoder(level zapcore.Level, enc zapcore.PrimitiveArrayEncoder) {
+	switch level {
+	case TraceLevel:
+		enc.AppendString("\x1b[90mTRACE\x1b[0m") // Gray color for trace
+	default:
+		zapcore.CapitalColorLevelEncoder(level, enc)
+	}
+}
+
+// Logger embeds zap.Logger and adds Trace level support
+type Logger struct {
+	*zap.Logger
+}
+
+// SugaredLogger embeds zap.SugaredLogger and adds Trace level support
+type SugaredLogger struct {
+	*zap.SugaredLogger
+}
+
+// New creates a new logger with the given name
+func New(name string) *Logger {
+	// Check if we should use development config (console format)
+	logFormat := os.Getenv("LOG_FORMAT")
+	isDevelopment := logFormat == "development" || logFormat == "console"
+
+	var cfg zap.Config
+	if isDevelopment {
+		cfg = zap.NewDevelopmentConfig()
+		cfg.Level = zap.NewAtomicLevelAt(zapcore.DebugLevel)
+		cfg.EncoderConfig.EncodeLevel = customColorLevelEncoder
+	} else {
+		cfg = zap.NewProductionConfig()
+		cfg.Level = zap.NewAtomicLevelAt(zapcore.InfoLevel)
+		cfg.EncoderConfig.EncodeLevel = customLowercaseLevelEncoder
+	}
+
+	// Set log level from environment (COG_LOG_LEVEL takes precedence, fallback to LOG_LEVEL)
+	logLevel := os.Getenv("COG_LOG_LEVEL")
+	if logLevel == "" {
+		logLevel = os.Getenv("LOG_LEVEL")
+	}
+	if logLevel != "" {
+		level, err := parseLevel(logLevel)
+		if err != nil {
+			fmt.Printf("Failed to parse log level \"%s\": %s\n", logLevel, err) //nolint:forbidigo // logger setup error reporting
+		} else {
+			cfg.Level = zap.NewAtomicLevelAt(level)
+		}
+	}
+
+	// Set output file if LOG_FILE is specified
+	logFile := os.Getenv("LOG_FILE")
+	if logFile != "" {
+		cfg.OutputPaths = []string{logFile}
+		cfg.ErrorOutputPaths = []string{logFile}
+	} else {
+		cfg.OutputPaths = []string{"stdout"}
+		cfg.ErrorOutputPaths = []string{"stderr"}
+	}
+
+	// Common encoder config
+	cfg.EncoderConfig.TimeKey = "timestamp"
+	cfg.EncoderConfig.LevelKey = "severity"
+	cfg.EncoderConfig.NameKey = "logger"
+	cfg.EncoderConfig.CallerKey = "caller"
+	cfg.EncoderConfig.MessageKey = "message"
+	cfg.EncoderConfig.StacktraceKey = "stacktrace"
+	cfg.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
+	cfg.EncoderConfig.EncodeDuration = zapcore.StringDurationEncoder
+	cfg.EncoderConfig.EncodeCaller = zapcore.ShortCallerEncoder
+
+	// Disable sampling for now (can be re-enabled later if needed)
+	cfg.Sampling = nil
+
+	zapLogger, err := cfg.Build()
+	if err != nil {
+		panic(fmt.Sprintf("Failed to build logger: %v", err))
+	}
+
+	return &Logger{Logger: zapLogger.Named(name)}
+}
+
+// parseLevel parses log level string including our custom "trace" level
+func parseLevel(level string) (zapcore.Level, error) {
+	switch strings.ToLower(level) {
+	case "trace":
+		return TraceLevel, nil
+	case "debug":
+		return zapcore.DebugLevel, nil
+	case "info":
+		return zapcore.InfoLevel, nil
+	case "warn", "warning":
+		return zapcore.WarnLevel, nil
+	case "error":
+		return zapcore.ErrorLevel, nil
+	default:
+		return zapcore.InfoLevel, fmt.Errorf("unknown log level: %s", level)
+	}
+}
+
+// Override Sugar to return our custom SugaredLogger
+func (l *Logger) Sugar() *SugaredLogger {
+	return &SugaredLogger{SugaredLogger: l.Logger.Sugar()}
+}
+
+// Override Named to return our custom Logger
+func (l *Logger) Named(name string) *Logger {
+	return &Logger{Logger: l.Logger.Named(name)}
+}
+
+// Override With to return our custom Logger
+func (l *Logger) With(fields ...zap.Field) *Logger {
+	return &Logger{Logger: l.Logger.With(fields...)}
+}
+
+// Override WithOptions to return our custom Logger
+func (l *Logger) WithOptions(opts ...zap.Option) *Logger {
+	return &Logger{Logger: l.Logger.WithOptions(opts...)}
+}
+
+// Add Trace method to Logger
+func (l *Logger) Trace(msg string, fields ...zap.Field) {
+	l.Log(TraceLevel, msg, fields...)
+}
+
+// Add Trace method to SugaredLogger
+func (s *SugaredLogger) Trace(args ...any) {
+	s.Log(TraceLevel, args...)
+}
+
+// Add Tracew method to SugaredLogger
+func (s *SugaredLogger) Tracew(msg string, keysAndValues ...any) {
+	s.Logw(TraceLevel, msg, keysAndValues...)
+}
+
+// Override With to return our custom SugaredLogger
+func (s *SugaredLogger) With(args ...any) *SugaredLogger {
+	return &SugaredLogger{SugaredLogger: s.SugaredLogger.With(args...)}
+}
+
+// Override Named to return our custom SugaredLogger
+func (s *SugaredLogger) Named(name string) *SugaredLogger {
+	return &SugaredLogger{SugaredLogger: s.SugaredLogger.Named(name)}
+}