diff --git a/cmd/picoclaw/internal/gateway/command.go b/cmd/picoclaw/internal/gateway/command.go index 66a56f9ce..d82b942c1 100644 --- a/cmd/picoclaw/internal/gateway/command.go +++ b/cmd/picoclaw/internal/gateway/command.go @@ -6,6 +6,7 @@ import ( func NewGatewayCommand() *cobra.Command { var debug bool + var logFilter string cmd := &cobra.Command{ Use: "gateway", @@ -13,11 +14,12 @@ func NewGatewayCommand() *cobra.Command { Short: "Start picoclaw gateway", Args: cobra.NoArgs, RunE: func(_ *cobra.Command, _ []string) error { - return gatewayCmd(debug) + return gatewayCmd(debug, logFilter) }, } cmd.Flags().BoolVarP(&debug, "debug", "d", false, "Enable debug logging") + cmd.Flags().StringVar(&logFilter, "log-filter", "", "Filter logs by component (comma separated)") return cmd } diff --git a/cmd/picoclaw/internal/gateway/helpers.go b/cmd/picoclaw/internal/gateway/helpers.go index baa489b92..a02eb853d 100644 --- a/cmd/picoclaw/internal/gateway/helpers.go +++ b/cmd/picoclaw/internal/gateway/helpers.go @@ -37,12 +37,17 @@ import ( "github.com/sipeed/picoclaw/pkg/tools" ) -func gatewayCmd(debug bool) error { +func gatewayCmd(debug bool, logFilter string) error { if debug { logger.SetLevel(logger.DEBUG) fmt.Println("🔍 Debug mode enabled") } + if logFilter != "" { + logger.SetComponentFilter(logFilter) + fmt.Printf("🔍 Log filter enabled: %s\n", logFilter) + } + cfg, err := internal.LoadConfig() if err != nil { return fmt.Errorf("error loading config: %w", err) diff --git a/pkg/logger/logger.go b/pkg/logger/logger.go index 56dc87a53..5383b1b07 100644 --- a/pkg/logger/logger.go +++ b/pkg/logger/logger.go @@ -34,6 +34,10 @@ var ( logger *Logger once sync.Once mu sync.RWMutex + + // componentFilter is a list of components to allow. + // If empty, all components are allowed. + componentFilter map[string]bool ) type Logger struct { @@ -67,6 +71,25 @@ func GetLevel() LogLevel { return currentLevel } +func SetComponentFilter(filter string) { + mu.Lock() + defer mu.Unlock() + + if filter == "" { + componentFilter = nil + return + } + + componentFilter = make(map[string]bool) + parts := strings.Split(filter, ",") + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + componentFilter[p] = true + } + } +} + func EnableFileLogging(filePath string) error { mu.Lock() defer mu.Unlock() @@ -97,7 +120,18 @@ func DisableFileLogging() { } func logMessage(level LogLevel, component string, message string, fields map[string]any) { - if level < currentLevel { + mu.RLock() + // Check filter first if it exists + if componentFilter != nil && component != "" { + if !componentFilter[component] { + mu.RUnlock() + return + } + } + lvl := currentLevel + mu.RUnlock() + + if level < lvl { return } diff --git a/pkg/logger/logger_filter_test.go b/pkg/logger/logger_filter_test.go new file mode 100644 index 000000000..97aed848a --- /dev/null +++ b/pkg/logger/logger_filter_test.go @@ -0,0 +1,60 @@ +package logger + +import ( + "bytes" + "log" + "os" + "strings" + "testing" +) + +func TestSetComponentFilter(t *testing.T) { + // Capture log output + var buf bytes.Buffer + log.SetOutput(&buf) + defer log.SetOutput(os.Stderr) + + // Reset filter + SetComponentFilter("") + + // Test 1: No filter + InfoC("comp1", "msg1") + if !strings.Contains(buf.String(), "msg1") { + t.Error("Expected msg1 to be logged") + } + buf.Reset() + + // Test 2: Filter comp1 + SetComponentFilter("comp1") + InfoC("comp1", "msg2") // Should be logged + InfoC("comp2", "msg3") // Should NOT be logged + + output := buf.String() + if !strings.Contains(output, "msg2") { + t.Error("Expected msg2 to be logged") + } + if strings.Contains(output, "msg3") { + t.Error("Expected msg3 NOT to be logged") + } + buf.Reset() + + // Test 3: Multiple filters + SetComponentFilter("comp1,comp2") + InfoC("comp1", "msg4") // Logged + InfoC("comp2", "msg5") // Logged + InfoC("comp3", "msg6") // Not logged + + output = buf.String() + if !strings.Contains(output, "msg4") { + t.Error("Expected msg4 to be logged") + } + if !strings.Contains(output, "msg5") { + t.Error("Expected msg5 to be logged") + } + if strings.Contains(output, "msg6") { + t.Error("Expected msg6 NOT to be logged") + } + + // Reset filter at end + SetComponentFilter("") +}