Add per-train lifecycle hook overrides on ServiceTrain

src/Trax.Effect/Extensions/BroadcasterExtensions.cs

@@ -2,7 +2,6 @@
 using Trax.Effect.Configuration.BroadcasterBuilder;
 using Trax.Effect.Configuration.TraxEffectBuilder;
 using Trax.Effect.Services.TrainEventBroadcaster;
-using Trax.Effect.Services.TrainLifecycleHookFactory;
 
 namespace Trax.Effect.Extensions;
 
@@ -31,14 +30,7 @@ Action<BroadcasterBuilder> configure
         var broadcasterBuilder = new BroadcasterBuilder(builder);
         configure(broadcasterBuilder);
 
-        builder.ServiceCollection.AddTransient<BroadcastLifecycleHook>();
-        builder
-            .ServiceCollection.AddSingleton<BroadcastLifecycleHookFactory>()
-            .AddSingleton<ITrainLifecycleHookFactory>(sp =>
-                sp.GetRequiredService<BroadcastLifecycleHookFactory>()
-            );
-
-        builder.EffectRegistry?.Register(typeof(BroadcastLifecycleHookFactory), toggleable: false);
+        builder.AddLifecycleHook<BroadcastLifecycleHook>(toggleable: false);
 
         builder.ServiceCollection.AddHostedService<TrainEventReceiverService>();
 

src/Trax.Effect/Extensions/ServiceExtensions.cs

@@ -10,6 +10,7 @@
 using Trax.Effect.Services.JunctionEffectProviderFactory;
 using Trax.Effect.Services.JunctionEffectRunner;
 using Trax.Effect.Services.LifecycleHookRunner;
+using Trax.Effect.Services.TrainLifecycleHook;
 using Trax.Effect.Services.TrainLifecycleHookFactory;
 
 namespace Trax.Effect.Extensions;
@@ -392,55 +393,87 @@ public static TraxEffectBuilder AddJunctionEffect<TJunctionEffectProviderFactory
     #region LifecycleHook
 
     /// <summary>
-    /// Registers a train lifecycle hook with both its interface and implementation type,
-    /// using an existing factory instance. Lifecycle hooks run at train start/completion/failure boundaries.
+    /// Registers a train lifecycle hook or hook factory.
+    /// <para>
+    /// If <typeparamref name="T"/> implements <see cref="ITrainLifecycleHook"/>, a factory is
+    /// created internally — no need to write a separate factory class. The hook is resolved from DI
+    /// on each train execution, so constructor-injected dependencies work.
+    /// </para>
+    /// <para>
+    /// If <typeparamref name="T"/> implements <see cref="ITrainLifecycleHookFactory"/>, the factory
+    /// is registered directly (advanced usage).
+    /// </para>
     /// </summary>
-    /// <typeparam name="TILifecycleHookFactory">The lifecycle hook factory interface.</typeparam>
-    /// <typeparam name="TLifecycleHookFactory">The concrete lifecycle hook factory type.</typeparam>
+    /// <typeparam name="T">
+    /// Either a concrete <see cref="ITrainLifecycleHook"/> type or an <see cref="ITrainLifecycleHookFactory"/> type.
+    /// </typeparam>
     /// <param name="builder">The effect builder.</param>
-    /// <param name="factory">The factory instance to register.</param>
     /// <param name="toggleable">Whether this hook can be toggled on/off at runtime. Defaults to <c>true</c>.</param>
     /// <returns>The effect builder for chaining.</returns>
-    public static TraxEffectBuilder AddLifecycleHook<TILifecycleHookFactory, TLifecycleHookFactory>(
+    public static TraxEffectBuilder AddLifecycleHook<T>(
         this TraxEffectBuilder builder,
-        TLifecycleHookFactory factory,
         bool toggleable = true
     )
-        where TILifecycleHookFactory : class, ITrainLifecycleHookFactory
-        where TLifecycleHookFactory : class, TILifecycleHookFactory
+        where T : class
     {
-        builder
-            .ServiceCollection.AddSingleton<TLifecycleHookFactory>(factory)
-            .AddSingleton<ITrainLifecycleHookFactory>(sp =>
-                sp.GetRequiredService<TLifecycleHookFactory>()
-            )
-            .AddSingleton<TILifecycleHookFactory>(sp =>
-                sp.GetRequiredService<TLifecycleHookFactory>()
+        if (typeof(ITrainLifecycleHookFactory).IsAssignableFrom(typeof(T)))
+        {
+            builder
+                .ServiceCollection.AddSingleton(typeof(T))
+                .AddSingleton<ITrainLifecycleHookFactory>(sp =>
+                    (ITrainLifecycleHookFactory)sp.GetRequiredService(typeof(T))
+                );
+
+            builder.EffectRegistry?.Register(typeof(T), toggleable: toggleable);
+        }
+        else if (typeof(ITrainLifecycleHook).IsAssignableFrom(typeof(T)))
+        {
+            builder.ServiceCollection.AddTransient(typeof(T));
+
+            var factoryType = typeof(LifecycleHookFactory<>).MakeGenericType(typeof(T));
+            builder.ServiceCollection.AddSingleton(factoryType);
+            builder.ServiceCollection.AddSingleton<ITrainLifecycleHookFactory>(sp =>
+                (ITrainLifecycleHookFactory)sp.GetRequiredService(factoryType)
             );
 
-        builder.EffectRegistry?.Register(typeof(TLifecycleHookFactory), toggleable: toggleable);
+            builder.EffectRegistry?.Register(factoryType, toggleable: toggleable);
+        }
+        else
+        {
+            throw new InvalidOperationException(
+                $"AddLifecycleHook<{typeof(T).Name}>() requires a type that implements "
+                    + $"ITrainLifecycleHook or ITrainLifecycleHookFactory."
+            );
+        }
 
         return builder;
     }
 
     /// <summary>
-    /// Registers a train lifecycle hook resolved from DI.
-    /// Lifecycle hooks run at train start/completion/failure boundaries.
+    /// Registers a train lifecycle hook with both its interface and implementation type,
+    /// using an existing factory instance. Lifecycle hooks run at train start/completion/failure boundaries.
     /// </summary>
+    /// <typeparam name="TILifecycleHookFactory">The lifecycle hook factory interface.</typeparam>
     /// <typeparam name="TLifecycleHookFactory">The concrete lifecycle hook factory type.</typeparam>
     /// <param name="builder">The effect builder.</param>
+    /// <param name="factory">The factory instance to register.</param>
     /// <param name="toggleable">Whether this hook can be toggled on/off at runtime. Defaults to <c>true</c>.</param>
     /// <returns>The effect builder for chaining.</returns>
-    public static TraxEffectBuilder AddLifecycleHook<TLifecycleHookFactory>(
+    public static TraxEffectBuilder AddLifecycleHook<TILifecycleHookFactory, TLifecycleHookFactory>(
         this TraxEffectBuilder builder,
+        TLifecycleHookFactory factory,
         bool toggleable = true
     )
-        where TLifecycleHookFactory : class, ITrainLifecycleHookFactory
+        where TILifecycleHookFactory : class, ITrainLifecycleHookFactory
+        where TLifecycleHookFactory : class, TILifecycleHookFactory
     {
         builder
-            .ServiceCollection.AddSingleton<TLifecycleHookFactory>()
+            .ServiceCollection.AddSingleton<TLifecycleHookFactory>(factory)
             .AddSingleton<ITrainLifecycleHookFactory>(sp =>
                 sp.GetRequiredService<TLifecycleHookFactory>()
+            )
+            .AddSingleton<TILifecycleHookFactory>(sp =>
+                sp.GetRequiredService<TLifecycleHookFactory>()
             );
 
         builder.EffectRegistry?.Register(typeof(TLifecycleHookFactory), toggleable: toggleable);

src/Trax.Effect/Services/ServiceTrain/ServiceTrain.cs

@@ -85,6 +85,37 @@ public abstract class ServiceTrain<TIn, TOut> : Train<TIn, TOut>, IServiceTrain<
         ?? GetType().FullName
         ?? throw new TrainException($"Could not find FullName for ({GetType().Name})");
 
+    /// <summary>
+    /// Called after the train's metadata is initialized and persisted, before RunInternal executes.
+    /// Override to add per-train startup logic. Exceptions are caught and logged — they will not
+    /// prevent the train from running.
+    /// </summary>
+    protected virtual Task OnStarted(Metadata metadata, CancellationToken ct) => Task.CompletedTask;
+
+    /// <summary>
+    /// Called after a successful run, after output is persisted and global hooks have fired.
+    /// Override to add per-train completion logic (e.g., notifications, cache invalidation).
+    /// Exceptions are caught and logged — they will not cause the train to report failure.
+    /// </summary>
+    protected virtual Task OnCompleted(Metadata metadata, CancellationToken ct) =>
+        Task.CompletedTask;
+
+    /// <summary>
+    /// Called after a failed run (non-cancellation exception), after failure state is persisted
+    /// and global hooks have fired. Override to add per-train failure handling (e.g., alerting).
+    /// Exceptions are caught and logged — they will not mask the original failure.
+    /// </summary>
+    protected virtual Task OnFailed(Metadata metadata, Exception exception, CancellationToken ct) =>
+        Task.CompletedTask;
+
+    /// <summary>
+    /// Called after cancellation (OperationCanceledException), after cancellation state is persisted
+    /// and global hooks have fired. Override to add per-train cancellation handling.
+    /// Exceptions are caught and logged.
+    /// </summary>
+    protected virtual Task OnCancelled(Metadata metadata, CancellationToken ct) =>
+        Task.CompletedTask;
+
     /// <summary>
     /// Overrides the base Train Run method to add database tracking and logging capabilities.
     /// </summary>
@@ -108,6 +139,19 @@ public override async Task<TOut> Run(TIn input, CancellationToken cancellationTo
 
         await LifecycleHookRunner.OnStarted(Metadata, CancellationToken);
 
+        try
+        {
+            await OnStarted(Metadata, CancellationToken);
+        }
+        catch (Exception hookEx)
+        {
+            Logger?.LogError(
+                hookEx,
+                "Train-level OnStarted hook threw for train ({TrainName}).",
+                TrainName
+            );
+        }
+
         try
         {
             Logger?.LogTrace("Running Train: ({TrainName})", TrainName);
@@ -126,10 +170,40 @@ public override async Task<TOut> Run(TIn input, CancellationToken cancellationTo
                 await EffectRunner.SaveChanges(CancellationToken);
 
                 if (exception is OperationCanceledException)
+                {
                     await LifecycleHookRunner.OnCancelled(Metadata, CancellationToken);
+
+                    try
+                    {
+                        await OnCancelled(Metadata, CancellationToken);
+                    }
+                    catch (Exception hookEx)
+                    {
+                        Logger?.LogError(
+                            hookEx,
+                            "Train-level OnCancelled hook threw for train ({TrainName}).",
+                            TrainName
+                        );
+                    }
+                }
                 else
+                {
                     await LifecycleHookRunner.OnFailed(Metadata, exception, CancellationToken);
 
+                    try
+                    {
+                        await OnFailed(Metadata, exception, CancellationToken);
+                    }
+                    catch (Exception hookEx)
+                    {
+                        Logger?.LogError(
+                            hookEx,
+                            "Train-level OnFailed hook threw for train ({TrainName}).",
+                            TrainName
+                        );
+                    }
+                }
+
                 exception.Rethrow();
             }
 
@@ -143,6 +217,19 @@ public override async Task<TOut> Run(TIn input, CancellationToken cancellationTo
 
             await LifecycleHookRunner.OnCompleted(Metadata, CancellationToken);
 
+            try
+            {
+                await OnCompleted(Metadata, CancellationToken);
+            }
+            catch (Exception hookEx)
+            {
+                Logger?.LogError(
+                    hookEx,
+                    "Train-level OnCompleted hook threw for train ({TrainName}).",
+                    TrainName
+                );
+            }
+
             return output;
         }
         catch (Exception e)
@@ -157,10 +244,40 @@ public override async Task<TOut> Run(TIn input, CancellationToken cancellationTo
             await EffectRunner.SaveChanges(CancellationToken);
 
             if (e is OperationCanceledException)
+            {
                 await LifecycleHookRunner.OnCancelled(Metadata, CancellationToken);
+
+                try
+                {
+                    await OnCancelled(Metadata, CancellationToken);
+                }
+                catch (Exception hookEx)
+                {
+                    Logger?.LogError(
+                        hookEx,
+                        "Train-level OnCancelled hook threw for train ({TrainName}).",
+                        TrainName
+                    );
+                }
+            }
             else
+            {
                 await LifecycleHookRunner.OnFailed(Metadata, e, CancellationToken);
 
+                try
+                {
+                    await OnFailed(Metadata, e, CancellationToken);
+                }
+                catch (Exception hookEx)
+                {
+                    Logger?.LogError(
+                        hookEx,
+                        "Train-level OnFailed hook threw for train ({TrainName}).",
+                        TrainName
+                    );
+                }
+            }
+
             throw;
         }
     }

src/Trax.Effect/Services/TrainLifecycleHookFactory/LifecycleHookFactory.cs

@@ -0,0 +1,17 @@
+using Microsoft.Extensions.DependencyInjection;
+using Trax.Effect.Services.TrainLifecycleHook;
+
+namespace Trax.Effect.Services.TrainLifecycleHookFactory;
+
+/// <summary>
+/// Generic factory that creates lifecycle hook instances via DI.
+/// Used internally by the <c>AddLifecycleHook&lt;THook&gt;()</c> overload
+/// so that users don't need to write their own factory classes.
+/// </summary>
+public class LifecycleHookFactory<THook>(IServiceProvider serviceProvider)
+    : ITrainLifecycleHookFactory
+    where THook : class, ITrainLifecycleHook
+{
+    public ITrainLifecycleHook Create() =>
+        ActivatorUtilities.CreateInstance<THook>(serviceProvider);
+}

tests/Trax.Effect.Tests.Broadcaster/UnitTests/BroadcasterRegistrationTests.cs

@@ -115,7 +115,7 @@ public void UseBroadcaster_RegistersHookAsNonToggleable()
         );
 
         // The broadcast hook should be registered and enabled (non-toggleable)
-        registry.IsEnabled(typeof(BroadcastLifecycleHookFactory)).Should().BeTrue();
+        registry.IsEnabled(typeof(LifecycleHookFactory<BroadcastLifecycleHook>)).Should().BeTrue();
     }
 
     [Test]