From 32ebcfe09145bc9f00c5d0b7f39d81b7b9fb151b Mon Sep 17 00:00:00 2001
From: Roman Pszonka <roman@pszonka.org>
Date: Wed, 3 Dec 2025 01:19:19 +0000
Subject: [PATCH] fix generating 2d and 3d maps

---
 .../flight_blender/flight_blender_client.py   |  6 +++-
 .../core/execution/scenario_runner.py         | 28 +++++++++++++++++++
 .../scenarios/registry.py                     |  4 +++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/openutm_verification/core/clients/flight_blender/flight_blender_client.py b/src/openutm_verification/core/clients/flight_blender/flight_blender_client.py
index 4c16d6dc..535e736a 100644
--- a/src/openutm_verification/core/clients/flight_blender/flight_blender_client.py
+++ b/src/openutm_verification/core/clients/flight_blender/flight_blender_client.py
@@ -12,7 +12,7 @@
     BaseBlenderAPIClient,
 )
 from openutm_verification.core.execution.config_models import DataFiles
-from openutm_verification.core.execution.scenario_runner import scenario_step
+from openutm_verification.core.execution.scenario_runner import ScenarioContext, scenario_step
 from openutm_verification.core.reporting.reporting_models import (
     Status,
     StepResult,
@@ -761,6 +761,10 @@ def setup_flight_declaration(self, flight_declaration_path: str, trajectory_path
 
         self.telemetry_states = telemetry_states
 
+        # Store data in ScenarioContext for reporting
+        ScenarioContext.set_flight_declaration_data(flight_declaration)
+        ScenarioContext.set_telemetry_data(telemetry_states)
+
         upload_result = self.upload_flight_declaration(flight_declaration)
 
         if upload_result.status == Status.FAIL:
diff --git a/src/openutm_verification/core/execution/scenario_runner.py b/src/openutm_verification/core/execution/scenario_runner.py
index 97f1402a..44d0b905 100644
--- a/src/openutm_verification/core/execution/scenario_runner.py
+++ b/src/openutm_verification/core/execution/scenario_runner.py
@@ -19,6 +19,8 @@
 class ScenarioState:
     steps: List[StepResult[Any]] = field(default_factory=list)
     active: bool = False
+    flight_declaration_data: Optional[Any] = None
+    telemetry_data: Optional[Any] = None
 
 
 _scenario_state: contextvars.ContextVar[Optional[ScenarioState]] = contextvars.ContextVar("scenario_state", default=None)
@@ -46,6 +48,18 @@ def add_result(cls, result: StepResult[Any]) -> None:
         if state and state.active:
             state.steps.append(result)
 
+    @classmethod
+    def set_flight_declaration_data(cls, data: Any) -> None:
+        state = _scenario_state.get()
+        if state and state.active:
+            state.flight_declaration_data = data
+
+    @classmethod
+    def set_telemetry_data(cls, data: Any) -> None:
+        state = _scenario_state.get()
+        if state and state.active:
+            state.telemetry_data = data
+
     @property
     def steps(self) -> List[StepResult[Any]]:
         if self._state:
@@ -53,6 +67,20 @@ def steps(self) -> List[StepResult[Any]]:
         state = _scenario_state.get()
         return state.steps if state else []
 
+    @property
+    def flight_declaration_data(self) -> Optional[Any]:
+        if self._state:
+            return self._state.flight_declaration_data
+        state = _scenario_state.get()
+        return state.flight_declaration_data if state else None
+
+    @property
+    def telemetry_data(self) -> Optional[Any]:
+        if self._state:
+            return self._state.telemetry_data
+        state = _scenario_state.get()
+        return state.telemetry_data if state else None
+
 
 class StepDecorator(Protocol):
     @overload
diff --git a/src/openutm_verification/scenarios/registry.py b/src/openutm_verification/scenarios/registry.py
index 487c95a4..65aec41f 100644
--- a/src/openutm_verification/scenarios/registry.py
+++ b/src/openutm_verification/scenarios/registry.py
@@ -39,6 +39,8 @@ def _run_scenario_simple(scenario_id: str, func: Callable, args, kwargs) -> Scen
                 return result
 
             steps = ctx.steps
+            flight_declaration_data = ctx.flight_declaration_data
+            telemetry_data = ctx.telemetry_data
 
         final_status = Status.PASS if all(s.status == Status.PASS for s in steps) else Status.FAIL
         total_duration = sum(s.duration for s in steps)
@@ -47,6 +49,8 @@ def _run_scenario_simple(scenario_id: str, func: Callable, args, kwargs) -> Scen
             status=final_status,
             duration_seconds=total_duration,
             steps=steps,
+            flight_declaration_data=flight_declaration_data,
+            telemetry_data=telemetry_data,
         )
 
     except Exception as e: