diff --git a/custom_components/ohme/api_client.py b/custom_components/ohme/api_client.py
index 00709a0..41ff9ab 100644
--- a/custom_components/ohme/api_client.py
+++ b/custom_components/ohme/api_client.py
@@ -194,7 +194,7 @@ async def async_max_charge(self):
         result = await self._put_request(f"/v1/chargeSessions/{self._serial}/rule?maxCharge=true")
         return bool(result)
 
-    async def async_apply_charge_rule(self, max_price=None, target_time=None, target_percent=None, pre_condition=None, pre_condition_length=None):
+    async def async_apply_session_rule(self, max_price=None, target_time=None, target_percent=None, pre_condition=None, pre_condition_length=None):
         """Apply rule to ongoing charge/stop max charge."""
         # Check every property. If we've provided it, use that. If not, use the existing.
         if max_price is None:
diff --git a/custom_components/ohme/number.py b/custom_components/ohme/number.py
index ef8ceb7..44d80ce 100644
--- a/custom_components/ohme/number.py
+++ b/custom_components/ohme/number.py
@@ -4,7 +4,7 @@
 from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.core import callback, HomeAssistant
 from .const import DOMAIN, DATA_CLIENT, DATA_COORDINATORS, COORDINATOR_CHARGESESSIONS, COORDINATOR_SCHEDULES
-
+from .utils import session_in_progress
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -50,15 +50,15 @@ def unique_id(self):
 
     async def async_set_native_value(self, value: float) -> None:
         """Update the current value."""
-        # If disconnected, update top rule. If not, apply rule to current session
-        if self.coordinator.data and self.coordinator.data['mode'] == "DISCONNECTED":
-            await self._client.async_update_schedule(target_percent=int(value))
+        # If session in progress, update this session, if not update the first schedule
+        if session_in_progress(self.coordinator.data):
+            await self._client.async_apply_session_rule(target_percent=int(value))
             await asyncio.sleep(1)
-            await self.coordinator_schedules.async_refresh()
+            await self.coordinator.async_refresh()
         else:
-            await self._client.async_apply_charge_rule(target_percent=int(value))
+            await self._client.async_update_schedule(target_percent=int(value))
             await asyncio.sleep(1)
-            await self.coordinator.async_refresh()
+            await self.coordinator_schedules.async_refresh()
 
     @property
     def icon(self):
@@ -68,9 +68,9 @@ def icon(self):
     @property
     def native_value(self):
         """Get value from data returned from API by coordinator"""
-        if self.coordinator.data and self.coordinator.data['appliedRule'] and self.coordinator.data['mode'] != "PENDING_APPROVAL" and self.coordinator.data['mode'] != "DISCONNECTED":
-            target = round(
-                self.coordinator.data['appliedRule']['targetPercent'])
+        # Set with the same logic as reading
+        if session_in_progress(self.coordinator.data):
+            target = round(self.coordinator.data['appliedRule']['targetPercent'])
         elif self.coordinator_schedules.data:
             target = round(self.coordinator_schedules.data['targetPercent'])
 
diff --git a/custom_components/ohme/switch.py b/custom_components/ohme/switch.py
index 7a37621..404ff66 100644
--- a/custom_components/ohme/switch.py
+++ b/custom_components/ohme/switch.py
@@ -161,7 +161,7 @@ async def async_turn_on(self):
     async def async_turn_off(self):
         """Stop max charging.
            We are not changing anything, just applying the last rule. No need to supply anything."""
-        await self._client.async_apply_charge_rule()
+        await self._client.async_apply_session_rule()
 
         await asyncio.sleep(1)
         await self.coordinator.async_refresh()
diff --git a/custom_components/ohme/time.py b/custom_components/ohme/time.py
index c8814f4..6eb85b4 100644
--- a/custom_components/ohme/time.py
+++ b/custom_components/ohme/time.py
@@ -5,6 +5,7 @@
 from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.core import callback, HomeAssistant
 from .const import DOMAIN, DATA_CLIENT, DATA_COORDINATORS, COORDINATOR_CHARGESESSIONS, COORDINATOR_SCHEDULES
+from .utils import session_in_progress
 from datetime import time as dt_time
 
 _LOGGER = logging.getLogger(__name__)
@@ -52,17 +53,15 @@ def unique_id(self):
 
     async def async_set_value(self, value: dt_time) -> None:
         """Update the current value."""
-        # If disconnected, update top rule. If not, apply rule to current session
-        if self.coordinator.data and self.coordinator.data['mode'] == "DISCONNECTED":
-            await self._client.async_update_schedule(target_time=(int(value.hour), int(value.minute)))
+        # If session in progress, update this session, if not update the first schedule
+        if session_in_progress(self.coordinator.data):
+            await self._client.async_apply_session_rule(target_time=(int(value.hour), int(value.minute)))
             await asyncio.sleep(1)
-            await self.coordinator_schedules.async_refresh()
+            await self.coordinator.async_refresh()
         else:
-            await self._client.async_apply_charge_rule(target_time=(int(value.hour), int(value.minute)))
+            await self._client.async_update_schedule(target_time=(int(value.hour), int(value.minute)))
             await asyncio.sleep(1)
-            await self.coordinator.async_refresh()
-
-        
+            await self.coordinator_schedules.async_refresh()
 
     @property
     def icon(self):
@@ -72,9 +71,9 @@ def icon(self):
     @property
     def native_value(self):
         """Get value from data returned from API by coordinator"""
-        # If we are not pending approval or disconnected, return in progress charge rule
+        # Set with the same logic as reading
         target = None
-        if self.coordinator.data and self.coordinator.data['appliedRule'] and self.coordinator.data['mode'] != "PENDING_APPROVAL" and self.coordinator.data['mode'] != "DISCONNECTED":
+        if session_in_progress(self.coordinator.data):
             target = self.coordinator.data['appliedRule']['targetTime']
         elif self.coordinator_schedules.data:
             target = self.coordinator_schedules.data['targetTime']
diff --git a/custom_components/ohme/utils.py b/custom_components/ohme/utils.py
index 17159fc..8053374 100644
--- a/custom_components/ohme/utils.py
+++ b/custom_components/ohme/utils.py
@@ -73,3 +73,16 @@ def time_next_occurs(hour, minute):
         target = target + timedelta(days=1)
 
     return target
+
+def session_in_progress(data):
+    """Is there a session in progress?
+       Used to check if we should update the current session rather than the first schedule."""
+    # Default to False with no data
+    if not data:
+        return False
+    
+    # Car disconnected or pending approval, we should update the schedule
+    if data['mode'] == "DISCONNECTED" or data['mode'] == "PENDING_APPROVAL":
+        return False
+    
+    return True