Improve tome automation scheduling window (#1527)

* Improve tome automation scheduling window

Modified tome automation logic to queue tasks if their schedule falls within the upcoming beacon check-in interval window (now to now + interval), rather than just checking the current minute. This prevents beacons with long sleep intervals from missing scheduled tasks.

Added `TestHandleTomeAutomation_IntervalWindow` to verify the fix and prevent regressions. Also fixed a missing sqlite3 driver import in `tome_automation_test.go`.

* Refactor tome automation to use time.Duration

Updated `handleTomeAutomation` signature to accept `time.Duration` for interval instead of `int`, addressing PR feedback. Updated call sites and tests to pass duration.

* Refactor tome automation for lookahead and range support

Updated `handleTomeAutomation` to support lookahead scheduling for standard cron expressions (queueing tasks scheduled within the upcoming beacon interval) while enforcing strict current-time matching for range-based schedules (e.g., `* 6-12 * * *`).

Refactored method to accept `time.Duration` for interval. Added comprehensive tests for interval window logic and cron range behavior. Fixed missing sqlite3 driver in tests.

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: google-labs-jules[bot]

tavern/internal/c2/api_claim_tasks.go

@@ -44,7 +44,7 @@ func init() {
 	prometheus.MustRegister(metricTomeAutomationErrors)
 }
 
-func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostID int, isNewBeacon bool, isNewHost bool, now time.Time) {
+func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostID int, isNewBeacon bool, isNewHost bool, now time.Time, interval time.Duration) {
 	// Tome Automation Logic
 	candidateTomes, err := srv.graph.Tome.Query().
 		Where(tome.Or(
@@ -62,7 +62,7 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
 
 	selectedTomes := make(map[int]*ent.Tome)
 	parser := cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow)
-	currentMinute := now.Truncate(time.Minute)
+	cutoff := now.Add(interval)
 
 	for _, t := range candidateTomes {
 		shouldRun := false
@@ -81,10 +81,27 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
 		if !shouldRun && t.RunOnSchedule != "" {
 			sched, err := parser.Parse(t.RunOnSchedule)
 			if err == nil {
-				// Check if schedule matches current time
-				// Next(now-1sec) == now?
-				next := sched.Next(currentMinute.Add(-1 * time.Second))
-				if next.Equal(currentMinute) {
+				isMatch := false
+				// If schedule contains a range (hyphen), checking for strict current match
+				// without factoring in callback interval.
+				if strings.Contains(t.RunOnSchedule, "-") {
+					currentMinute := now.Truncate(time.Minute)
+					// Verify if 'currentMinute' is a valid trigger time.
+					// We check if the next trigger after (currentMinute - 1s) is exactly currentMinute.
+					next := sched.Next(currentMinute.Add(-1 * time.Second))
+					if next.Equal(currentMinute) {
+						isMatch = true
+					}
+				} else {
+					// Check if any point between now and the next expected check-in matches the run schedule
+					// Next(now-1sec) <= now + interval?
+					next := sched.Next(now.Add(-1 * time.Second))
+					if !next.After(cutoff) {
+						isMatch = true
+					}
+				}
+
+				if isMatch {
 					// Check scheduled_hosts constraint
 					hostCount, err := t.QueryScheduledHosts().Count(ctx)
 					if err != nil {
@@ -294,7 +311,7 @@ func (srv *Server) ClaimTasks(ctx context.Context, req *c2pb.ClaimTasksRequest)
 	}
 
 	// Run Tome Automation (non-blocking, best effort)
-	srv.handleTomeAutomation(ctx, beaconID, hostID, isNewBeacon, isNewHost, now)
+	srv.handleTomeAutomation(ctx, beaconID, hostID, isNewBeacon, isNewHost, now, time.Duration(req.GetBeacon().GetActiveTransport().GetInterval())*time.Second)
 
 	// Load Tasks
 	tasks, err := srv.graph.Task.Query().

tavern/internal/c2/tome_automation_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	_ "github.com/mattn/go-sqlite3"
 	"realm.pub/tavern/internal/c2/c2pb"
 	"realm.pub/tavern/internal/ent"
 	"realm.pub/tavern/internal/ent/enttest"
@@ -138,7 +139,9 @@ func TestHandleTomeAutomation(t *testing.T) {
 			client.Task.Delete().ExecX(ctx)
 			client.Quest.Delete().ExecX(ctx)
 
-			srv.handleTomeAutomation(ctx, b.ID, h.ID, tt.isNewBeacon, tt.isNewHost, now)
+			// Use 0 interval to match "current minute" logic (cutoff = now)
+			// effectively checking sched.Next(now-1s) <= now
+			srv.handleTomeAutomation(ctx, b.ID, h.ID, tt.isNewBeacon, tt.isNewHost, now, 0)
 
 			// Verify Tasks
 			tasks := client.Task.Query().WithQuest(func(q *ent.QuestQuery) {
@@ -185,9 +188,98 @@ func TestHandleTomeAutomation_Deduplication(t *testing.T) {
 		SaveX(ctx)
 
 	// Trigger all conditions
-	srv.handleTomeAutomation(ctx, b.ID, h.ID, true, true, now)
+	srv.handleTomeAutomation(ctx, b.ID, h.ID, true, true, now, 0)
 
 	// Should only have 1 task
 	count := client.Task.Query().CountX(ctx)
 	assert.Equal(t, 1, count, "Should only create one task despite multiple triggers matching")
 }
+
+func TestHandleTomeAutomation_IntervalWindow(t *testing.T) {
+	ctx := context.Background()
+	client := enttest.Open(t, "sqlite3", "file:ent?mode=memory&cache=shared&_fk=1")
+	defer client.Close()
+
+	srv := &Server{graph: client}
+	// Date: 3:00:56 PM
+	now := time.Date(2023, 10, 27, 15, 0, 56, 0, time.UTC)
+
+	h := client.Host.Create().
+		SetIdentifier("test-host").
+		SetPlatform(c2pb.Host_PLATFORM_LINUX).
+		SaveX(ctx)
+	b := client.Beacon.Create().
+		SetIdentifier("test-beacon").
+		SetHost(h).
+		SetTransport(c2pb.ActiveTransport_TRANSPORT_HTTP1).
+		SaveX(ctx)
+
+	// Schedule: 3:01:00 PM -> "1 15 * * *"
+	client.Tome.Create().
+		SetName("Scheduled Tome").
+		SetDescription("Test").
+		SetAuthor("Test Author").
+		SetEldritch("print('schedule')").
+		SetRunOnSchedule("1 15 * * *").
+		SaveX(ctx)
+
+	// 1. First Check-in at 3:00:56 PM. Interval 120s.
+	// Window: [3:00:56, 3:02:56].
+	// Schedule 3:01:00 is in window.
+	srv.handleTomeAutomation(ctx, b.ID, h.ID, false, false, now, 120*time.Second)
+
+	count := client.Task.Query().CountX(ctx)
+	assert.Equal(t, 1, count, "Tome should be queued at 3:00:56 for 3:01:00 schedule")
+
+	// Clear tasks
+	client.Task.Delete().ExecX(ctx)
+	client.Quest.Delete().ExecX(ctx)
+
+	// 2. Second Check-in at 3:02:56 PM. Interval 120s.
+	// Next checkin: 3:04:56.
+	// Schedule 3:01:00 (tomorrow) is NOT in window.
+	nextCheckin := now.Add(2 * time.Minute) // 3:02:56
+	srv.handleTomeAutomation(ctx, b.ID, h.ID, false, false, nextCheckin, 120*time.Second)
+
+	count = client.Task.Query().CountX(ctx)
+	assert.Equal(t, 0, count, "Tome should NOT be queued at 3:02:56")
+}
+
+func TestHandleTomeAutomation_CronRange(t *testing.T) {
+	ctx := context.Background()
+	client := enttest.Open(t, "sqlite3", "file:ent?mode=memory&cache=shared&_fk=1")
+	defer client.Close()
+
+	srv := &Server{graph: client}
+	// 06:05:00
+	now := time.Date(2023, 10, 27, 6, 5, 0, 0, time.UTC)
+
+	h := client.Host.Create().SetIdentifier("h").SetPlatform(c2pb.Host_PLATFORM_LINUX).SaveX(ctx)
+	b := client.Beacon.Create().SetIdentifier("b").SetHost(h).SetTransport(c2pb.ActiveTransport_TRANSPORT_HTTP1).SaveX(ctx)
+
+	// Range Schedule: "* 6-12 * * *" (Every minute of hours 6-12)
+	client.Tome.Create().
+		SetName("Range Tome").
+		SetDescription("Test").
+		SetAuthor("Test").
+		SetEldritch("print('range')").
+		SetRunOnSchedule("* 6-12 * * *").
+		SaveX(ctx)
+
+	// 1. Current time 06:05:00. Matches range. Should run.
+	// Interval irrelevant (using 1h to prove lookahead is ignored if range logic applies)
+	srv.handleTomeAutomation(ctx, b.ID, h.ID, false, false, now, 1*time.Hour)
+	count := client.Task.Query().CountX(ctx)
+	assert.Equal(t, 1, count, "Range tome should run at 06:05:00")
+
+	client.Task.Delete().ExecX(ctx)
+	client.Quest.Delete().ExecX(ctx)
+
+	// 2. Current time 05:55:00. Does NOT match range.
+	// Lookahead (1h) would see 06:00:00.
+	// But Range logic should IGNORE lookahead.
+	early := now.Add(-10 * time.Minute) // 05:55:00
+	srv.handleTomeAutomation(ctx, b.ID, h.ID, false, false, early, 1*time.Hour)
+	count = client.Task.Query().CountX(ctx)
+	assert.Equal(t, 0, count, "Range tome should NOT run at 05:55:00 even with lookahead")
+}