Improve async log handling and fix test compatibility

blink-so[bot] · kylecarbs · blink-so[bot] · commit f6b43a0c9dea · 2025-06-10T01:53:17.000Z
- Add proper context management to logger goroutines
- Fix log level conversion in fake API PatchLogs method
- Improve goroutine lifecycle with cancellation support
- Use ticker-based approach for more reliable log sending
- Continue on PatchLogs errors instead of exiting goroutine

Co-authored-by: kylecarbs &lt;7122116+kylecarbs@users.noreply.github.com&gt;
diff --git a/logger.go b/logger.go
@@ -517,6 +517,7 @@ type agentLoggerLifecycle struct {
 	client *agentsdk.Client
 	sourceID uuid.UUID
 	timer  *quartz.Timer
+	cancel context.CancelFunc
 }
 
 func (lq *logQueuer) work(ctx context.Context) {
@@ -563,6 +564,9 @@ func (lq *logQueuer) ensureLogger(ctx context.Context, token string) {
 		lq.logger.Debug(ctx, "failed to create log source", slog.Error(err))
 	}
 
+	// Create a context for this logger that can be cancelled
+	loggerCtx, cancel := context.WithCancel(ctx)
+
 	timer := lq.clock.AfterFunc(lq.loggerTTL, func() {
 		lq.deleteLogger(token)
 	})
@@ -571,23 +575,31 @@ func (lq *logQueuer) ensureLogger(ctx context.Context, token string) {
 		client:   agentClient,
 		sourceID: sourceID,
 		timer:    timer,
+		cancel:   cancel,
 	}
 
 	go func() {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+		
 		for {
-			logs := lq.logCache.flush(token)
-			if len(logs) == 0 {
-				time.Sleep(time.Second)
-				continue
-			}
-
-			err := agentClient.PatchLogs(ctx, agentsdk.PatchLogs{
-				LogSourceID: sourceID,
-				Logs:        logs,
-			})
-			if err != nil {
-				lq.logger.Error(ctx, "patch agent logs", slog.Error(err))
+			select {
+			case <-loggerCtx.Done():
 				return
+			case <-ticker.C:
+				logs := lq.logCache.flush(token)
+				if len(logs) == 0 {
+					continue
+				}
+
+				err := agentClient.PatchLogs(loggerCtx, agentsdk.PatchLogs{
+					LogSourceID: sourceID,
+					Logs:        logs,
+				})
+				if err != nil {
+					lq.logger.Error(loggerCtx, "patch agent logs", slog.Error(err))
+					// Don't return on error, keep trying
+				}
 			}
 		}
 	}()
@@ -603,6 +615,7 @@ func (lq *logQueuer) deleteLogger(token string) {
 	}
 
 	lifecycle.timer.Stop()
+	lifecycle.cancel() // Cancel the context to stop the goroutine
 	delete(lq.loggers, token)
 }
 
diff --git a/logger_test.go b/logger_test.go
@@ -475,14 +475,32 @@ func (f *fakeAgentAPI) PatchLogs(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
+
 	
 	// Convert agentsdk.Log to proto.Log for the channel
 	protoLogs := make([]*proto.Log, len(req.Logs))
 	for i, log := range req.Logs {
+		// Simple log level mapping
+		var level proto.Log_Level
+		switch string(log.Level) {
+		case "trace":
+			level = 1 // Assuming TRACE = 1
+		case "debug":
+			level = 2 // Assuming DEBUG = 2
+		case "info":
+			level = 3 // Assuming INFO = 3
+		case "warn":
+			level = 4 // Assuming WARN = 4
+		case "error":
+			level = 5 // Assuming ERROR = 5
+		default:
+			level = 3 // Default to INFO
+		}
+		
 		protoLogs[i] = &proto.Log{
 			CreatedAt: timestamppb.New(log.CreatedAt),
 			Output:    log.Output,
-			Level:     proto.Log_Level(proto.Log_Level_value[string(log.Level)]),
+			Level:     level,
 		}
 	}
 	
