11package com .dbaagent .controller ;
22
3+ import com .dbaagent .model .SavedDashboard ;
34import com .dbaagent .service .DashboardAgentService ;
5+ import com .dbaagent .service .SavedDashboardService ;
46import com .dbaagent .service .security .AccessControlService ;
57import lombok .RequiredArgsConstructor ;
68import lombok .extern .slf4j .Slf4j ;
9+ import org .springframework .dao .OptimisticLockingFailureException ;
710import org .springframework .http .HttpStatus ;
811import org .springframework .http .MediaType ;
912import org .springframework .http .ResponseEntity ;
1821
1922import java .io .IOException ;
2023import java .util .Map ;
24+ import java .util .UUID ;
2125import java .util .concurrent .Executors ;
2226import java .util .concurrent .ScheduledExecutorService ;
2327import java .util .concurrent .TimeUnit ;
@@ -44,6 +48,7 @@ public class DashboardGenerationController {
4448
4549 private final DashboardAgentService dashboardAgentService ;
4650 private final AccessControlService accessControlService ;
51+ private final SavedDashboardService savedDashboardService ;
4752
4853 @ PostMapping ("/generate" )
4954 public ResponseEntity <?> generate (@ RequestBody GenerateRequest request ) {
@@ -68,13 +73,47 @@ public ResponseEntity<?> generate(@RequestBody GenerateRequest request) {
6873 @ PostMapping (value = "/generate/stream" , produces = MediaType .TEXT_EVENT_STREAM_VALUE )
6974 public SseEmitter generateStream (@ RequestBody GenerateRequest request ) {
7075 requireValid (request );
71- accessControlService .assertCanReadConnectionContent (request .connectionId ());
76+ // This path creates/updates a SavedDashboard row (beginGenerationTurn etc.)
77+ // on every call, not just reads — a VIEWER (read-only) must not be able to
78+ // mint or mutate drafts via chat.
79+ accessControlService .assertCanManageConnectionContent (request .connectionId ());
80+ SseEmitter emitter = new SseEmitter (600_000L );
81+
82+ // Resolve (or create) the target dashboard and record the user's message
83+ // SYNCHRONOUSLY, before any slow agent work starts — this is what lets a
84+ // reload mid-generation see "still working" (generationStatus=RUNNING)
85+ // instead of nothing at all, even for a brand-new, never-saved dashboard.
86+ // See SavedDashboardService's "Server-owned chat-turn persistence" section.
87+ final SavedDashboard dashboard ;
88+ try {
89+ dashboard = savedDashboardService .beginGenerationTurn (
90+ request .dashboardId (), request .connectionId (), request .prompt ());
91+ } catch (IllegalArgumentException | IllegalStateException e ) {
92+ sendErrorAndComplete (emitter , e .getMessage ());
93+ return emitter ;
94+ } catch (OptimisticLockingFailureException e ) {
95+ // Lost the race to another concurrent submit on the same dashboard —
96+ // same user-facing shape as the "already running" case above.
97+ sendErrorAndComplete (emitter , "A generation is already running for this dashboard." );
98+ return emitter ;
99+ }
100+ // The frontend needs this id right away (not just at the end) so a
101+ // brand-new dashboard is addressable — e.g. by a reload — well before
102+ // the potentially multi-minute build finishes.
103+ try {
104+ emitter .send (SseEmitter .event ().name ("created" )
105+ .data (Map .of ("dashboardId" , dashboard .getId ().toString ())));
106+ } catch (IOException ignore ) {
107+ // Client already gone before we even started streaming — fine, the
108+ // turn is already durably recorded; the work below still runs and
109+ // persists its result regardless of this connection.
110+ }
111+
72112 // Coding a whole dashboard (ground + verify every query + write the HTML) can
73113 // run for minutes. Give it real headroom (10 min) and keep the stream alive
74114 // with a heartbeat — otherwise it emits only 3 step events and the long idle
75115 // gap gets cut by nginx/emitter timeouts before `done`, surfacing to the user
76116 // as "Generation ended unexpectedly".
77- SseEmitter emitter = new SseEmitter (600_000L );
78117 ScheduledExecutorService heartbeat = Executors .newSingleThreadScheduledExecutor (r -> {
79118 Thread t = new Thread (r , "dashboard-generate-hb" );
80119 t .setDaemon (true );
@@ -106,21 +145,47 @@ public SseEmitter generateStream(@RequestBody GenerateRequest request) {
106145 // `done` event with a real artifact — the FE's done handler always
107146 // appends "Done — built…" and auto-saves. A dedicated `chat` event
108147 // keeps that path from swallowing out-of-context messages.
109- if (Boolean .TRUE .equals (config .get ("chat" ))) {
110- emitter .send (SseEmitter .event ().name ("chat" )
111- .data (Map .of (
112- "success" , true ,
113- "reply" , String .valueOf (config .getOrDefault ("reply" , "" )),
114- "dashboardConfig" , config )));
115- } else {
116- emitter .send (SseEmitter .event ().name ("done" )
117- .data (Map .of ("success" , true , "dashboardConfig" , config )));
148+ boolean chatOnly = Boolean .TRUE .equals (config .get ("chat" ));
149+ // Persist BEFORE attempting to notify the client — a client that's
150+ // gone by now must never turn an already-successful result into a
151+ // recorded failure (see the catch block below, which only ever
152+ // handles a real dashboardAgentService.generate() failure, not a
153+ // dead SSE connection at delivery time).
154+ try {
155+ if (chatOnly ) {
156+ savedDashboardService .appendAgentReply (
157+ dashboard .getId (), String .valueOf (config .getOrDefault ("reply" , "" )));
158+ } else {
159+ savedDashboardService .completeBuildTurn (dashboard .getId (), config );
160+ }
161+ } catch (Exception persistErr ) {
162+ log .error ("Failed to persist completed dashboard turn {}" , dashboard .getId (), persistErr );
163+ }
164+ try {
165+ if (chatOnly ) {
166+ emitter .send (SseEmitter .event ().name ("chat" )
167+ .data (Map .of (
168+ "success" , true ,
169+ "reply" , String .valueOf (config .getOrDefault ("reply" , "" )),
170+ "dashboardConfig" , config )));
171+ } else {
172+ emitter .send (SseEmitter .event ().name ("done" )
173+ .data (Map .of ("success" , true , "dashboardConfig" , config )));
174+ }
175+ } catch (IOException ignore ) {
176+ // Client gone by the time the result was ready — already
177+ // persisted above, so this is a no-op, not a failure.
118178 }
119179 emitter .complete ();
120180 } catch (ClientGoneException gone ) {
121181 emitter .complete ();
122182 } catch (Exception e ) {
123183 log .warn ("Streamed dashboard generation failed: {}" , e .getMessage ());
184+ try {
185+ savedDashboardService .appendErrorReply (dashboard .getId (), safe (e ));
186+ } catch (Exception persistErr ) {
187+ log .error ("Failed to persist dashboard generation error {}" , dashboard .getId (), persistErr );
188+ }
124189 try {
125190 emitter .send (SseEmitter .event ().name ("error" )
126191 .data (Map .of ("success" , false , "error" , safe (e ))));
@@ -134,6 +199,14 @@ public SseEmitter generateStream(@RequestBody GenerateRequest request) {
134199 return emitter ;
135200 }
136201
202+ private static void sendErrorAndComplete (SseEmitter emitter , String message ) {
203+ try {
204+ emitter .send (SseEmitter .event ().name ("error" )
205+ .data (Map .of ("success" , false , "error" , message == null ? "Request failed" : message )));
206+ } catch (IOException ignore ) { }
207+ emitter .complete ();
208+ }
209+
137210 private static void requireValid (GenerateRequest request ) {
138211 if (request == null || request .connectionId () == null
139212 || request .prompt () == null || request .prompt ().isBlank ()) {
@@ -149,5 +222,8 @@ private static final class ClientGoneException extends RuntimeException {
149222 ClientGoneException (Throwable cause ) { super (cause ); }
150223 }
151224
152- public record GenerateRequest (String connectionId , String prompt , Object currentConfig ) { }
225+ // dashboardId is optional — omitting it (a brand-new, never-saved dashboard)
226+ // always creates a new SavedDashboard row, matching the pre-existing default
227+ // behavior for new dashboards.
228+ public record GenerateRequest (String connectionId , String prompt , Object currentConfig , UUID dashboardId ) { }
153229}
0 commit comments