From 4c595d2430a07c014f523927d5dec2beb3b16b4f Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Mon, 15 Apr 2024 23:09:57 +0200
Subject: [PATCH] Progress

---
 core/src/Ribasim.jl | 2 +-
 core/src/main.jl    | 7 +++----
 core/src/model.jl   | 9 +++++++--
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
index dc2b07ed4..835d68ef8 100644
--- a/core/src/Ribasim.jl
+++ b/core/src/Ribasim.jl
@@ -53,10 +53,10 @@ using OrdinaryDiffEq: OrdinaryDiffEq, OrdinaryDiffEqRosenbrockAdaptiveAlgorithm
 using PreallocationTools: DiffCache, get_tmp
 using SciMLBase:
     init,
-    check_error!,
     solve!,
     step!,
     SciMLBase,
+    ReturnCode,
     successful_retcode,
     CallbackSet,
     ODEFunction,
diff --git a/core/src/main.jl b/core/src/main.jl
index 78d4823ad..1e75e721f 100644
--- a/core/src/main.jl
+++ b/core/src/main.jl
@@ -54,8 +54,7 @@ function main(ARGS::Vector{String})::Cint
         config = Config(arg)
         mkpath(results_path(config, "."))
         open(results_path(config, "ribasim.log"), "w") do io
-            logger =
-                Ribasim.setup_logger(; verbosity = config.logging.verbosity, stream = io)
+            logger = setup_logger(; verbosity = config.logging.verbosity, stream = io)
             with_logger(logger) do
                 cli = (; ribasim_version = string(pkgversion(Ribasim)))
                 (; starttime, endtime) = config
@@ -63,13 +62,13 @@ function main(ARGS::Vector{String})::Cint
                     @warn "The Ribasim version in the TOML config file does not match the used Ribasim CLI version." config.ribasim_version cli.ribasim_version
                 end
                 @info "Starting a Ribasim simulation." cli.ribasim_version starttime endtime
-                model = Ribasim.run(config)
+                model = run(config)
                 if successful_retcode(model)
                     @info "The model finished successfully"
                     return 0
                 end
 
-                t = Ribasim.datetime_since(model.integrator.t, starttime)
+                t = datetime_since(model.integrator.t, starttime)
                 retcode = model.integrator.sol.retcode
                 @error "The model exited at model time $t with return code $retcode.\nSee https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/#retcodes"
                 return 1
diff --git a/core/src/model.jl b/core/src/model.jl
index 649038049..7d04ec562 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -206,7 +206,7 @@ function SciMLBase.step!(model::Model, dt::Float64)::Model
     if ntimes > 0 && round(ntimes) ≈ ntimes
         update_allocation!(integrator)
     end
-    step!(integrator, dt; stop_at_tdt = true)
+    step!(integrator, dt, true)
     return model
 end
 
@@ -222,9 +222,14 @@ function SciMLBase.solve!(model::Model)::Model
         for _ in TimeChoiceIterator(integrator, times)
             update_allocation!(integrator)
         end
+
+        # https://github.com/SciML/SciMLBase.jl/issues/669
+        if integrator.sol.retcode != ReturnCode.Default
+            return integrator.sol
+        end
+        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol, ReturnCode.Success)
     else
         solve!(integrator)
     end
-    check_error!(integrator)
     return model
 end