diff --git a/library.json b/library.json
index e3a027b..0162e48 100644
--- a/library.json
+++ b/library.json
@@ -1,6 +1,6 @@
 {
     "name": "Wifi",
-    "version": "1.0.0",
+    "version": "1.0.1",
     "description": "Helper classes for handling wifi connectiviy and OTA",
     "repository":
     {
diff --git a/src/ota.h b/src/ota.h
index 148c5ee..def614f 100644
--- a/src/ota.h
+++ b/src/ota.h
@@ -2,16 +2,13 @@
 
 namespace  Ota {
 
-    void loop();
-    Task tLoop(TASK_IMMEDIATE, TASK_FOREVER, loop, &ts, true);
+    Task tLoop(TASK_IMMEDIATE, TASK_FOREVER, [] {ArduinoOTA.handle();}, &ts, true);
 
-    void setup() {
+    void setup(void (*onStart)() = nullptr, void (*onEnd)() = nullptr, void (*onError)() = nullptr) {
         ArduinoOTA.onStart(
-            []() {
+            [onStart]() {
                 Serial.println("Starting OTA");
-                Mqtt::publishCleanupConfig();
-                delay(2000);
-                Mqtt::disconnect();
+                if (onStart) onStart();
             });
         ArduinoOTA.onEnd(
             []() {
@@ -32,8 +29,4 @@ namespace  Ota {
             });
         ArduinoOTA.begin();
     }
-
-    void loop() {
-        ArduinoOTA.handle();
-    }
 }