diff --git a/cmd/gateway_cron.go b/cmd/gateway_cron.go index 42554d00..2075f1f6 100644 --- a/cmd/gateway_cron.go +++ b/cmd/gateway_cron.go @@ -17,6 +17,10 @@ import ( // makeCronJobHandler creates a cron job handler that routes through the scheduler's cron lane. // This ensures per-session concurrency control (same job can't run concurrently) // and integration with /stop, /stopall commands. +// cronHeartbeatWakeFn holds the heartbeat wake function, set after ticker creation. +// Safe because cron jobs only fire after Start(), well after this is set. +var cronHeartbeatWakeFn func(agentID string) + func makeCronJobHandler(sched *scheduler.Scheduler, msgBus *bus.MessageBus, cfg *config.Config, channelMgr *channels.Manager) func(job *store.CronJob) (*store.CronJobResult, error) { return func(job *store.CronJob) (*store.CronJobResult, error) { agentID := job.AgentID @@ -102,6 +106,11 @@ func makeCronJobHandler(sched *scheduler.Scheduler, msgBus *bus.MessageBus, cfg cronResult.OutputTokens = result.Usage.CompletionTokens } + // wakeMode: trigger heartbeat after cron job completes + if job.Payload.WakeHeartbeat && cronHeartbeatWakeFn != nil { + cronHeartbeatWakeFn(agentID) + } + return cronResult, nil } } diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 4fe804db..41224eb8 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -363,6 +363,7 @@ type RunRequest struct { TraceName string // override trace name (default: "chat ") TraceTags []string // additional tags for the trace (e.g. "cron") MaxIterations int // per-request override (0 = use agent default, must be lower) + LightContext bool // skip loading context files (only inject ExtraSystemPrompt) // Run classification RunKind string // "delegation", "announce" — empty for user-initiated runs diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index cf59767f..41973a40 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -3,6 +3,7 @@ package scheduler import ( "context" "log/slog" + "strings" "sync" "sync/atomic" @@ -143,6 +144,20 @@ func (s *Scheduler) Stop() { s.lanes.StopAll() } +// HasActiveSessionsForAgent checks if any session queue for the given agent has active runs. +// Used by heartbeat ticker to skip agents currently processing other requests. +func (s *Scheduler) HasActiveSessionsForAgent(agentID string) bool { + prefix := "agent:" + agentID + ":" + s.mu.RLock() + defer s.mu.RUnlock() + for key, sq := range s.sessions { + if strings.HasPrefix(key, prefix) && sq.IsActive() { + return true + } + } + return false +} + // LaneStats returns utilization metrics for all lanes. func (s *Scheduler) LaneStats() []LaneStats { return s.lanes.AllStats() diff --git a/internal/store/cron_store.go b/internal/store/cron_store.go index 9f53fcc4..29e942f9 100644 --- a/internal/store/cron_store.go +++ b/internal/store/cron_store.go @@ -28,12 +28,13 @@ type CronSchedule struct { // CronPayload describes what a job does when triggered. type CronPayload struct { - Kind string `json:"kind"` - Message string `json:"message"` - Command string `json:"command,omitempty"` - Deliver bool `json:"deliver"` - Channel string `json:"channel,omitempty"` - To string `json:"to,omitempty"` + Kind string `json:"kind"` + Message string `json:"message"` + Command string `json:"command,omitempty"` + Deliver bool `json:"deliver"` + Channel string `json:"channel,omitempty"` + To string `json:"to,omitempty"` + WakeHeartbeat bool `json:"wake_heartbeat,omitempty"` // trigger heartbeat after job completes } // CronJobState tracks runtime state for a job. @@ -75,6 +76,7 @@ type CronJobPatch struct { Channel *string `json:"channel,omitempty"` To *string `json:"to,omitempty"` DeleteAfterRun *bool `json:"deleteAfterRun,omitempty"` + WakeHeartbeat *bool `json:"wakeHeartbeat,omitempty"` } // CronEvent represents a job lifecycle event sent to subscribers. diff --git a/internal/store/pg/cron_update.go b/internal/store/pg/cron_update.go index 150857be..01e1f545 100644 --- a/internal/store/pg/cron_update.go +++ b/internal/store/pg/cron_update.go @@ -148,7 +148,7 @@ func (s *PGCronStore) UpdateJob(jobID string, patch store.CronJobPatch) (*store. } // Update payload JSONB — fetch current, merge patch fields, re-serialize - needsPayloadUpdate := patch.Message != "" || patch.Deliver != nil || patch.Channel != nil || patch.To != nil + needsPayloadUpdate := patch.Message != "" || patch.Deliver != nil || patch.Channel != nil || patch.To != nil || patch.WakeHeartbeat != nil if needsPayloadUpdate { var payloadJSON []byte if scanErr := s.db.QueryRow("SELECT payload FROM cron_jobs WHERE id = $1", id).Scan(&payloadJSON); scanErr == nil { @@ -167,6 +167,9 @@ func (s *PGCronStore) UpdateJob(jobID string, patch store.CronJobPatch) (*store. if patch.To != nil { payload.To = *patch.To } + if patch.WakeHeartbeat != nil { + payload.WakeHeartbeat = *patch.WakeHeartbeat + } merged, _ := json.Marshal(payload) updates["payload"] = merged diff --git a/internal/tools/cron.go b/internal/tools/cron.go index 026c8710..a5a179fd 100644 --- a/internal/tools/cron.go +++ b/internal/tools/cron.go @@ -267,6 +267,14 @@ func (t *CronTool) handleAdd(ctx context.Context, args map[string]any, agentID, return ErrorResult(fmt.Sprintf("failed to create cron job: %v", err)) } + // Set wake_heartbeat if requested (triggers heartbeat after cron job completes) + if wh, _ := jobObj["wake_heartbeat"].(bool); wh { + wakeTrue := true + if updated, uErr := t.cronStore.UpdateJob(job.ID, store.CronJobPatch{WakeHeartbeat: &wakeTrue}); uErr == nil { + job = updated + } + } + data, _ := json.MarshalIndent(map[string]any{"job": job}, "", " ") return NewResult(string(data)) }