Skip to content
This repository was archived by the owner on Aug 10, 2026. It is now read-only.

Commit fe563e0

Browse files
pedronauckpedronauck
authored
refactor: orchestration improvements (#230)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added task run recovery capability for escalated runs requiring operator attention. * Added scheduler starvation detection with multi-stage escalation for long-queued runs. * Introduced "needs_attention" task run status for runs requiring intervention. * Added scheduler metrics displaying starved and attention-requiring run counts. * Added `recover` CLI command and corresponding API endpoint for run recovery. <!-- review_stack_entry_start --> [![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/compozy/agh/pull/230?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent d3b24f9 commit fe563e0

94 files changed

Lines changed: 6909 additions & 1008 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎internal/api/contract/scheduler.go‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,16 @@ import (
88

99
// SchedulerStatusPayload exposes scheduler-wide pause state and queue pressure.
1010
type SchedulerStatusPayload struct {
11-
Paused bool `json:"paused"`
12-
PausedBy string `json:"paused_by,omitempty"`
13-
PausedAt *time.Time `json:"paused_at,omitempty"`
14-
PausedReason string `json:"paused_reason,omitempty"`
15-
ActiveClaimCount int `json:"active_claim_count"`
16-
QueuedRunCount int `json:"queued_run_count"`
17-
PausedTaskCount int `json:"paused_task_count"`
18-
AsOf time.Time `json:"as_of"`
11+
Paused bool `json:"paused"`
12+
PausedBy string `json:"paused_by,omitempty"`
13+
PausedAt *time.Time `json:"paused_at,omitempty"`
14+
PausedReason string `json:"paused_reason,omitempty"`
15+
ActiveClaimCount int `json:"active_claim_count"`
16+
QueuedRunCount int `json:"queued_run_count"`
17+
PausedTaskCount int `json:"paused_task_count"`
18+
StarvedRunCount int `json:"starved_run_count"`
19+
NeedsAttentionRunCount int `json:"needs_attention_run_count"`
20+
AsOf time.Time `json:"as_of"`
1921
}
2022

2123
// SchedulerStatusResponse wraps scheduler status.

‎internal/api/contract/tasks.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -752,6 +752,12 @@ type RetryTaskRunRequest struct {
752752
Metadata json.RawMessage `json:"metadata,omitempty"`
753753
}
754754

755+
// RecoverTaskRunRequest is the shared recovery request payload for a needs_attention run.
756+
type RecoverTaskRunRequest struct {
757+
Reason string `json:"reason,omitempty"`
758+
Metadata json.RawMessage `json:"metadata,omitempty"`
759+
}
760+
755761
// PauseTaskRequest captures one per-task pause request.
756762
type PauseTaskRequest struct {
757763
Reason string `json:"reason"`

‎internal/api/core/scheduler.go‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -247,14 +247,16 @@ func (h *BaseHandlers) GetSchedulerBacklog(c *gin.Context) {
247247
// SchedulerStatusPayloadFromDomain keeps transports on one scheduler DTO shape.
248248
func SchedulerStatusPayloadFromDomain(status taskpkg.SchedulerStatus) contract.SchedulerStatusPayload {
249249
return contract.SchedulerStatusPayload{
250-
Paused: status.Paused,
251-
PausedBy: status.PausedBy,
252-
PausedAt: optionalTime(status.PausedAt),
253-
PausedReason: status.PausedReason,
254-
ActiveClaimCount: status.ActiveClaimCount,
255-
QueuedRunCount: status.QueuedRunCount,
256-
PausedTaskCount: status.PausedTaskCount,
257-
AsOf: status.AsOf,
250+
Paused: status.Paused,
251+
PausedBy: status.PausedBy,
252+
PausedAt: optionalTime(status.PausedAt),
253+
PausedReason: status.PausedReason,
254+
ActiveClaimCount: status.ActiveClaimCount,
255+
QueuedRunCount: status.QueuedRunCount,
256+
PausedTaskCount: status.PausedTaskCount,
257+
StarvedRunCount: status.StarvedRunCount,
258+
NeedsAttentionRunCount: status.NeedsAttentionRunCount,
259+
AsOf: status.AsOf,
258260
}
259261
}
260262

‎internal/api/core/tasks.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ const (
4141
taskActionForceReleaseRun = "force_release_run"
4242
taskActionForceFailRun = "force_fail_run"
4343
taskActionRetryRun = "retry_run"
44+
taskActionRecoverRun = "recover_run"
4445
taskActionBulkReleaseRuns = "bulk_release_runs"
4546
taskActionBulkFailRuns = "bulk_fail_runs"
4647
taskActionCancelRun = "cancel_run"
@@ -1467,6 +1468,44 @@ func (h *BaseHandlers) RetryTaskRun(c *gin.Context) {
14671468
c.JSON(http.StatusCreated, RetryTaskRunResponseFromResult(result))
14681469
}
14691470

1471+
// RecoverTaskRun terminalizes one needs_attention run and queues a fresh child to resume work.
1472+
func (h *BaseHandlers) RecoverTaskRun(c *gin.Context) {
1473+
manager, ok := h.requireTaskManager(c)
1474+
if !ok {
1475+
return
1476+
}
1477+
runID, err := requiredPathID(c.Param("id"), "run id")
1478+
if err != nil {
1479+
h.respondError(c, StatusForTaskError(err), err)
1480+
return
1481+
}
1482+
var req contract.RecoverTaskRunRequest
1483+
if err := c.ShouldBindJSON(&req); err != nil && !errors.Is(err, io.EOF) {
1484+
h.respondError(
1485+
c,
1486+
http.StatusBadRequest,
1487+
NewTaskValidationError(fmt.Errorf("%s: decode recover run request: %w", h.transportName(), err)),
1488+
)
1489+
return
1490+
}
1491+
actor, err := h.taskActorContext(c, taskActionRecoverRun)
1492+
if err != nil {
1493+
h.respondError(c, StatusForTaskError(err), err)
1494+
return
1495+
}
1496+
result, err := manager.RecoverRun(
1497+
c.Request.Context(),
1498+
runID,
1499+
taskpkg.RecoverRunRequest{Reason: req.Reason, Metadata: req.Metadata},
1500+
actor,
1501+
)
1502+
if err != nil {
1503+
h.respondError(c, StatusForTaskError(err), err)
1504+
return
1505+
}
1506+
c.JSON(http.StatusCreated, RetryTaskRunResponseFromResult(result))
1507+
}
1508+
14701509
// BulkForceReleaseTaskRuns force releases a bounded set of runs.
14711510
func (h *BaseHandlers) BulkForceReleaseTaskRuns(c *gin.Context) {
14721511
h.bulkForceTaskRuns(c, taskActionBulkReleaseRuns, false)

‎internal/api/httpapi/handlers_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,7 @@ func TestRegisterRoutesCoversTechSpecEndpoints(t *testing.T) {
264264
"POST /api/providers/:provider_id/auth/probe",
265265
"POST /api/runs/:id/fail",
266266
"POST /api/runs/:id/release",
267+
"POST /api/runs/:id/recover",
267268
"POST /api/runs/:id/retry",
268269
"POST /api/runs/bulk/fail",
269270
"POST /api/runs/bulk/release",
@@ -410,6 +411,7 @@ func TestRegisterTaskRoutesUseSharedHandlerBindings(t *testing.T) {
410411
"GET /api/scheduler": "GetScheduler",
411412
"GET /api/scheduler/backlog": "GetSchedulerBacklog",
412413
"POST /api/runs/:id/fail": "ForceFailTaskRun",
414+
"POST /api/runs/:id/recover": "RecoverTaskRun",
413415
"POST /api/runs/:id/release": "ForceReleaseTaskRun",
414416
"POST /api/runs/:id/retry": "RetryTaskRun",
415417
"POST /api/runs/bulk/fail": "BulkForceFailTaskRuns",

‎internal/api/httpapi/routes.go‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -311,13 +311,7 @@ func registerTaskRoutes(api gin.IRouter, handlers *Handlers) {
311311
taskRuns.POST("/:id/fail", handlers.FailTaskRun)
312312
taskRuns.POST("/:id/cancel", handlers.CancelTaskRun)
313313

314-
runs := api.Group("/runs")
315-
runs.POST("/bulk/release", handlers.BulkForceReleaseTaskRuns)
316-
runs.POST("/bulk/fail", handlers.BulkForceFailTaskRuns)
317-
runs.GET("/:id/inspect", handlers.InspectRun)
318-
runs.POST("/:id/release", handlers.ForceReleaseTaskRun)
319-
runs.POST("/:id/fail", handlers.ForceFailTaskRun)
320-
runs.POST("/:id/retry", handlers.RetryTaskRun)
314+
registerRunForceRoutes(api, handlers)
321315

322316
scheduler := api.Group("/scheduler")
323317
scheduler.GET("", handlers.GetScheduler)
@@ -331,6 +325,17 @@ func registerTaskRoutes(api gin.IRouter, handlers *Handlers) {
331325
taskReviews.POST("/:id/verdict", handlers.SubmitTaskRunReviewVerdict)
332326
}
333327

328+
func registerRunForceRoutes(api gin.IRouter, handlers *Handlers) {
329+
runs := api.Group("/runs")
330+
runs.POST("/bulk/release", handlers.BulkForceReleaseTaskRuns)
331+
runs.POST("/bulk/fail", handlers.BulkForceFailTaskRuns)
332+
runs.GET("/:id/inspect", handlers.InspectRun)
333+
runs.POST("/:id/release", handlers.ForceReleaseTaskRun)
334+
runs.POST("/:id/fail", handlers.ForceFailTaskRun)
335+
runs.POST("/:id/retry", handlers.RetryTaskRun)
336+
runs.POST("/:id/recover", handlers.RecoverTaskRun)
337+
}
338+
334339
func registerSkillRoutes(api gin.IRouter, handlers *Handlers) {
335340
privileged := handlers.privilegedMutationGuard()
336341
skillsGroup := api.Group("/skills")

‎internal/api/spec/spec.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ const (
6666
specAPIRunsIDInspectPath = "/api/runs/{id}/inspect"
6767
specAPIRunsIDReleasePath = "/api/runs/{id}/release"
6868
specAPIRunsIDRetryPath = "/api/runs/{id}/retry"
69+
specAPIRunsIDRecoverPath = "/api/runs/{id}/recover"
6970
specAPISchedulerPath = "/api/scheduler"
7071
specAPISchedulerBacklogPath = "/api/scheduler/backlog"
7172
specAPISchedulerDrainPath = "/api/scheduler/drain"
@@ -4138,6 +4139,28 @@ var operationRegistry = []OperationSpec{
41384139
{Status: 500, Description: specInternalServerErrorDescription, Body: contract.ErrorPayload{}},
41394140
},
41404141
},
4142+
{
4143+
Method: httpMethodPost,
4144+
Path: specAPIRunsIDRecoverPath,
4145+
OperationID: "recoverTaskRun",
4146+
Summary: "Recover one needs_attention task run",
4147+
Tags: []string{specTasksKey},
4148+
Transports: []Transport{TransportHTTP, TransportUDS},
4149+
Parameters: []ParameterSpec{
4150+
pathParam("id", "Task run id"),
4151+
},
4152+
RequestBody: contract.RecoverTaskRunRequest{},
4153+
Responses: []ResponseSpec{
4154+
{Status: 201, Description: specCreatedDescription, Body: contract.RetryTaskRunResponse{}},
4155+
{Status: 403, Description: specForceOperationForbiddenDescription, Body: contract.ErrorPayload{}},
4156+
{Status: 404, Description: specTaskRunNotFoundDescription, Body: contract.ErrorPayload{}},
4157+
{Status: 409, Description: "Task-run recovery conflict", Body: contract.ErrorPayload{}},
4158+
{Status: 422, Description: "Invalid recovery request", Body: contract.ErrorPayload{}},
4159+
{Status: 429, Description: specForceOperationRateLimitExceededDescription, Body: contract.ErrorPayload{}},
4160+
{Status: 503, Description: specTaskServiceIsNotConfiguredDescription, Body: contract.ErrorPayload{}},
4161+
{Status: 500, Description: specInternalServerErrorDescription, Body: contract.ErrorPayload{}},
4162+
},
4163+
},
41414164
{
41424165
Method: httpMethodPost,
41434166
Path: specAPIRunsBulkReleasePath,
@@ -6333,6 +6356,7 @@ func taskRunStatusValues() []string {
63336356
string(taskpkg.TaskRunStatusCompleted),
63346357
string(taskpkg.TaskRunStatusFailed),
63356358
string(taskpkg.TaskRunStatusCanceled),
6359+
string(taskpkg.TaskRunStatusNeedsAttention),
63366360
}
63376361
}
63386362

‎internal/api/spec/spec_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1181,6 +1181,7 @@ func TestDocumentTracksRequiredFieldsAndEnums(t *testing.T) {
11811181
string(taskpkg.TaskRunStatusCompleted),
11821182
string(taskpkg.TaskRunStatusFailed),
11831183
string(taskpkg.TaskRunStatusCanceled),
1184+
string(taskpkg.TaskRunStatusNeedsAttention),
11841185
)
11851186
assertEnumValues(t, propertySchema(t, propertySchema(t, runSchema, "origin"), "kind"),
11861187
string(taskpkg.OriginKindCLI),
@@ -1639,6 +1640,7 @@ func TestEnumHelpersReturnStableValues(t *testing.T) {
16391640
string(taskpkg.TaskRunStatusCompleted),
16401641
string(taskpkg.TaskRunStatusFailed),
16411642
string(taskpkg.TaskRunStatusCanceled),
1643+
string(taskpkg.TaskRunStatusNeedsAttention),
16421644
},
16431645
},
16441646
{

‎internal/api/testutil/task_stub.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,13 @@ type retryRunFunc func(
2828
taskpkg.ActorContext,
2929
) (*taskpkg.RetryRunResult, error)
3030

31+
type recoverRunFunc func(
32+
context.Context,
33+
string,
34+
taskpkg.RecoverRunRequest,
35+
taskpkg.ActorContext,
36+
) (*taskpkg.RetryRunResult, error)
37+
3138
type bulkForceRunFunc func(
3239
context.Context,
3340
taskpkg.BulkForceRunRequest,
@@ -136,6 +143,7 @@ type StubTaskManager struct {
136143
ForceReleaseRunFn forceReleaseRunFunc
137144
ForceFailRunFn forceFailRunFunc
138145
RetryRunFn retryRunFunc
146+
RecoverRunFn recoverRunFunc
139147
BulkForceReleaseRunsFn bulkForceRunFunc
140148
BulkForceFailRunsFn bulkForceRunFunc
141149
PauseTaskFn pauseTaskFunc
@@ -574,6 +582,18 @@ func (s StubTaskManager) RetryRun(
574582
return nil, taskpkg.ErrTaskRunNotFound
575583
}
576584

585+
func (s StubTaskManager) RecoverRun(
586+
ctx context.Context,
587+
runID string,
588+
req taskpkg.RecoverRunRequest,
589+
actor taskpkg.ActorContext,
590+
) (*taskpkg.RetryRunResult, error) {
591+
if s.RecoverRunFn != nil {
592+
return s.RecoverRunFn(ctx, runID, req, actor)
593+
}
594+
return nil, taskpkg.ErrTaskRunNotFound
595+
}
596+
577597
func (s StubTaskManager) BulkForceReleaseRuns(
578598
ctx context.Context,
579599
req taskpkg.BulkForceRunRequest,

‎internal/api/udsapi/handlers_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -391,6 +391,7 @@ func TestRegisterRoutesCoversTechSpecEndpoints(t *testing.T) {
391391
"POST /api/model-catalog/*catalog_path",
392392
"POST /api/providers/:provider_id/auth/probe",
393393
"POST /api/runs/:id/fail",
394+
"POST /api/runs/:id/recover",
394395
"POST /api/runs/:id/release",
395396
"POST /api/runs/:id/retry",
396397
"POST /api/runs/bulk/fail",
@@ -816,6 +817,7 @@ func TestRegisterTaskRoutesUseSharedHandlerBindings(t *testing.T) {
816817
"GET /api/observe/tasks/inbox": "TaskInbox",
817818
"GET /api/runs/:id/inspect": "InspectRun",
818819
"POST /api/runs/:id/fail": "ForceFailTaskRun",
820+
"POST /api/runs/:id/recover": "RecoverTaskRun",
819821
"POST /api/runs/:id/release": "ForceReleaseTaskRun",
820822
"POST /api/runs/:id/retry": "RetryTaskRun",
821823
"POST /api/runs/bulk/fail": "BulkForceFailTaskRuns",

0 commit comments

Comments
 (0)