Add session resume flow and configurable steering dimension

Author: apartsin

app/core/config_yaml.py

@@ -18,6 +18,7 @@
     # updater: winner_average | winner_copy | linear_preference
     # feedback_mode: scalar_rating | pairwise | top_k | winner_only | approve_reject
     # seed_policy: fixed-per-round | fixed-per-candidate | fixed-per-candidate-role
+    # steering_dimension: low-dimensional steering vector size, for example 3 or 5
     # image_size: WIDTHxHEIGHT, for example 512x512
     # guidance_scale: classifier-free guidance strength, for example 7.5
     # num_inference_steps: diffusion denoising steps, for example 15 or 30

app/core/jobs.py

@@ -3,6 +3,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from enum import Enum
+from inspect import signature
 from threading import Lock
 from typing import Any, Callable
 from uuid import uuid4
@@ -62,9 +63,16 @@ async def get(self, job_id: str) -> JobRecord | None:
     def _run_sync(self, job_id: str, fn: Callable[[], Any]) -> None:
         """Execute one submitted job and persist its status transitions."""
 
-        self._update(job_id, state=JobState.running, progress=15, status_message="Running")
+        self._update(job_id, state=JobState.running, progress=12, status_message="Starting work")
         try:
-            result = fn()
+            try:
+                arity = len(signature(fn).parameters)
+            except (TypeError, ValueError):
+                arity = 0
+            if arity >= 1:
+                result = fn(lambda progress, message: self.update_progress(job_id, progress, message))
+            else:
+                result = fn()
         except Exception as exc:
             self._update(
                 job_id,
@@ -80,11 +88,16 @@ def _run_sync(self, job_id: str, fn: Callable[[], Any]) -> None:
             job_id,
             state=JobState.succeeded,
             progress=100,
-            status_message="Completed",
+            status_message="Completed successfully",
             result=payload,
             error=None,
         )
 
+    def update_progress(self, job_id: str, progress: int, status_message: str) -> None:
+        """Expose safe phase-level progress updates to long-running jobs."""
+
+        self._update(job_id, state=JobState.running, progress=progress, status_message=status_message)
+
     def _update(self, job_id: str, **changes: Any) -> None:
         """Mutate one job record in place under the manager lock."""
 

app/core/schema.py

@@ -85,6 +85,7 @@ class StrategyConfig(BaseModel):
     feedback_mode: FeedbackType = FeedbackType.scalar_rating
     seed_policy: SeedPolicy = SeedPolicy.fixed_per_round
     steering_mode: SteeringMode = SteeringMode.low_dimensional
+    steering_dimension: int = Field(default=3, ge=1, le=16)
     candidate_count: int = Field(default=5, ge=1, le=12)
     image_size: str = "512x512"
     trust_radius: float = Field(default=0.35, gt=0.0, le=1.0)
@@ -189,7 +190,7 @@ class Session(BaseModel):
     status: SessionStatus = SessionStatus.created
     basis_type: str = "random_orthonormal"
     current_round: int = 0
-    current_z: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0])
+    current_z: list[float] = Field(default_factory=list)
     incumbent_candidate_id: str | None = None
     final_selected_candidate: str | None = None
     base_embedding_cache_key: str = Field(default_factory=lambda: new_id("emb"))

app/engine/orchestrator.py

@@ -3,6 +3,7 @@
 from copy import deepcopy
 import hashlib
 import math
+from typing import Callable
 
 from app.core.config import settings
 from app.core.schema import (
@@ -62,6 +63,13 @@ def __init__(
             "linear_preference": LinearPreferenceUpdater(),
         }
 
+    @staticmethod
+    def _report_progress(progress_callback: Callable[[int, str], None] | None, progress: int, message: str) -> None:
+        """Emit a phase-level progress update when a callback is available."""
+
+        if progress_callback is not None:
+            progress_callback(progress, message)
+
     def create_experiment(self, request: ExperimentCreate) -> Experiment:
         """Create and persist a reusable experiment definition."""
 
@@ -102,6 +110,7 @@ def create_session(self, request: SessionCreate) -> Session:
             negative_prompt=request.negative_prompt,
             model_name=config.model_name,
             config=config,
+            current_z=[0.0 for _ in range(config.steering_dimension)],
             status=SessionStatus.ready,
         )
         logger.info("Created session %s for experiment %s", session.id, session.experiment_id)
@@ -118,17 +127,28 @@ def get_session(self, session_id: str) -> Session | None:
 
         return self.repository.get_session(session_id)
 
+    def list_sessions(self) -> list[Session]:
+        """Return all stored sessions ordered by recent activity."""
+
+        return self.repository.list_sessions()
+
     def get_session_rounds(self, session_id: str) -> list[Round]:
         """Return ordered rounds for a given session."""
 
         return self.repository.list_rounds_for_session(session_id)
 
-    def generate_round(self, session_id: str) -> RoundResponse:
+    def generate_round(
+        self,
+        session_id: str,
+        progress_callback: Callable[[int, str], None] | None = None,
+    ) -> RoundResponse:
         """Propose, render, persist, and return the next round of candidates."""
 
+        self._report_progress(progress_callback, 14, "Checking session readiness")
         session = self._assert_round_generation_allowed(session_id)
         sampler = self.samplers[session.config.sampler]
         round_index = session.current_round + 1
+        self._report_progress(progress_callback, 24, "Preparing round state")
         round_obj = Round(
             session_id=session.id,
             round_index=round_index,
@@ -144,6 +164,7 @@ def generate_round(self, session_id: str) -> RoundResponse:
         carried_forward = self._build_carried_forward_candidate(session)
         baseline_candidate = self._build_baseline_prompt_candidate(session)
         sampler_seed = self._seed_token(session.id, round_index, "sampler")
+        self._report_progress(progress_callback, 36, f"Sampling {session.config.candidate_count} candidate directions")
         proposed_candidates = sampler.propose(session, sampler_seed)
         proposed_candidates = self._widen_first_round_candidates(session, proposed_candidates)
         candidates = self._compose_round_candidates(
@@ -152,6 +173,7 @@ def generate_round(self, session_id: str) -> RoundResponse:
             candidate_count=session.config.candidate_count,
         )
         self._assign_candidate_seeds(session, round_index, candidates)
+        self._report_progress(progress_callback, 52, "Rendering candidate images on the model backend")
         # Render each candidate independently so future versions can tolerate
         # partial round failures without changing the orchestration contract.
         for candidate in candidates:
@@ -164,6 +186,7 @@ def generate_round(self, session_id: str) -> RoundResponse:
         round_obj.candidates = candidates
         round_obj.render_status = RenderStatus.succeeded
         round_obj.latency_ms = 15 * len(candidates)
+        self._report_progress(progress_callback, 76, "Saving rendered candidates and round state")
         session.current_round = round_index
         session.status = SessionStatus.awaiting_feedback
         session.updated_at = utc_now()
@@ -179,7 +202,9 @@ def generate_round(self, session_id: str) -> RoundResponse:
                 "candidates": [self._candidate_trace_payload(candidate) for candidate in round_obj.candidates],
             },
         )
+        self._report_progress(progress_callback, 90, "Refreshing trace report and replay data")
         self.generate_trace_report(session.id)
+        self._report_progress(progress_callback, 98, "Round ready for review")
         return RoundResponse(
             round_id=round_obj.id,
             candidate_metadata=round_obj.candidates,
@@ -191,20 +216,29 @@ def generate_round(self, session_id: str) -> RoundResponse:
             },
         )
 
-    def submit_feedback(self, round_id: str, request: FeedbackRequest) -> FeedbackResponse:
+    def submit_feedback(
+        self,
+        round_id: str,
+        request: FeedbackRequest,
+        progress_callback: Callable[[int, str], None] | None = None,
+    ) -> FeedbackResponse:
         """Normalize feedback, update state, and persist the new incumbent."""
 
+        self._report_progress(progress_callback, 14, "Checking round readiness for feedback")
         round_obj, session = self._assert_feedback_submission_allowed(round_id, request)
+        self._report_progress(progress_callback, 30, "Normalizing and validating user preferences")
         feedback = normalize_feedback(round_id, request)
         self._validate_feedback_against_round(round_obj, feedback)
         updater = self.updaters[session.config.updater]
+        self._report_progress(progress_callback, 52, "Updating the steering model from your feedback")
         next_z, update_summary = updater.update(session, round_obj.candidates, feedback)
         round_obj.feedback_events.append(feedback)
         round_obj.update_summary = update_summary
         session.current_z = next_z
         session.incumbent_candidate_id = update_summary["winner_candidate_id"]
         session.status = SessionStatus.ready
         session.updated_at = utc_now()
+        self._report_progress(progress_callback, 72, "Saving updated session state")
         self.repository.save_round(round_obj)
         self.repository.save_session(session)
         logger.info("Applied feedback to round %s for session %s", round_obj.id, session.id)
@@ -221,7 +255,9 @@ def submit_feedback(self, round_id: str, request: FeedbackRequest) -> FeedbackRe
                 "next_incumbent_state": next_z,
             },
         )
+        self._report_progress(progress_callback, 90, "Refreshing trace report with the new preference outcome")
         self.generate_trace_report(session.id)
+        self._report_progress(progress_callback, 98, "Feedback applied and next round unlocked")
         return FeedbackResponse(update_summary=update_summary, next_incumbent_state=next_z)
 
     def export_replay(self, session_id: str) -> dict:

app/frontend/static/styles.css

@@ -44,6 +44,10 @@ h1, h2, h3, p { margin-top: 0; }
 .card { background: var(--paper); border: 1px solid var(--line); border-radius: 20px; padding: 22px; margin-bottom: 20px; box-shadow: 0 10px 30px rgba(31,26,23,0.05); }
 .section-head { display: flex; align-items: baseline; justify-content: space-between; gap: 12px; margin-bottom: 16px; }
 .actions { display: flex; gap: 12px; flex-wrap: wrap; margin-top: 16px; }
+.inline-actions {
+  margin-top: 0;
+  gap: 8px;
+}
 .setup-steps {
   display: grid;
   grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
@@ -218,6 +222,38 @@ textarea { min-height: 100px; resize: vertical; }
   color: var(--accent);
   margin-bottom: 10px;
 }
+.round-block {
+  margin-top: 14px;
+  border: 1px solid var(--line);
+  border-radius: 18px;
+  background: white;
+  overflow: hidden;
+}
+.round-summary {
+  display: flex;
+  justify-content: space-between;
+  gap: 12px;
+  padding: 14px 16px;
+  cursor: pointer;
+  list-style: none;
+  font-weight: 600;
+}
+.round-summary::-webkit-details-marker {
+  display: none;
+}
+.round-summary-title {
+  color: var(--ink);
+}
+.round-summary-meta {
+  color: var(--muted);
+  font-weight: 400;
+}
+.round-body {
+  padding: 0 16px 16px;
+}
+.round-footer-actions {
+  margin-top: 18px;
+}
 code { background: #f2ebdf; padding: 2px 6px; border-radius: 8px; }
 .trace-card { background: linear-gradient(180deg, #fffaf2, #f4efe7); }
 .trace-log {

app/frontend/templates/index.html

@@ -17,6 +17,41 @@ <h1>Interactive prompt-embedding steering research prototype</h1>
           <a class="button secondary" href="/diagnostics/view">Runtime diagnostics</a>
         </div>
       </section>
+      <section class="card">
+        <div class="section-head">
+          <h2>Resume sessions</h2>
+          <span>{{ sessions|length }} recent</span>
+        </div>
+        {% if sessions %}
+          <table class="table">
+            <thead>
+              <tr>
+                <th>Prompt</th>
+                <th>Status</th>
+                <th>Round</th>
+                <th>Continue</th>
+              </tr>
+            </thead>
+            <tbody>
+            {% for session in sessions %}
+              <tr>
+                <td>{{ session.prompt }}</td>
+                <td>{{ session.status }}</td>
+                <td>{{ session.current_round }}</td>
+                <td>
+                  <div class="actions inline-actions">
+                    <a class="button secondary" href="/sessions/{{ session.id }}/view">Resume session</a>
+                    <a class="button secondary" href="/sessions/{{ session.id }}/replay-view">Replay</a>
+                  </div>
+                </td>
+              </tr>
+            {% endfor %}
+            </tbody>
+          </table>
+        {% else %}
+          <p>No sessions yet. Create your first session to make it resumable from this dashboard.</p>
+        {% endif %}
+      </section>
       <section class="card">
         <div class="section-head">
           <h2>Experiments</h2>