Forecast a Sud schedule through its confirm steps, not just up to the first one

A user_wait_for_continue step used to stop SudForecastEstimator.estimate()
dead, so the GUI's forecast plot only ever showed the first segment on Load.
Model the wait as a zero-delay auto-confirm instead, so the whole schedule's
projected curve is visible right away; correct that assumption piecewise as
real confirmations actually happen, anchored at the real elapsed time and
temperature.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DkkuG48uHFCGKe6dPSERFk

tasks/sud.py

@@ -1,4 +1,5 @@
 import asyncio
+import bisect
 from tasks import ATask
 from ws.message import MsgIo
 from utils.value import ChangedFloat
@@ -29,18 +30,29 @@ class SudTask(ATask):
 		# build a SudTask without one still work; the server always passes
 		# one.
 		self.forecast_estimator = forecast_estimator
-		# Accumulated forecast across all piecewise segments computed so
-		# far (see send_forecast()/_continue_forecast_after_confirm()) -
-		# T in simulated seconds, Theta in degrees, parallel lists, both
-		# growing monotonically as segments get appended; never re-walked
-		# from scratch once a run is in progress, except for the one
-		# genuinely unforecastable case (a real user confirmation).
+		# The forecast as last computed/corrected (see send_forecast()/
+		# _continue_forecast_after_confirm()) - T in simulated seconds,
+		# Theta in degrees, parallel lists, both monotonically growing as
+		# corrections get spliced in. Covers the whole schedule from the
+		# very first Load, including through steps requiring user
+		# confirmation - those are simulated as a zero-delay auto-confirm
+		# (see components/sud_forecast.py's SudForecastEstimator.
+		# estimate()) rather than left unforecast.
 		self.forecast_t = []
 		self.forecast_theta = []
-		# Whether the most recently computed segment stopped at a step
-		# requiring user confirmation rather than reaching the
-		# schedule's actual end - see _continue_forecast_after_confirm().
-		self.forecast_waiting_for_confirm = False
+		# Whether the forecast above actually reaches the schedule's real
+		# end (sent as 'Finished') - False only in the pathological case
+		# of a step whose target can never be reached (see
+		# components/sud_forecast.py's MAX_TICKS).
+		self.forecast_finished = True
+		# Where each user-confirmation step's zero-delay assumption sits
+		# in the forecast timeline above, keyed by the schedule's
+		# (absolute) step index - see SudForecastEstimator.estimate()'s
+		# confirm_points. Consulted by _continue_forecast_after_confirm()
+		# to know where to cut the optimistic guess loose and splice in a
+		# freshly anchored simulation once that confirmation actually
+		# happens for real.
+		self.forecast_confirm_marks = {}
 		msg_handler.set_recv_handler(self.recv)
 
 	def apply_plant_params(self, step):
@@ -119,12 +131,18 @@ class SudTask(ATask):
 		asyncio.create_task(self.send({'Elapsed': value}))
 
 	async def send_forecast(self, doc):
-		"""Computes and sends the first piecewise segment of doc's
-		forecast - from the very start, up to either the schedule's end
-		or its first step requiring user confirmation (see
-		components/sud_forecast.py's SudForecastEstimator.estimate()).
-		Always a fresh start: resets any segments accumulated for
-		whatever was loaded before."""
+		"""Computes and sends the full forecast for doc, start to finish -
+		including through every step requiring user confirmation, which
+		is modeled as a zero-delay auto-confirm rather than left
+		unforecast (see components/sud_forecast.py's
+		SudForecastEstimator.estimate()) - so the whole schedule's
+		projected curve is visible right away instead of stopping at the
+		first one. Always a fresh start: discards whatever forecast was
+		accumulated for the previously loaded schedule.
+
+		Those zero-delay assumptions get corrected piecewise as real
+		confirmations actually happen - see
+		_continue_forecast_after_confirm()."""
 		if self.forecast_estimator is None:
 			return
 		# Runs the simulation in a worker thread - it's CPU-bound and can
@@ -132,34 +150,53 @@ class SudTask(ATask):
 		# otherwise stall every other task (heater, sensor, ...) for that
 		# whole window.
 		loop = asyncio.get_event_loop()
-		t, theta, final_state = await loop.run_in_executor(None, self.forecast_estimator.estimate, doc)
+		t, theta, final_state, confirm_points = await loop.run_in_executor(None, self.forecast_estimator.estimate, doc)
 		self.forecast_t = t
 		self.forecast_theta = theta
-		self.forecast_waiting_for_confirm = (final_state == SudState.WAIT_USER)
+		self.forecast_finished = (final_state == SudState.DONE)
+		self.forecast_confirm_marks = dict(confirm_points)
 		await self._send_forecast()
 
 	async def _send_forecast(self):
 		await self.send({'Forecast': {
 			'T': self.forecast_t,
 			'Theta': self.forecast_theta,
-			'Finished': not self.forecast_waiting_for_confirm,
+			'Finished': self.forecast_finished,
 		}})
 
-	async def _continue_forecast_after_confirm(self):
-		"""Computes and appends the next piecewise segment once a real
-		user confirmation has actually happened, rather than guessing at
-		the delay - see SudForecastEstimator.estimate()'s docstring for
-		why. Anchored at the real elapsed time (self.sud.elapsed), so the
-		unforecastable wait shows up honestly as a gap in the forecast
-		rather than as zero delay, and at the real current temperature
-		(more trustworthy than whatever the now-superseded previous
-		segment's simulation predicted it would be by now)."""
-		if self.forecast_estimator is None or not self.forecast_waiting_for_confirm:
+	async def _continue_forecast_after_confirm(self, confirmed_index):
+		"""Corrects the optimistic, zero-delay guess send_forecast() (or a
+		previous call to this method) made at confirmed_index's
+		user-confirmation step, now that the confirmation has actually
+		happened for real - see SudForecastEstimator.estimate()'s
+		docstring for why that assumption can't just be trusted as-is.
+		Truncates the forecast right back to that point and splices in a
+		freshly anchored simulation of the rest of the schedule, anchored
+		at the real elapsed time (self.sud.elapsed) - so an actual delay
+		shows up honestly as a gap rather than as the assumed zero - and
+		the real current temperature (more trustworthy than whatever the
+		now-superseded guess predicted it would be by now).
+
+		No-op if confirmed_index was never part of a computed forecast in
+		the first place (e.g. the estimator isn't configured)."""
+		if self.forecast_estimator is None:
 			return
+		mark_t = self.forecast_confirm_marks.pop(confirmed_index, None)
+		if mark_t is None:
+			return
+		# Drop the now-stale tail (everything beyond the confirmation
+		# point) - both the curve itself and any confirm marks that fell
+		# within it, since they're about to be replaced by fresh ones.
+		cut = bisect.bisect_right(self.forecast_t, mark_t)
+		self.forecast_t = self.forecast_t[:cut]
+		self.forecast_theta = self.forecast_theta[:cut]
+		self.forecast_confirm_marks = {idx: t for idx, t in self.forecast_confirm_marks.items() if t <= mark_t}
+
 		schedule = self.sud.schedule
 		index = self.sud.index
 		if not (0 <= index < len(schedule)):
-			self.forecast_waiting_for_confirm = False
+			self.forecast_finished = True
+			await self._send_forecast()
 			return
 		doc = {
 			'Name': self.sud.name,
@@ -171,17 +208,22 @@ class SudTask(ATask):
 		start_theta = self.tc.get_theta_ist()
 		real_elapsed = self.sud.elapsed
 		loop = asyncio.get_event_loop()
-		t, theta, final_state = await loop.run_in_executor(None, self.forecast_estimator.estimate, doc, start_theta)
-		# Bridge the gap between where the previous segment's own
-		# simulated time left off and the real elapsed time the confirm
-		# actually happened at, holding flat at its last temperature -
-		# then append the new segment, offset to start exactly there.
+		t, theta, final_state, confirm_points = await loop.run_in_executor(None, self.forecast_estimator.estimate, doc, start_theta)
+		# Bridge the gap between the confirmation point and the real
+		# elapsed time it actually happened at, holding flat at its last
+		# temperature - then append the new segment, offset to start
+		# exactly there.
 		if self.forecast_t and self.forecast_t[-1] < real_elapsed:
 			self.forecast_t.append(real_elapsed)
 			self.forecast_theta.append(self.forecast_theta[-1])
 		self.forecast_t.extend(real_elapsed + seconds for seconds in t)
 		self.forecast_theta.extend(theta)
-		self.forecast_waiting_for_confirm = (final_state == SudState.WAIT_USER)
+		# confirm_points' indices/times are relative to this sub-schedule
+		# (starting fresh at doc['steps'][0]) - rebase both onto the real
+		# schedule's absolute indices and the master forecast timeline.
+		self.forecast_confirm_marks.update(
+			(index + local_index, real_elapsed + local_t) for local_index, local_t in confirm_points)
+		self.forecast_finished = (final_state == SudState.DONE)
 		await self._send_forecast()
 
 	async def recv(self, data):
@@ -189,8 +231,9 @@ class SudTask(ATask):
 			if 'Start' in pair[0]:
 				self.sud.start()
 			elif 'Confirm' in pair[0]:
+				confirmed_index = self.sud.index
 				self.sud.confirm()
-				asyncio.create_task(self._continue_forecast_after_confirm())
+				asyncio.create_task(self._continue_forecast_after_confirm(confirmed_index))
 			elif 'Pause' in pair[0]:
 				self.sud.pause()
 			elif 'Stop' in pair[0]: