Support for persistent sessions using local and remote (LaunchIT) MariaDB by amarathe84 · Pull Request #168 · FLASK-LLNL/flask-copilot

amarathe84 · 2026-02-27T00:02:34Z

Description

This PR adds database-backed session persistence to the Flask Copilot:

Support for persistent sessions local and remote
Automatic checkpointing - Saves experiment state to MariaDB each time a new molecule is generated
Manual checkpointing - "Checkpoint" dropdown button allows on-demand saves and toggling auto-save on/off
Session restore on reconnect - Automatically restores the last session (molecules, graph layout, reasoning messages) when reopening the app
Project/experiment tracking - Each checkpoint stores project ID, experiment ID, and full experiment metadata
Complete state preservation - Saves tree nodes, edges, sidebar reasoning messages, metrics history, zoom/pan settings, and property optimization config
Duplicate prevention - Prevents duplicate molecules when restoring mid-computation sessions
Graceful resume handling - Detects interrupted computations and attempts to resume or gracefully resets

Minor code hardening:

Moved credentials to environment variables (no hard-coding)
-Add automatic connection retry in case of network connection failure or stalled session. Also added a timeout for the overall experiment
-Fix the SSL configuration (no more fake_flag_to_enable_tls used previously)
Add cleanup for sync_engine on failure - Currently only async_engine is disposed in main()
Unit testing (CLI) for state consistency
Code changes for integration with LC user management

Unit/integration testing:

Checkpoint Creation & Updates

Test checkpoint creation on node arrival
Test checkpoint updates incrementally as nodes arrive
Test checkpoint created when no prior state exists
Test checkpoint merges with existing partial state

Checkpoint Persistence

 5. Test checkpoint persists correct node count after disconnect 
 6. Test checkpoint persists correct edge count after disconnect 
 7. Test serverSessionId is correctly stored in checkpoint 
 8. Test isComputing flag is set to false in checkpoint after disconnect 
 9. Test checkpoint contains all required state fields (smiles, problemType, etc.)
  10. Test checkpoint timestamp is updated on each save

Checkpoint Restoration

 11. Test checkpoint restoration via /sessions/latest endpoint
 12. Test checkpoint restoration via /sessions/{id} endpoint 
 13. Test resumed computation continues from checkpoint

Some Corner Case Handling

 14. Test checkpoint handles malformed node data gracefully
  15. Test checkpoint with maximum depth computation (compute complete and restore)

tbennun · 2026-02-27T23:52:03Z

+    SSL_KW = {}
+
+# Create sync engine (SSL required by server; skip cert verification because tunnel is trusted)
+try:


using the global namespace to run this code is problematic. uvicorn uses fork and so sometimes these objects are not preserved well. The engines need to be created lazily in a function context

To address this, engine creation is now lazy and happens with get_async_engine() rather than at module import time. The module-level globals are only placeholders and are not constructed engine objects. Let me know if this looks good

engine is still assigned in the global context...

please also refactor this lengthy code into a function

These were due to incorrect force-push from another branch and conflict resolution which didn't go as planned. Both of these (local cache instead of global context and refactor) have been handled now.

tbennun · 2026-02-27T23:52:16Z

+    with sync_engine.connect() as _conn:
+        _conn.execute(text('SELECT 1'))
+
+    # Also create async engine for compatibility


why two engines?

The sync engine functionality was initially added for CLI/debugging at my end with a couple of standalone test scripts to check consistency of database, functionality flow etc.. I’ve removed the sync engine since it isn’t used (the tests scripts are not part of this PR). I have a separate branch where I'm keeping the functionality required for debugging.

I can still see it above. Is there a merge conflict?

Yes, I had messed it up with the last push. I've correctly merged and pushed from my local branch

tbennun · 2026-02-27T23:54:12Z

+    sidebar_state: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
+
+    # Experiment context
+    experiment_context: Mapped[Optional[str]] = mapped_column(Text, nullable=True)


why is this Text and not JSON?

FIxed this. experiment_context is now JSON based. Was originally plaintext for readability but I've switched to JSON

tbennun · 2026-02-27T23:57:04Z

+
+
+# Pydantic models
+class ExperimentData(BaseModel):


this repeats what we have in database/models, but looks better? Consolidate the two classes please

Yes makes sense. Now these are combined into schemas.py. Code in projects.py imports and uses the shared schema definitions instead of redefining mapping classes locally.

tbennun · 2026-03-13T16:27:53Z

+    with sync_engine.connect() as _conn:
+        _conn.execute(text("SELECT 1"))


only one process will execute this!

Good catch, silly mistake (based on my incorrect scope understanding earlier). New update to the code takes out eager validation and startup time DB side effect issue.

- Add db_backend/ package with SQLAlchemy models, schemas, and REST API routers for projects and sessions (CRUD + auto-save/restore) - Add useSessionPersistence hook for frontend auto-checkpoint logic - Update App.tsx with experiment persistence, restore-on-reload, and problemType fallback when DB row is not yet populated - Update project_sidebar.tsx with save/checkpoint UI controls - Update useProjectData.ts to integrate with the new DB backend - Add db_config.py for database connection configuration - Add .env.example with all required environment variables - Update mock_server.py with session save/restore, immediate persist on compute start, and background checkpoint support - Update .gitignore with Python, env, and build artifact patterns - Add sqlalchemy to requirements.txt

amarathe84 requested review from bvanessen and tbennun February 27, 2026 00:02

amarathe84 force-pushed the database-v3 branch from 801fda2 to d2a502e Compare March 5, 2026 23:58

tbennun requested changes Mar 6, 2026

View reviewed changes

amarathe84 force-pushed the database-v3 branch from 22b71e2 to 9360807 Compare March 10, 2026 20:14

amarathe84 requested a review from tbennun March 10, 2026 20:36

amarathe84 force-pushed the database-v3 branch from 9360807 to 8bdfd5d Compare March 12, 2026 23:32

tbennun reviewed Mar 13, 2026

View reviewed changes

amarathe84 force-pushed the database-v3 branch from 8bdfd5d to c666da0 Compare March 17, 2026 19:10

amarathe84 requested a review from tbennun March 17, 2026 20:02

amarathe84 force-pushed the database-v3 branch from cb1ee54 to 7241d86 Compare March 19, 2026 22:41

amarathe84 added 5 commits April 14, 2026 15:19

Ran pre-commit to minus linting/formatting

3d75785

Unlinked sessions.py; refactored to use websocket

5ca1faf

Addressed review comments

395a1ee

Removed global context assignment to the engine

1eb8e38

amarathe84 force-pushed the database-v3 branch from 7241d86 to 1eb8e38 Compare April 14, 2026 22:43

		with sync_engine.connect() as _conn:
		_conn.execute(text("SELECT 1"))

Uh oh!

Conversation

amarathe84 commented Feb 27, 2026

Description

Minor code hardening:

Unit/integration testing:

Checkpoint Creation & Updates

Checkpoint Persistence

Checkpoint Restoration

Some Corner Case Handling

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tbennun Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tbennun Mar 13, 2026 •

edited

Loading