Skip to content

Feature: Move Tome schedule configuration to new ScheduledTask ent #2072

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hulto merged 7 commits into from
Mar 14, 2026
Merged

Feature: Move Tome schedule configuration to new ScheduledTask ent #2072

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
81e4f26
feat: create ScheduledTask ent and move scheduling off of Tomes
google-labs-jules[bot] Mar 14, 2026
b638e07
feat: create ScheduledTask ent and remove scheduling from Tomes
google-labs-jules[bot] Mar 14, 2026
64df866
refactor: handle scheduled hosts via sql query instead of loop
google-labs-jules[bot] Mar 14, 2026
69ae2c5
generate off main
hulto Mar 14, 2026
e94250f
Set tome automation to use schtasks not tome
hulto Mar 14, 2026
1b3279b
Add disabling and quust traverse
hulto Mar 14, 2026
d1d99c3
Merge branch 'main' into jules-scheduled-task-ent-12776578495355315753
hulto Mar 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
345 changes: 100 additions & 245 deletions tavern/internal/builder/builderpb/builder.pb.go
View file
Open in desktop

Large diffs are not rendered by default.

180 changes: 63 additions & 117 deletions tavern/internal/builder/builderpb/builder_grpc.pb.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

37 changes: 19 additions & 18 deletions tavern/internal/c2/api_claim_tasks.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import (
"realm.pub/tavern/internal/ent/shelltask"
"realm.pub/tavern/internal/ent/tag"
"realm.pub/tavern/internal/ent/task"
"realm.pub/tavern/internal/ent/tome"
"realm.pub/tavern/internal/ent/scheduledtask"
"realm.pub/tavern/internal/namegen"
"realm.pub/tavern/internal/redirectors"
)
Expand Down Expand Up @@ -49,16 +49,17 @@ func init() {

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(
tome.RunOnNewBeaconCallback(true),
tome.RunOnFirstHostCallback(true),
tome.RunOnScheduleNEQ(""),
candidateTasks, err := srv.graph.ScheduledTask.Query().
Where(scheduledtask.Or(
scheduledtask.RunOnNewBeaconCallback(true),
scheduledtask.RunOnFirstHostCallback(true),
scheduledtask.RunOnScheduleNEQ(""),
)).
WithTome().
All(ctx)

if err != nil {
slog.ErrorContext(ctx, "failed to query candidate tomes for automation", "err", err)
slog.ErrorContext(ctx, "failed to query candidate scheduled tasks for automation", "err", err)
metricTomeAutomationErrors.Inc()
return
}
Expand All @@ -67,7 +68,7 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
parser := cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow)
cutoff := now.Add(interval)

for _, t := range candidateTomes {
for _, t := range candidateTasks {
shouldRun := false

// Check RunOnNewBeaconCallback
Expand Down Expand Up @@ -108,7 +109,7 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
// Check scheduled_hosts constraint
hostCount, err := t.QueryScheduledHosts().Count(ctx)
if err != nil {
slog.ErrorContext(ctx, "failed to count scheduled hosts for automation", "err", err, "tome_id", t.ID)
slog.ErrorContext(ctx, "failed to count scheduled hosts for automation", "err", err, "scheduled_task_id", t.ID)
metricTomeAutomationErrors.Inc()
continue
}
Expand All @@ -119,7 +120,7 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
Where(host.ID(hostID)).
Exist(ctx)
if err != nil {
slog.ErrorContext(ctx, "failed to check host existence for automation", "err", err, "tome_id", t.ID)
slog.ErrorContext(ctx, "failed to check host existence for automation", "err", err, "scheduled_task_id", t.ID)
metricTomeAutomationErrors.Inc()
continue
}
Expand All @@ -134,21 +135,21 @@ func (srv *Server) handleTomeAutomation(ctx context.Context, beaconID int, hostI
}
}

if shouldRun {
selectedTomes[t.ID] = t
if shouldRun && t.Edges.Tome != nil {
selectedTomes[t.Edges.Tome.ID] = t.Edges.Tome
}
}

// Create Quest and Task for each selected Tome
for _, t := range selectedTomes {
for _, tome := range selectedTomes {
q, err := srv.graph.Quest.Create().
SetName(fmt.Sprintf("Automated: %s", t.Name)).
SetTome(t).
SetParamDefsAtCreation(t.ParamDefs).
SetEldritchAtCreation(t.Eldritch).
SetName(fmt.Sprintf("Automated: %s", tome.Name)).
SetTome(tome).
SetParamDefsAtCreation(tome.ParamDefs).
SetEldritchAtCreation(tome.Eldritch).
Save(ctx)
if err != nil {
slog.ErrorContext(ctx, "failed to create automated quest", "err", err, "tome_id", t.ID)
slog.ErrorContext(ctx, "failed to create automated quest", "err", err, "tome_id", tome.ID)
metricTomeAutomationErrors.Inc()
continue
}
Expand Down
Loading
Loading