fix: Prevent panic on context cancellation during agent shutdown

The leader election callback in the agent was previously panicking if the agent encountered an error, including context.Canceled which is expected during graceful shutdown.

This commit modifies the callback to check if the error is context.Canceled. If so, it logs an informational message and exits gracefully. Other unexpected errors will still cause a panic.

Author: google-labs-jules[bot]

cmd/agent/run.go

@@ -42,9 +42,19 @@ func run(options *config.AgentOptions) error {
 
 	callback := func(ctx context.Context) {
 		if err := agent.Run(ctx); err != nil {
-			panic(err)
+			// Check if the error is context.Canceled, which is expected on graceful shutdown.
+			if errors.Is(err, context.Canceled) {
+				logrus.Info("Agent run completed due to context cancellation.")
+			} else {
+				// For any other error, it's unexpected, so panic.
+				logrus.Errorf("Agent run failed with unexpected error: %v", err)
+				panic(err)
+			}
 		}
+		// Wait for the context to be done, ensuring the leader election logic
+		// holds the leadership until the context is fully cancelled.
 		<-ctx.Done()
+		logrus.Info("Leader election callback completed as context is done.")
 	}
 
 	httpServerOptions := config.HTTPServerOptions{