From 4f57425224282870f99c7765d5fdfe8a315c4e29 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 12:47:35 -0400 Subject: [PATCH 1/8] fix(control-plane): admit executions before persistence on every dispatch lane MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-agent concurrency limit and the LLM circuit breaker were checked before persistence only on the async lane. On the sync, restart and MCP lanes the check ran *after* prepare had already written an executions row, a workflow_executions row and an input payload blob, so a gate-rejected request was charged a failed execution for work that was never attempted. There is now a single admission point, ahead of persistence, on all four lanes: prepareExecutionForTargetWithAdmission takes acquireSlot=true from the sync handler, the restart handler and the MCP start_run path, and the duplicate post-prepare gate blocks are gone. findReplayHit moves ahead of the gate so a replay hit — which never dials the agent — is never rejected by it and consumes no slot; the async lane no longer acquires and releases a slot for one. preparedExecution.slotHeld records whether a plan actually owns a slot, so every release site releases exactly what it took. Two other holes in the rejection contract close with it: - handleAsync abandoned the already-persisted row in "running" when submitReserved found a stopped pool. It now terminates it through failForControlPlaneShutdown (failed + status_reason control_plane_shutdown on both tables), on a detached context because the request context is very likely being cancelled by the same drain. - The restart and MCP queue-full paths reserved pool capacity only after prepare, so a queue-full burst wrote rows and then failed them, and the restart lane persisted status_reason internal_error while answering concurrency_limit. reserve() is hoisted ahead of prepare on both, and a single typed executionPreconditionError now feeds both failExecution and the response. Retry-After is completed at the same time: writeExecutionError takes the value stamped on the error, else a per-category default, so llm_unavailable advertises the circuit breaker's remaining recovery window (default 30s, floor 1s) via the new LLMHealthMonitor.RetryAfterSeconds — circuitOpenedAt is unexported, so the window has to be computed inside services — while concurrency_limit and node_unavailable stay at 1 and non-retryable rejections (413, agent_pending_approval) still carry nothing. The stale reservation comment above pool.reserve() is corrected: the worker releases the reservation when its job returns, not on dequeue, so a reservation covers preparation, queue wait and the whole dispatch. Both gate halves are opt-in and off by default (AGENTFIELD_MAX_CONCURRENT_PER_AGENT=0, llm_health.enabled=false), so a stock deployment sees no behaviour change. Refs #986 Co-Authored-By: Claude Fable 5 --- control-plane/internal/handlers/execute.go | 27 +++++---- .../internal/handlers/execute_helpers.go | 18 ++++-- .../internal/handlers/execute_lifecycle.go | 1 + .../internal/handlers/execute_prepare.go | 58 ++++++++++--------- .../internal/handlers/execute_restart.go | 34 ++++++----- .../internal/handlers/execution_guards.go | 23 ++++---- control-plane/internal/handlers/mcp.go | 27 +++++---- .../internal/services/llm_health_monitor.go | 24 ++++++++ 8 files changed, 134 insertions(+), 78 deletions(-) diff --git a/control-plane/internal/handlers/execute.go b/control-plane/internal/handlers/execute.go index b04893572..e96202ef6 100644 --- a/control-plane/internal/handlers/execute.go +++ b/control-plane/internal/handlers/execute.go @@ -266,6 +266,9 @@ func (c *executionController) handleSync(ctx *gin.Context) { return } plan.executionMode = "sync" + if plan.slotHeld { + defer ReleaseExecutionSlot(plan.target.NodeID) + } if plan.replayHit != nil { if err := c.completeReplayHit(reqCtx, plan); err != nil { @@ -287,14 +290,6 @@ func (c *executionController) handleSync(ctx *gin.Context) { return } - // Check LLM health and per-agent concurrency limits before proceeding - if err := CheckExecutionPreconditions(plan.target.NodeID, plan.llmEndpoint); err != nil { - _ = c.failExecution(reqCtx, plan, err, 0, nil) - writeExecutionError(ctx, err) - return - } - defer ReleaseExecutionSlot(plan.target.NodeID) - // Emit execution started event with full reasoner context c.publishExecutionStartedEvent(plan) @@ -434,8 +429,10 @@ func (c *executionController) handleSync(ctx *gin.Context) { func (c *executionController) handleAsync(ctx *gin.Context) { reqCtx := ctx.Request.Context() pool := getAsyncWorkerPool() - // A reservation covers both preparation and time waiting in the queue. The - // worker releases it on dequeue so only not-yet-started work consumes capacity. + // A reservation covers preparation, queue wait, and the worker's dispatch. It + // is released when the worker job returns (early on an agent HTTP 202 ACK). + // Workers plus queue capacity bound admitted work; a paused execution pins a + // worker and reservation for up to 24 hours. if !pool.reserve() { writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue is full; retry later") return @@ -455,7 +452,6 @@ func (c *executionController) handleAsync(ctx *gin.Context) { plan.executionMode = "async" if plan.replayHit != nil { - ReleaseExecutionSlot(plan.target.NodeID) if err := c.completeReplayHit(reqCtx, plan); err != nil { writeExecutionError(ctx, err) return @@ -496,7 +492,14 @@ func (c *executionController) handleAsync(ctx *gin.Context) { } if ok := pool.submitReserved(job); !ok { - ReleaseExecutionSlot(plan.target.NodeID) // Release since process() won't run + // The pool only refuses a reserved submission once it has stopped, i.e. + // the control plane is draining. Persist the outcome on a detached + // context: the request context is very likely being cancelled by the + // same shutdown, and a cancelled write would strand this row in + // "running" — exactly the orphan this branch exists to prevent. + persistCtx, persistCancel := shutdownPersistenceContext() + job.failForControlPlaneShutdown(persistCtx) + persistCancel() writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue stopped; retry later") return } diff --git a/control-plane/internal/handlers/execute_helpers.go b/control-plane/internal/handlers/execute_helpers.go index 955788821..e91d4527d 100644 --- a/control-plane/internal/handlers/execute_helpers.go +++ b/control-plane/internal/handlers/execute_helpers.go @@ -573,9 +573,17 @@ func writeExecutionError(ctx *gin.Context, err error) { body["error"] = code body["message"] = pe.Error() } - if pe.Category() == ErrorCategoryConcurrencyLimit || pe.Category() == ErrorCategoryNodeUnavailable { - ctx.Header("Retry-After", "1") - body["retry_after"] = 1 + retryAfter := pe.retryAfter + if retryAfter <= 0 { + retryAfter = map[ErrorCategory]int{ + ErrorCategoryConcurrencyLimit: 1, + ErrorCategoryNodeUnavailable: 1, + ErrorCategoryLLMUnavailable: 30, + }[pe.Category()] + } + if retryAfter > 0 { + ctx.Header("Retry-After", strconv.Itoa(retryAfter)) + body["retry_after"] = retryAfter } ctx.JSON(pe.HTTPStatusCode(), body) return @@ -840,7 +848,7 @@ func (j asyncExecutionJob) process() { func (j asyncExecutionJob) processWithContext(workerCtx context.Context) { // Release the per-agent concurrency slot when this job finishes - if j.plan.target != nil { + if j.plan.slotHeld && j.plan.target != nil { defer ReleaseExecutionSlot(j.plan.target.NodeID) } @@ -1051,7 +1059,7 @@ func shutdownPersistenceContext() (context.Context, context.CancelFunc) { } func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context) { - if j.plan.target != nil { + if j.plan.slotHeld && j.plan.target != nil { ReleaseExecutionSlot(j.plan.target.NodeID) } shutdownErr := &executionPreconditionError{ diff --git a/control-plane/internal/handlers/execute_lifecycle.go b/control-plane/internal/handlers/execute_lifecycle.go index ae4639fea..29ffab754 100644 --- a/control-plane/internal/handlers/execute_lifecycle.go +++ b/control-plane/internal/handlers/execute_lifecycle.go @@ -197,6 +197,7 @@ type preparedExecution struct { targetType string executionMode string llmEndpoint string + slotHeld bool webhookRegistered bool webhookError *string // DID context forwarded to the target agent. diff --git a/control-plane/internal/handlers/execute_prepare.go b/control-plane/internal/handlers/execute_prepare.go index 8430891d2..b98130e39 100644 --- a/control-plane/internal/handlers/execute_prepare.go +++ b/control-plane/internal/handlers/execute_prepare.go @@ -17,7 +17,7 @@ import ( ) func (c *executionController) prepareExecution(ctx context.Context, ginCtx *gin.Context) (*preparedExecution, error) { - return c.prepareExecutionWithAdmission(ctx, ginCtx, false) + return c.prepareExecutionWithAdmission(ctx, ginCtx, true) } func (c *executionController) prepareAsyncExecution(ctx context.Context, ginCtx *gin.Context) (*preparedExecution, error) { @@ -41,10 +41,6 @@ func (c *executionController) prepareExecutionWithAdmission(ctx context.Context, ) } -func (c *executionController) prepareExecutionForTarget(ctx context.Context, targetParam string, req ExecuteRequest, headers executionHeaders, callerDID, targetDID string) (*preparedExecution, error) { - return c.prepareExecutionForTargetWithAdmission(ctx, targetParam, req, headers, callerDID, targetDID, false) -} - func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context.Context, targetParam string, req ExecuteRequest, headers executionHeaders, callerDID, targetDID string, acquireSlot bool) (_ *preparedExecution, retErr error) { target, err := parseTarget(targetParam) if err != nil { @@ -145,10 +141,37 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context } target.TargetType = targetType + clientPayload := map[string]interface{}{ + "input": req.Input, + } + if len(req.Context) > 0 { + clientPayload["context"] = req.Context + } + + storedPayload, err := json.Marshal(clientPayload) + if err != nil { + return nil, fmt.Errorf("encode execution payload: %w", err) + } + + hit, err := c.findReplayHit(ctx, headers, target, storedPayload) + if err != nil { + return nil, err + } + + runID := headers.runID + if runID == "" { + runID = utils.GenerateRunID() + } + llmEndpoint := extractRequestedLLMEndpoint(req) slotAcquired := false - if acquireSlot { + if acquireSlot && hit == nil { if err := CheckExecutionPreconditions(target.NodeID, llmEndpoint); err != nil { + logger.Logger.Warn(). + Str("node_id", target.NodeID). + Str("error_category", string(classifyExecutionError(err))). + Str("run_id", runID). + Msg("execution rejected by admission gate") return nil, err } slotAcquired = true @@ -159,26 +182,9 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context }() } - runID := headers.runID - if runID == "" { - runID = utils.GenerateRunID() - } - executionID := utils.GenerateExecutionID() now := time.Now().UTC() - clientPayload := map[string]interface{}{ - "input": req.Input, - } - if len(req.Context) > 0 { - clientPayload["context"] = req.Context - } - - storedPayload, err := json.Marshal(clientPayload) - if err != nil { - return nil, fmt.Errorf("encode execution payload: %w", err) - } - exec := &types.Execution{ ExecutionID: executionID, RunID: runID, @@ -252,11 +258,6 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context c.ensureWorkflowExecutionRecord(ctx, exec, target, storedPayload) - hit, err := c.findReplayHit(ctx, headers, target, storedPayload) - if err != nil { - return nil, err - } - return &preparedExecution{ exec: exec, requestBody: agentPayloadBytes, @@ -264,6 +265,7 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context target: target, targetType: targetType, llmEndpoint: llmEndpoint, + slotHeld: slotAcquired, webhookRegistered: webhookRegistered, webhookError: webhookError, callerDID: callerDID, diff --git a/control-plane/internal/handlers/execute_restart.go b/control-plane/internal/handlers/execute_restart.go index 82e7da070..3fb0047e0 100644 --- a/control-plane/internal/handlers/execute_restart.go +++ b/control-plane/internal/handlers/execute_restart.go @@ -151,22 +151,28 @@ func (c *executionController) handleRestart(ctx *gin.Context) { } target := fmt.Sprintf("%s.%s", restartExec.NodeID, restartExec.ReasonerID) - plan, err := c.prepareExecutionForTarget(reqCtx, target, ExecuteRequest{ + pool := getAsyncWorkerPool() + if !pool.reserve() { + writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue is full; retry later") + return + } + reserved := true + defer func() { + if reserved { + pool.releaseReservation() + } + }() + + plan, err := c.prepareExecutionForTargetWithAdmission(reqCtx, target, ExecuteRequest{ Input: input, Context: contextPayload, Webhook: req.Webhook, - }, headers, "", "") + }, headers, "", "", true) if err != nil { writeExecutionError(ctx, err) return } - if err := CheckExecutionPreconditions(plan.target.NodeID, plan.llmEndpoint); err != nil { - _ = c.failExecution(reqCtx, plan, err, 0, nil) - writeExecutionError(ctx, err) - return - } - kind := "restart" if req.Fork || req.Input != nil || req.Context != nil { kind = "fork" @@ -175,20 +181,22 @@ func (c *executionController) handleRestart(ctx *gin.Context) { c.publishExecutionStartedEvent(plan) - pool := getAsyncWorkerPool() job := asyncExecutionJob{ controller: c, plan: *plan, } - if ok := pool.submit(job); !ok { - ReleaseExecutionSlot(plan.target.NodeID) - queueErr := errors.New("async execution queue is full; retry later") + if ok := pool.submitReserved(job); !ok { + if plan.slotHeld { + ReleaseExecutionSlot(plan.target.NodeID) + } + queueErr := &executionPreconditionError{code: http.StatusServiceUnavailable, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} if updateErr := c.failExecution(reqCtx, plan, queueErr, 0, nil); updateErr != nil { logger.Logger.Error().Err(updateErr).Str("execution_id", plan.exec.ExecutionID).Msg("restart: failed to persist queue saturation") } - ctx.JSON(http.StatusServiceUnavailable, gin.H{"error": queueErr.Error(), "error_category": "concurrency_limit"}) + writeExecutionError(ctx, queueErr) return } + reserved = false createdAt := plan.exec.CreatedAt.UTC().Format(time.RFC3339) var replayBefore *string diff --git a/control-plane/internal/handlers/execution_guards.go b/control-plane/internal/handlers/execution_guards.go index ff1661e71..0a347601e 100644 --- a/control-plane/internal/handlers/execution_guards.go +++ b/control-plane/internal/handlers/execution_guards.go @@ -69,14 +69,14 @@ func checkLLMEndpointHealth(monitor *services.LLMHealthMonitor, llmEndpoint stri if status.CircuitState != services.CircuitOpen { return nil } - return newLLMUnavailableError(fmt.Sprintf("LLM backend %q unavailable", status.Name), status.LastError) + return newLLMUnavailableError(fmt.Sprintf("LLM backend %q unavailable", status.Name), status.LastError, monitor.RetryAfterSeconds(status.Name)) } } statuses := monitor.GetAllStatuses() if monitor.EndpointCount() == 1 { if unavailable := firstUnavailableEndpoint(statuses); unavailable != nil { - return newLLMUnavailableError(fmt.Sprintf("LLM backend %q unavailable", unavailable.Name), unavailable.LastError) + return newLLMUnavailableError(fmt.Sprintf("LLM backend %q unavailable", unavailable.Name), unavailable.LastError, monitor.RetryAfterSeconds(unavailable.Name)) } return nil } @@ -85,6 +85,7 @@ func checkLLMEndpointHealth(monitor *services.LLMHealthMonitor, llmEndpoint stri return newLLMUnavailableError( fmt.Sprintf("LLM backend health is degraded and request backend could not be determined (endpoint %q unavailable)", unavailable.Name), unavailable.LastError, + monitor.RetryAfterSeconds(unavailable.Name), ) } @@ -100,14 +101,15 @@ func firstUnavailableEndpoint(statuses []services.LLMEndpointStatus) *services.L return nil } -func newLLMUnavailableError(message, lastErr string) error { +func newLLMUnavailableError(message, lastErr string, retryAfter int) error { if strings.TrimSpace(lastErr) != "" { message += ": " + lastErr } return &executionPreconditionError{ - code: 503, - message: message, - category: ErrorCategoryLLMUnavailable, + code: 503, + message: message, + category: ErrorCategoryLLMUnavailable, + retryAfter: retryAfter, } } @@ -152,10 +154,11 @@ const ( // (reasoners, skills, permission middleware) for conditions like // agent_pending_approval. type executionPreconditionError struct { - code int - message string - category ErrorCategory - errorCode string + code int + message string + category ErrorCategory + errorCode string + retryAfter int } func (e *executionPreconditionError) Error() string { diff --git a/control-plane/internal/handlers/mcp.go b/control-plane/internal/handlers/mcp.go index f03fbe286..a54bd0cf9 100644 --- a/control-plane/internal/handlers/mcp.go +++ b/control-plane/internal/handlers/mcp.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "errors" "fmt" "io" "net/http" @@ -512,26 +511,34 @@ func (s *mcpServer) toolWaitRun(c *gin.Context, rawArgs json.RawMessage) (map[st // the MCP tool. It returns as soon as the job is enqueued. func (s *mcpServer) startAsyncRun(ctx context.Context, target string, input map[string]interface{}, headers executionHeaders, callerDID, targetDID string) (runID, execID string, err error) { controller := newExecutionController(s.store, s.payloads, s.webhooks, s.timeout, s.internalToken) - plan, err := controller.prepareExecutionForTarget(ctx, target, ExecuteRequest{Input: input}, headers, callerDID, targetDID) - if err != nil { - return "", "", err + pool := getAsyncWorkerPool() + if !pool.reserve() { + return "", "", &executionPreconditionError{code: 503, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} } + reserved := true + defer func() { + if reserved { + pool.releaseReservation() + } + }() - if err := CheckExecutionPreconditions(plan.target.NodeID, plan.llmEndpoint); err != nil { - _ = controller.failExecution(ctx, plan, err, 0, nil) + plan, err := controller.prepareExecutionForTargetWithAdmission(ctx, target, ExecuteRequest{Input: input}, headers, callerDID, targetDID, true) + if err != nil { return "", "", err } controller.publishExecutionStartedEvent(plan) - pool := getAsyncWorkerPool() job := asyncExecutionJob{controller: controller, plan: *plan} - if ok := pool.submit(job); !ok { - ReleaseExecutionSlot(plan.target.NodeID) - queueErr := errors.New("async execution queue is full; retry later") + if ok := pool.submitReserved(job); !ok { + if plan.slotHeld { + ReleaseExecutionSlot(plan.target.NodeID) + } + queueErr := &executionPreconditionError{code: 503, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} _ = controller.failExecution(ctx, plan, queueErr, 0, nil) return "", "", queueErr } + reserved = false return plan.exec.RunID, plan.exec.ExecutionID, nil } diff --git a/control-plane/internal/services/llm_health_monitor.go b/control-plane/internal/services/llm_health_monitor.go index dc653f1c5..eee1c5740 100644 --- a/control-plane/internal/services/llm_health_monitor.go +++ b/control-plane/internal/services/llm_health_monitor.go @@ -3,6 +3,7 @@ package services import ( "context" "fmt" + "math" "net/http" "strings" "sync" @@ -102,6 +103,29 @@ func (m *LLMHealthMonitor) EndpointCount() int { return len(m.endpoints) } +// RetryAfterSeconds returns the remaining circuit-breaker recovery window. +func (m *LLMHealthMonitor) RetryAfterSeconds(name string) int { + if m == nil { + return 30 + } + m.mu.RLock() + defer m.mu.RUnlock() + + recovery := m.config.RecoveryTimeout + if recovery <= 0 { + recovery = 30 * time.Second + } + remaining := recovery + if endpoint, ok := m.endpoints[normalizeLLMEndpointName(name)]; ok && endpoint.CircuitState == CircuitOpen && !endpoint.circuitOpenedAt.IsZero() { + remaining = recovery - time.Since(endpoint.circuitOpenedAt) + } + seconds := int(math.Ceil(remaining.Seconds())) + if seconds < 1 { + return 1 + } + return seconds +} + // Start begins the health monitoring loop. func (m *LLMHealthMonitor) Start() { if !m.config.Enabled || len(m.config.Endpoints) == 0 { From 1eb0aefdad12770f9f9363e22a12b0dbc095d5b3 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 12:47:54 -0400 Subject: [PATCH 2/8] test(control-plane): cover the execute admission gate and rejection contract One test per observable behaviour of the new admission point, written from the caller's side (HTTP status, headers, body, and what the store holds afterwards) rather than from the implementation: - sync concurrency and llm_unavailable rejections persist no executions row, no workflow_executions row and no payload blob; - a replay hit against an agent already at its cap still returns 200/202 with X-AgentField-Replay-Hit, never dials the agent, and consumes no slot; - the per-agent running count is 1 during a successful sync call and back to 0 after success, an agent 5xx and a pre-gate precondition rejection; - restart rejections (gate and queue-full) persist nothing and carry Retry-After plus retry_after; - an async request whose pool stops between reserve() and submitReserved ends as failed/control_plane_shutdown on both tables with the slot released exactly once; - writeExecutionError's Retry-After table, including that 413 and agent_pending_approval carry neither header nor field; - LLMHealthMonitor.RetryAfterSeconds counts the window down, floors at 1, and falls back to the configured recovery timeout (30s) when the circuit is closed, the endpoint is unknown, or the receiver is nil; - an MCP start_run rejected by the gate persists nothing. Two existing fixtures build a preparedExecution by hand after acquiring a slot themselves; they now set slotHeld so the job still releases what they took. TestPrepareExecution_AdditionalCoverage pins the process-global limiter to nil, because prepareExecution now acquires a slot and its four direct calls would otherwise leak counts into unrelated tests. Refs #986 Co-Authored-By: Claude Fable 5 --- .../coverage_handlers_90_additional_test.go | 6 + .../handlers/coverage_raise_88_test.go | 3 + .../internal/handlers/execute_async_test.go | 50 ++++ .../handlers/execute_sync_admission_test.go | 271 ++++++++++++++++++ control-plane/internal/handlers/mcp_test.go | 20 ++ .../services/llm_health_monitor_retry_test.go | 36 +++ 6 files changed, 386 insertions(+) create mode 100644 control-plane/internal/handlers/execute_sync_admission_test.go create mode 100644 control-plane/internal/services/llm_health_monitor_retry_test.go diff --git a/control-plane/internal/handlers/coverage_handlers_90_additional_test.go b/control-plane/internal/handlers/coverage_handlers_90_additional_test.go index 8559384f1..b2d4fe310 100644 --- a/control-plane/internal/handlers/coverage_handlers_90_additional_test.go +++ b/control-plane/internal/handlers/coverage_handlers_90_additional_test.go @@ -217,6 +217,9 @@ func TestExecutionController_CompletionAndFailureCoverage(t *testing.T) { } func TestPrepareExecution_AdditionalCoverage(t *testing.T) { + previousLimiter := concurrencyLimiter + concurrencyLimiter = nil + t.Cleanup(func() { concurrencyLimiter = previousLimiter }) gin.SetMode(gin.TestMode) t.Run("versioned serverless agent registers webhook and preserves headers", func(t *testing.T) { @@ -433,6 +436,9 @@ func TestAsyncExecutionJob_ProcessFallbackCoverage(t *testing.T) { requestBody: []byte(`{"prompt":"hello"}`), agent: agent, target: &parsedTarget{NodeID: agent.ID, TargetName: "reasoner-a", TargetType: "reasoner"}, + // The slot was acquired above on behalf of this plan, so the + // job owns it and process() must release it exactly once. + slotHeld: true, }, } diff --git a/control-plane/internal/handlers/coverage_raise_88_test.go b/control-plane/internal/handlers/coverage_raise_88_test.go index 6e04027fd..5ab8b3b9c 100644 --- a/control-plane/internal/handlers/coverage_raise_88_test.go +++ b/control-plane/internal/handlers/coverage_raise_88_test.go @@ -934,6 +934,9 @@ func TestAsyncExecutionJob_ProcessAdditionalCoverage(t *testing.T) { plan: preparedExecution{ exec: &types.Execution{ExecutionID: "exec-1"}, target: &parsedTarget{NodeID: "node-1"}, + // The slot was acquired above on behalf of this plan, so the job + // owns it and process() must release it exactly once. + slotHeld: true, }, } diff --git a/control-plane/internal/handlers/execute_async_test.go b/control-plane/internal/handlers/execute_async_test.go index a94fffd3b..05bfd076c 100644 --- a/control-plane/internal/handlers/execute_async_test.go +++ b/control-plane/internal/handlers/execute_async_test.go @@ -36,6 +36,56 @@ func useAsyncPoolForTest(t *testing.T, pool *asyncWorkerPool) { }) } +func TestExecuteAsyncHandler_PoolStoppedTerminatesPersistedRow(t *testing.T) { + gin.SetMode(gin.TestMode) + pool := newAsyncWorkerPool(0, 4) + useAsyncPoolForTest(t, pool) + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 2} + defer func() { concurrencyLimiter = oldLimiter }() + + agent := testRestartAgent("http://agent.example") + baseStore := newTestExecutionStorage(agent) + store := &stopPoolOnCreateStorage{testExecutionStorage: baseStore, pool: pool} + router := gin.New() + router.POST("/api/v1/execute/async/:target", ExecuteAsyncHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/async/node-1.reasoner-a", strings.NewReader(`{"input":{"foo":"bar"}}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + require.Equal(t, http.StatusServiceUnavailable, resp.Code) + require.Contains(t, resp.Body.String(), "async execution queue stopped") + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, float64(1), body["retry_after"]) + records, err := baseStore.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Len(t, records, 1) + require.Equal(t, types.ExecutionStatusFailed, records[0].Status) + require.NotNil(t, records[0].StatusReason) + require.Equal(t, "control_plane_shutdown", *records[0].StatusReason) + workflows, err := baseStore.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Len(t, workflows, 1) + require.NotNil(t, workflows[0].StatusReason) + require.Equal(t, "control_plane_shutdown", *workflows[0].StatusReason) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) +} + +type stopPoolOnCreateStorage struct { + *testExecutionStorage + pool *asyncWorkerPool +} + +func (s *stopPoolOnCreateStorage) CreateExecutionRecord(ctx context.Context, execution *types.Execution) error { + s.pool.mu.Lock() + s.pool.stopped = true + s.pool.mu.Unlock() + return s.testExecutionStorage.CreateExecutionRecord(ctx, execution) +} + func TestExecuteAsyncHandler_QueueSaturation(t *testing.T) { gin.SetMode(gin.TestMode) useAsyncPoolForTest(t, newAsyncWorkerPool(1, 1)) diff --git a/control-plane/internal/handlers/execute_sync_admission_test.go b/control-plane/internal/handlers/execute_sync_admission_test.go new file mode 100644 index 000000000..6903b4125 --- /dev/null +++ b/control-plane/internal/handlers/execute_sync_admission_test.go @@ -0,0 +1,271 @@ +package handlers + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Agent-Field/agentfield/control-plane/internal/config" + "github.com/Agent-Field/agentfield/control-plane/internal/services" + "github.com/Agent-Field/agentfield/control-plane/pkg/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestExecuteHandler_ConcurrencyRejectionHasNoPersistence(t *testing.T) { + gin.SetMode(gin.TestMode) + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + + agent := &types.AgentNode{ID: "node-1", BaseURL: "http://agent.example", Reasoners: []types.ReasonerDefinition{{ID: "reasoner-a"}}} + store := newTestExecutionStorage(agent) + payloadDir := t.TempDir() + router := gin.New() + router.POST("/api/v1/execute/:target", ExecuteHandler(store, services.NewFilePayloadStore(payloadDir), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/node-1.reasoner-a", strings.NewReader(`{"input":{"foo":"bar"}}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + require.Equal(t, http.StatusTooManyRequests, resp.Code) + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, "concurrency_limit", body["error_category"]) + require.Equal(t, float64(1), body["retry_after"]) + records, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Empty(t, records) + workflows, err := store.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Empty(t, workflows) + files, err := filepath.Glob(filepath.Join(payloadDir, "*")) + require.NoError(t, err) + require.Empty(t, files) +} + +func TestExecuteHandler_LLMUnavailableRejectionHasNoPersistence(t *testing.T) { + failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusServiceUnavailable) })) + defer failing.Close() + monitor := services.NewLLMHealthMonitor(config.LLMHealthConfig{Enabled: true, CheckInterval: 10 * time.Millisecond, CheckTimeout: 100 * time.Millisecond, FailureThreshold: 1, RecoveryTimeout: 30 * time.Second, Endpoints: []config.LLMEndpoint{{Name: "primary", URL: failing.URL}}}, nil) + go monitor.Start() + defer monitor.Stop() + SetLLMHealthMonitor(monitor) + defer SetLLMHealthMonitor(nil) + require.Eventually(t, func() bool { + s, ok := monitor.GetStatus("primary") + return ok && s.CircuitState == services.CircuitOpen + }, 2*time.Second, 10*time.Millisecond) + + store := newTestExecutionStorage(testRestartAgent("http://agent.example")) + payloadDir := t.TempDir() + router := gin.New() + router.POST("/api/v1/execute/:target", ExecuteHandler(store, services.NewFilePayloadStore(payloadDir), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/node-1.reasoner-a", strings.NewReader(`{"input":{"foo":"bar"}}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + require.Equal(t, http.StatusServiceUnavailable, resp.Code) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, "llm_unavailable", body["error_category"]) + require.NotEmpty(t, resp.Header().Get("Retry-After")) + require.Equal(t, fmt.Sprint(body["retry_after"]), resp.Header().Get("Retry-After")) + records, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Empty(t, records) + workflows, err := store.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Empty(t, workflows) + files, err := filepath.Glob(filepath.Join(payloadDir, "*")) + require.NoError(t, err) + require.Empty(t, files) +} + +func TestExecuteHandler_ReplayHitNotGatedBySaturatedAgent(t *testing.T) { + for _, tc := range []struct { + name, route string + status int + }{{"sync", "/api/v1/execute/:target", http.StatusOK}, {"async", "/api/v1/execute/async/:target", http.StatusAccepted}} { + t.Run(tc.name, func(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + useAsyncPoolForTest(t, newAsyncWorkerPool(1, 2)) + var calls int32 + agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { atomic.AddInt32(&calls, 1); w.WriteHeader(http.StatusOK) })) + defer agentServer.Close() + store := newTestExecutionStorage(testRestartAgent(agentServer.URL)) + now := time.Now().UTC() + seedExecutionRecord(t, store, &types.Execution{ExecutionID: "source", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", ReasonerID: "reasoner-a", Status: types.ExecutionStatusSucceeded, InputPayload: json.RawMessage(`{"input":{"foo":"bar"}}`), ResultPayload: json.RawMessage(`{"answer":42}`), StartedAt: now.Add(-time.Minute), CreatedAt: now.Add(-time.Minute), UpdatedAt: now.Add(-time.Minute)}) + seedExecutionRecord(t, store, &types.Execution{ExecutionID: "marker", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", ReasonerID: "reasoner-b", Status: types.ExecutionStatusFailed, InputPayload: json.RawMessage(`{}`), StartedAt: now, CreatedAt: now, UpdatedAt: now}) + router := gin.New() + if tc.name == "sync" { + router.POST(tc.route, ExecuteHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + } else { + router.POST(tc.route, ExecuteAsyncHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + } + req := httptest.NewRequest(http.MethodPost, strings.Replace(tc.route, ":target", "node-1.reasoner-a", 1), strings.NewReader(`{"input":{"foo":"bar"}}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Run-ID", "new-run") + req.Header.Set("X-Parent-Execution-ID", "new-parent") + req.Header.Set("X-AgentField-Replay-Source-Run-ID", "old-run") + req.Header.Set("X-AgentField-Replay-Before-Execution-ID", "marker") + req.Header.Set("X-AgentField-Replay-Mode", "succeeded-before") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + require.Equal(t, tc.status, resp.Code, resp.Body.String()) + require.Equal(t, "source", resp.Header().Get("X-AgentField-Replay-Hit")) + require.EqualValues(t, 1, concurrencyLimiter.GetRunningCount("node-1")) + require.Zero(t, atomic.LoadInt32(&calls)) + }) + } +} + +func TestExecuteHandler_SlotBalancedAcrossOutcomes(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 2} + defer func() { concurrencyLimiter = oldLimiter }() + var during int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + atomic.StoreInt32(&during, int32(concurrencyLimiter.GetRunningCount("node-1"))) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + store := newTestExecutionStorage(testRestartAgent(server.URL)) + router := gin.New() + router.POST("/api/v1/execute/:target", ExecuteHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + request := func(r *gin.Engine) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/node-1.reasoner-a", strings.NewReader(`{"input":{}}`)) + req.Header.Set("Content-Type", "application/json") + out := httptest.NewRecorder() + r.ServeHTTP(out, req) + return out + } + require.Equal(t, http.StatusOK, request(router).Code) + require.EqualValues(t, 1, atomic.LoadInt32(&during)) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) + server.Close() + errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) + defer errorServer.Close() + errorRouter := gin.New() + errorRouter.POST("/api/v1/execute/:target", ExecuteHandler(newTestExecutionStorage(testRestartAgent(errorServer.URL)), services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + require.Equal(t, http.StatusBadGateway, request(errorRouter).Code) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) + inactive := testRestartAgent("http://agent.example") + inactive.HealthStatus = types.HealthStatusInactive + inactiveRouter := gin.New() + inactiveRouter.POST("/api/v1/execute/:target", ExecuteHandler(newTestExecutionStorage(inactive), services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + require.Equal(t, http.StatusServiceUnavailable, request(inactiveRouter).Code) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) +} + +func TestWriteExecutionError_RetryAfterPerCategory(t *testing.T) { + tests := []struct { + category ErrorCategory + retry int + }{ + {ErrorCategoryConcurrencyLimit, 1}, + {ErrorCategoryNodeUnavailable, 1}, + {ErrorCategoryLLMUnavailable, 17}, + } + for _, test := range tests { + t.Run(string(test.category), func(t *testing.T) { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + err := &executionPreconditionError{code: http.StatusServiceUnavailable, message: "retry", category: test.category} + if test.category == ErrorCategoryLLMUnavailable { + err.retryAfter = test.retry + } + writeExecutionError(ctx, err) + require.Equal(t, fmt.Sprint(test.retry), recorder.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body)) + require.Equal(t, float64(test.retry), body["retry_after"]) + }) + } + t.Run("body too large is not retryable", func(t *testing.T) { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + writeExecutionError(ctx, &http.MaxBytesError{}) + require.Equal(t, http.StatusRequestEntityTooLarge, recorder.Code) + require.Empty(t, recorder.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body)) + require.NotContains(t, body, "retry_after") + }) + t.Run("pending approval is not retryable", func(t *testing.T) { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + writeExecutionError(ctx, &executionPreconditionError{code: http.StatusServiceUnavailable, message: "pending", category: ErrorCategoryAgentError, errorCode: "agent_pending_approval"}) + require.Equal(t, http.StatusServiceUnavailable, recorder.Code) + require.Empty(t, recorder.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body)) + require.NotContains(t, body, "retry_after") + }) +} + +func TestRestartHandler_ConcurrencyRejectionHasNoPersistence(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + defer func() { concurrencyLimiter = oldLimiter }() + useAsyncPoolForTest(t, newAsyncWorkerPool(0, 2)) + store := newTestExecutionStorage(testRestartAgent("http://agent.example")) + now := time.Now().UTC() + seedExecutionRecord(t, store, &types.Execution{ExecutionID: "source", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", ReasonerID: "reasoner-a", Status: types.ExecutionStatusFailed, InputPayload: json.RawMessage(`{"input":{"foo":"bar"}}`), StartedAt: now, CreatedAt: now, UpdatedAt: now}) + router := gin.New() + router.POST("/api/v1/executions/:execution_id/restart", RestartExecutionHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + require.Equal(t, http.StatusTooManyRequests, resp.Code, resp.Body.String()) + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, "concurrency_limit", body["error_category"]) + require.Equal(t, float64(1), body["retry_after"]) + records, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Len(t, records, 1) + workflows, err := store.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Empty(t, workflows) +} + +func TestRestartHandler_QueueFullCarriesRetryAfter(t *testing.T) { + pool := newAsyncWorkerPool(0, 1) + useAsyncPoolForTest(t, pool) + require.True(t, pool.reserve()) + store := newTestExecutionStorage(testRestartAgent("http://agent.example")) + now := time.Now().UTC() + seedExecutionRecord(t, store, &types.Execution{ExecutionID: "source", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", ReasonerID: "reasoner-a", Status: types.ExecutionStatusFailed, InputPayload: json.RawMessage(`{"input":{}}`), StartedAt: now, CreatedAt: now, UpdatedAt: now}) + router := gin.New() + router.POST("/api/v1/executions/:execution_id/restart", RestartExecutionHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + require.Equal(t, http.StatusServiceUnavailable, resp.Code, resp.Body.String()) + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, "concurrency_limit", body["error_category"]) + require.Equal(t, float64(1), body["retry_after"]) + records, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Len(t, records, 1) +} diff --git a/control-plane/internal/handlers/mcp_test.go b/control-plane/internal/handlers/mcp_test.go index d5e002385..4fa5926ce 100644 --- a/control-plane/internal/handlers/mcp_test.go +++ b/control-plane/internal/handlers/mcp_test.go @@ -342,6 +342,26 @@ func TestMCP_ExecuteReasoner(t *testing.T) { }) } +func TestMCP_ExecuteReasonerConcurrencyRejectionHasNoPersistence(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + require.NoError(t, concurrencyLimiter.Acquire("planner")) + defer func() { concurrencyLimiter = oldLimiter }() + useAsyncPoolForTest(t, newAsyncWorkerPool(0, 2)) + store := newMCPTestStore(mcpActiveAgent()) + router := newMCPTestRouter(t, store) + + payload, isErr := mcpCallTool(t, router, "execute_reasoner", map[string]interface{}{ + "target": "planner.plan", + "input": map[string]interface{}{"goal": "ship"}, + }) + require.True(t, isErr) + require.Contains(t, payload["text"], "reached max concurrent executions") + execs, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Empty(t, execs) +} + func TestMCP_ExecuteReasonerAuthorizesAndBindsRunToVerifiedCaller(t *testing.T) { store := newMCPTestStore(mcpActiveAgent()) var gotCaller, gotTarget string diff --git a/control-plane/internal/services/llm_health_monitor_retry_test.go b/control-plane/internal/services/llm_health_monitor_retry_test.go new file mode 100644 index 000000000..6aeada44d --- /dev/null +++ b/control-plane/internal/services/llm_health_monitor_retry_test.go @@ -0,0 +1,36 @@ +package services + +import ( + "testing" + "time" + + "github.com/Agent-Field/agentfield/control-plane/internal/config" + "github.com/stretchr/testify/assert" +) + +func TestLLMHealthMonitor_RetryAfterSeconds(t *testing.T) { + monitor := NewLLMHealthMonitor(config.LLMHealthConfig{ + RecoveryTimeout: 10 * time.Second, + Endpoints: []config.LLMEndpoint{{Name: "primary"}}, + }, nil) + + assert.Equal(t, 10, monitor.RetryAfterSeconds("unknown")) + assert.Equal(t, 10, monitor.RetryAfterSeconds("primary")) + + monitor.mu.Lock() + monitor.endpoints["primary"].CircuitState = CircuitOpen + monitor.endpoints["primary"].circuitOpenedAt = time.Now().Add(-4 * time.Second) + monitor.mu.Unlock() + remaining := monitor.RetryAfterSeconds("primary") + assert.GreaterOrEqual(t, remaining, 5) + assert.LessOrEqual(t, remaining, 6) + + monitor.mu.Lock() + monitor.endpoints["primary"].circuitOpenedAt = time.Now().Add(-20 * time.Second) + monitor.mu.Unlock() + assert.Equal(t, 1, monitor.RetryAfterSeconds("primary")) + + defaults := NewLLMHealthMonitor(config.LLMHealthConfig{}, nil) + assert.Equal(t, 30, defaults.RetryAfterSeconds("unknown")) + assert.Equal(t, 30, (*LLMHealthMonitor)(nil).RetryAfterSeconds("unknown")) +} From 907e40a0d7140bef657af66e7ca90573af07d0f4 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 12:48:07 -0400 Subject: [PATCH 3/8] docs: describe the real execute admission model, not a lease-based queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit README advertised "a durable PostgreSQL queue with lease-based processing, so a crash or a restart resumes where it left off". No such thing exists: the lease columns in migrations 011 and 013 are inert plumbing, and there is no acquisition, renewal or expiry-reclaim code anywhere in the tree. What the control plane actually does is admit work into a bounded in-process queue with backpressure (429/503 plus Retry-After) and, on graceful shutdown, terminate in-flight executions with status_reason control_plane_shutdown instead of silently dropping them. Both README claims now say that. Alongside it: - docs/api/EXECUTE.md said llm_unavailable carried no Retry-After. It now does, advertising the circuit breaker's remaining recovery window, and a new line under the table states that these pre-dispatch rejections persist no rows — with the one exception of a request rejected after preparation because the pool has already stopped. - docs/api/EXECUTION_RESTART.md records that the restart lane runs the same admission checks before persistence and returns Retry-After on queue-full. - AGENTFIELD_EXEC_ASYNC_QUEUE_CAPACITY was documented as the number of executions "waiting for a worker". The admission bound is really workers + queue_capacity and a reservation is held across preparation, queue wait and the worker's dispatch (up to 24h for a paused execution). Refs #986 Co-Authored-By: Claude Fable 5 --- README.md | 6 +++--- docs/ENVIRONMENT_VARIABLES.md | 2 +- docs/api/EXECUTE.md | 4 +++- docs/api/EXECUTION_RESTART.md | 2 ++ 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 0d62d4be7..bfe88288b 100644 --- a/README.md +++ b/README.md @@ -231,12 +231,12 @@ Most agent tools help you **write** agent logic. AgentField is what **runs** it ## How it scales -The control plane is a stateless Go service. You put more of them behind a load balancer and the fleet grows horizontally. Work lands in a durable PostgreSQL queue with lease-based processing, so a crash or a restart resumes where it left off instead of dropping the job. +The control plane is a stateless Go service. You put more of them behind a load balancer and the fleet grows horizontally. Work is admitted into a bounded in-process queue with backpressure (`429`/`503` plus `Retry-After`). On graceful shutdown, in-flight executions are terminated with `status_reason` `control_plane_shutdown` rather than silently dropped. | Property | What it means | |---|---| | Stateless Go control plane | Horizontal scaling behind a load balancer. Add replicas to add capacity. | -| Durable PostgreSQL queue | Lease-based processing. Jobs survive crashes and restarts. | +| Bounded in-process admission | Backpressure returns `429`/`503` with `Retry-After`; graceful shutdown records `control_plane_shutdown`. | | Async execution | Webhooks and SSE, no timeout limits. A single run can go for hours or days. | | Backpressure | Queue-depth limits and circuit breakers keep a fan-out from overwhelming downstream agents. | | Routing overhead | Roughly 100-200ms per cross-agent hop. It matters when a branch does little work per hop, so keep hops coarse when latency is tight. | @@ -318,7 +318,7 @@ Two examples already run at this load. The [deep-research engine](https://agentf | Progress updates mid-execution | Intermediate payloads during long tasks | | Auto retries + exponential backoff | Transparent - control plane handles | | Backpressure + queue depth limits | Fair scheduling, circuit breakers | -| Durable queue (PostgreSQL) | Atomic lease-based processing | +| Bounded in-process queue | Backpressure and explicit graceful-shutdown termination | #### Memory (Distributed State) diff --git a/docs/ENVIRONMENT_VARIABLES.md b/docs/ENVIRONMENT_VARIABLES.md index af62ca315..788cce493 100644 --- a/docs/ENVIRONMENT_VARIABLES.md +++ b/docs/ENVIRONMENT_VARIABLES.md @@ -113,7 +113,7 @@ The telemetry payload does not include prompts, inputs, outputs, logs, secrets, - `AGENTFIELD_MAX_CONCURRENT_PER_AGENT` (default: `0`): Maximum concurrent executions dispatched to one agent; `0` means unlimited. - `AGENTFIELD_EXEC_ASYNC_WORKERS` (default: the greater of the available CPU count and `16`): Worker count for asynchronous execution and restart jobs, which are I/O-bound; non-positive values use the default. -- `AGENTFIELD_EXEC_ASYNC_QUEUE_CAPACITY` (default: `1024`): Maximum number of asynchronous executions waiting for a worker; non-positive values use the default. Requests arriving once the queue is saturated are rejected with `503`, a `Retry-After` header and a `retry_after` field, and no execution row is persisted for them. +- `AGENTFIELD_EXEC_ASYNC_QUEUE_CAPACITY` (default: `1024`): Additional admitted asynchronous work beyond the worker count; `workers + queue_capacity` bounds work across preparation, queue wait, and worker dispatch. A paused execution pins a worker and reservation for up to 24 hours. Non-positive values use the default. Requests arriving once admission is saturated are rejected with `503`, a `Retry-After` header and a `retry_after` field, and no execution row is persisted for them. - `AGENTFIELD_MAX_EXECUTE_BODY_BYTES` (default: `33554432`, 32 MiB): Maximum request body size, in bytes, for POST routes under `/api/v1/execute`. Oversize requests are rejected with `413` before any execution is persisted; other routes are not capped by this setting. - `AGENTFIELD_MAX_REGISTER_BODY_BYTES` (default: `8388608`, 8 MiB): Maximum request body size, in bytes, for node registration POST routes (`/api/v1/nodes`, `/api/v1/nodes/register`, and `/api/v1/nodes/register-serverless`). Oversize requests are rejected with `413` before registration handling begins. - `AGENTFIELD_SHUTDOWN_TIMEOUT` (default: `30s`): Grace period for draining the control plane HTTP server during shutdown. diff --git a/docs/api/EXECUTE.md b/docs/api/EXECUTE.md index a6b624811..810e3bd80 100644 --- a/docs/api/EXECUTE.md +++ b/docs/api/EXECUTE.md @@ -88,7 +88,9 @@ Execute requests can be rejected before dispatch: | `503` | Async dispatch queue full | `Retry-After: 1` and `{"error":"async execution queue is full; retry later","error_category":"concurrency_limit","retry_after":1}` | | `503` | Control plane shutting down (async pool stopped) | `Retry-After: 1` and `{"error":"...","error_category":"concurrency_limit","retry_after":1}`; the `error` text distinguishes it from queue-full | | `503` | Target node known to be down (after the drain hold expires) | `Retry-After: 1` and `{"error":"...","error_category":"node_unavailable","retry_after":1}` | -| `503` | Required LLM unavailable | `{"error":"...","error_category":"llm_unavailable"}` with no `Retry-After` header | +| `503` | Required LLM unavailable | `Retry-After: ` and `{"error":"...","error_category":"llm_unavailable","retry_after":}`; the window is the circuit breaker's remaining recovery timeout (default 30s, floor 1s) | | `413` | Body exceeds `AGENTFIELD_MAX_EXECUTE_BODY_BYTES` (default 32 MiB) | `{"error":"request body too large"}` | +These pre-dispatch rejections persist no `executions` or `workflow_executions` row and no payload; the exception is a request rejected because the async pool has already stopped after preparation, which is terminated as `failed` with `status_reason` `control_plane_shutdown`. + The execute routes do not accept an idempotency key. Retrying a request can create another execution; use the [restart/replay API](EXECUTION_RESTART.md) when replaying an existing run. diff --git a/docs/api/EXECUTION_RESTART.md b/docs/api/EXECUTION_RESTART.md index 062a6f456..3a59b3aed 100644 --- a/docs/api/EXECUTION_RESTART.md +++ b/docs/api/EXECUTION_RESTART.md @@ -47,4 +47,6 @@ Operators polling execution state should branch on the stable category before an | `agent_client_error:` | The agent reported a client-facing HTTP 4xx failure. | | `llm_unavailable`, `concurrency_limit`, `agent_timeout`, `agent_error`, `agent_unreachable`, `bad_response`, `internal_error`, `validation`, `permission_denied`, `node_unavailable`, `target_not_found` | Canonical failure categories used for operator routing and HTTP mapping. | +Concurrency and LLM-circuit admission checks run before restart persistence. A rejected restart creates no execution or workflow-execution row. Queue-full restart responses return `503` with matching `Retry-After` and `retry_after` values. + Do not emulate restart by re-submitting `/execute`: execute has no idempotency key, creates unrelated executions, and cannot establish restart lineage or replay boundaries. From 7ab39445b29a1990f61de57ce9365aa2f2cebd8a Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 13:15:17 -0400 Subject: [PATCH 4/8] test(control-plane): assert the restart lane persists the category it answers with The restart handler can still lose the race between reserve() and submitReserved when the pool stops in between. That branch is the one that used to answer concurrency_limit while writing status_reason internal_error, because the queue error was an untyped errors.New. Drive it through the same CreateExecutionRecord seam the async pool-stopped test uses and assert the persisted status_reason equals the error_category in the body, that Retry-After and retry_after are both present, and that the per-agent slot is released. Refs #986 Co-Authored-By: Claude Fable 5 --- .../handlers/execute_sync_admission_test.go | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/control-plane/internal/handlers/execute_sync_admission_test.go b/control-plane/internal/handlers/execute_sync_admission_test.go index 6903b4125..360d312ad 100644 --- a/control-plane/internal/handlers/execute_sync_admission_test.go +++ b/control-plane/internal/handlers/execute_sync_admission_test.go @@ -269,3 +269,55 @@ func TestRestartHandler_QueueFullCarriesRetryAfter(t *testing.T) { require.NoError(t, err) require.Len(t, records, 1) } + +// A restart admitted past reserve() but refused by submitReserved (the pool +// stopped in between) must persist a status_reason that matches the +// error_category it answers with — it used to answer concurrency_limit while +// recording internal_error, because the queue error was untyped. +func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { + gin.SetMode(gin.TestMode) + pool := newAsyncWorkerPool(0, 4) + useAsyncPoolForTest(t, pool) + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 2} + defer func() { concurrencyLimiter = oldLimiter }() + + base := newTestExecutionStorage(testRestartAgent("http://agent.example")) + now := time.Now().UTC() + seedExecutionRecord(t, base, &types.Execution{ + ExecutionID: "source", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", + ReasonerID: "reasoner-a", Status: types.ExecutionStatusFailed, + InputPayload: json.RawMessage(`{"input":{"foo":"bar"}}`), + StartedAt: now, CreatedAt: now, UpdatedAt: now, + }) + store := &stopPoolOnCreateStorage{testExecutionStorage: base, pool: pool} + + router := gin.New() + router.POST("/api/v1/executions/:execution_id/restart", RestartExecutionHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + require.Equal(t, http.StatusServiceUnavailable, resp.Code, resp.Body.String()) + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]any + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, "concurrency_limit", body["error_category"]) + require.Equal(t, float64(1), body["retry_after"]) + + records, err := base.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Len(t, records, 2) + var restarted *types.Execution + for _, record := range records { + if record.ExecutionID != "source" { + restarted = record + } + } + require.NotNil(t, restarted) + require.Equal(t, types.ExecutionStatusFailed, restarted.Status) + require.NotNil(t, restarted.StatusReason) + require.Equal(t, string(ErrorCategoryConcurrencyLimit), *restarted.StatusReason) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) +} From 36ee9f671a7aabb3f1f6188a669ab4374c5809d4 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 14:40:40 -0400 Subject: [PATCH 5/8] fix(control-plane): terminalize pool-stopped restart/MCP admissions through a detached context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The restart and MCP submit-failure paths persisted the terminal state with the request context — during shutdown that context is likely already cancelled, stranding the freshly created rows in running (the same bug class #1001 fixed on the async lane). MCP also discarded the persistence error entirely. All three lanes now share one helper: detached bounded persistence context, failed/control_plane_shutdown on both tables, and a warn log carrying node_id and execution_id when persistence itself fails. Co-Authored-By: Claude Fable 5 --- control-plane/internal/handlers/execute.go | 4 +- .../internal/handlers/execute_async_test.go | 9 ++- .../internal/handlers/execute_helpers.go | 20 +++++- .../internal/handlers/execute_restart.go | 7 +- .../handlers/execute_sync_admission_test.go | 13 +++- control-plane/internal/handlers/mcp.go | 5 +- control-plane/internal/handlers/mcp_test.go | 71 +++++++++++++++++++ 7 files changed, 108 insertions(+), 21 deletions(-) diff --git a/control-plane/internal/handlers/execute.go b/control-plane/internal/handlers/execute.go index e96202ef6..1bf3d0b2d 100644 --- a/control-plane/internal/handlers/execute.go +++ b/control-plane/internal/handlers/execute.go @@ -497,9 +497,7 @@ func (c *executionController) handleAsync(ctx *gin.Context) { // context: the request context is very likely being cancelled by the // same shutdown, and a cancelled write would strand this row in // "running" — exactly the orphan this branch exists to prevent. - persistCtx, persistCancel := shutdownPersistenceContext() - job.failForControlPlaneShutdown(persistCtx) - persistCancel() + job.terminateForControlPlaneShutdown() writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue stopped; retry later") return } diff --git a/control-plane/internal/handlers/execute_async_test.go b/control-plane/internal/handlers/execute_async_test.go index 05bfd076c..7a0775009 100644 --- a/control-plane/internal/handlers/execute_async_test.go +++ b/control-plane/internal/handlers/execute_async_test.go @@ -76,14 +76,19 @@ func TestExecuteAsyncHandler_PoolStoppedTerminatesPersistedRow(t *testing.T) { type stopPoolOnCreateStorage struct { *testExecutionStorage - pool *asyncWorkerPool + pool *asyncWorkerPool + cancel context.CancelFunc } func (s *stopPoolOnCreateStorage) CreateExecutionRecord(ctx context.Context, execution *types.Execution) error { s.pool.mu.Lock() s.pool.stopped = true s.pool.mu.Unlock() - return s.testExecutionStorage.CreateExecutionRecord(ctx, execution) + err := s.testExecutionStorage.CreateExecutionRecord(ctx, execution) + if s.cancel != nil { + s.cancel() + } + return err } func TestExecuteAsyncHandler_QueueSaturation(t *testing.T) { diff --git a/control-plane/internal/handlers/execute_helpers.go b/control-plane/internal/handlers/execute_helpers.go index e91d4527d..e469302fb 100644 --- a/control-plane/internal/handlers/execute_helpers.go +++ b/control-plane/internal/handlers/execute_helpers.go @@ -1058,16 +1058,26 @@ func shutdownPersistenceContext() (context.Context, context.CancelFunc) { return context.WithTimeout(context.Background(), 5*time.Second) } +func (j asyncExecutionJob) terminateForControlPlaneShutdown() { + persistCtx, cancel := shutdownPersistenceContext() + defer cancel() + j.failForControlPlaneShutdown(persistCtx) +} + func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context) { + nodeID := "" + if j.plan.target != nil { + nodeID = j.plan.target.NodeID + } if j.plan.slotHeld && j.plan.target != nil { - ReleaseExecutionSlot(j.plan.target.NodeID) + ReleaseExecutionSlot(nodeID) } shutdownErr := &executionPreconditionError{ message: "execution was not started before the control plane shut down", category: ErrorCategoryControlPlaneShutdown, } if err := j.controller.failExecution(ctx, &j.plan, shutdownErr, 0, nil); err != nil { - logger.Logger.Error().Err(err).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to terminate queued execution during shutdown") + logger.Logger.Warn().Err(err).Str("node_id", nodeID).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to terminate queued execution during shutdown") return } reason := string(ErrorCategoryControlPlaneShutdown) @@ -1075,10 +1085,14 @@ func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context) { if current == nil { return nil, fmt.Errorf("workflow execution %s not found", j.plan.exec.ExecutionID) } + now := time.Now().UTC() + current.Status = string(types.ExecutionStatusFailed) current.StatusReason = &reason + current.CompletedAt = &now + current.UpdatedAt = now return current, nil }); err != nil { - logger.Logger.Error().Err(err).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to record shutdown reason on queued workflow execution") + logger.Logger.Warn().Err(err).Str("node_id", nodeID).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to record shutdown reason on queued workflow execution") } } diff --git a/control-plane/internal/handlers/execute_restart.go b/control-plane/internal/handlers/execute_restart.go index 3fb0047e0..a4bae714f 100644 --- a/control-plane/internal/handlers/execute_restart.go +++ b/control-plane/internal/handlers/execute_restart.go @@ -186,13 +186,8 @@ func (c *executionController) handleRestart(ctx *gin.Context) { plan: *plan, } if ok := pool.submitReserved(job); !ok { - if plan.slotHeld { - ReleaseExecutionSlot(plan.target.NodeID) - } + job.terminateForControlPlaneShutdown() queueErr := &executionPreconditionError{code: http.StatusServiceUnavailable, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} - if updateErr := c.failExecution(reqCtx, plan, queueErr, 0, nil); updateErr != nil { - logger.Logger.Error().Err(updateErr).Str("execution_id", plan.exec.ExecutionID).Msg("restart: failed to persist queue saturation") - } writeExecutionError(ctx, queueErr) return } diff --git a/control-plane/internal/handlers/execute_sync_admission_test.go b/control-plane/internal/handlers/execute_sync_admission_test.go index 360d312ad..5eb6edab5 100644 --- a/control-plane/internal/handlers/execute_sync_admission_test.go +++ b/control-plane/internal/handlers/execute_sync_admission_test.go @@ -290,11 +290,12 @@ func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { InputPayload: json.RawMessage(`{"input":{"foo":"bar"}}`), StartedAt: now, CreatedAt: now, UpdatedAt: now, }) - store := &stopPoolOnCreateStorage{testExecutionStorage: base, pool: pool} + reqCtx, cancel := context.WithCancel(context.Background()) + store := &stopPoolOnCreateStorage{testExecutionStorage: base, pool: pool, cancel: cancel} router := gin.New() router.POST("/api/v1/executions/:execution_id/restart", RestartExecutionHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) - req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)) + req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)).WithContext(reqCtx) req.Header.Set("Content-Type", "application/json") resp := httptest.NewRecorder() router.ServeHTTP(resp, req) @@ -318,6 +319,12 @@ func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { require.NotNil(t, restarted) require.Equal(t, types.ExecutionStatusFailed, restarted.Status) require.NotNil(t, restarted.StatusReason) - require.Equal(t, string(ErrorCategoryConcurrencyLimit), *restarted.StatusReason) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *restarted.StatusReason) + workflows, err := base.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Len(t, workflows, 1) + require.Equal(t, string(types.ExecutionStatusFailed), workflows[0].Status) + require.NotNil(t, workflows[0].StatusReason) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *workflows[0].StatusReason) require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) } diff --git a/control-plane/internal/handlers/mcp.go b/control-plane/internal/handlers/mcp.go index a54bd0cf9..e1b92772d 100644 --- a/control-plane/internal/handlers/mcp.go +++ b/control-plane/internal/handlers/mcp.go @@ -531,11 +531,8 @@ func (s *mcpServer) startAsyncRun(ctx context.Context, target string, input map[ job := asyncExecutionJob{controller: controller, plan: *plan} if ok := pool.submitReserved(job); !ok { - if plan.slotHeld { - ReleaseExecutionSlot(plan.target.NodeID) - } + job.terminateForControlPlaneShutdown() queueErr := &executionPreconditionError{code: 503, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} - _ = controller.failExecution(ctx, plan, queueErr, 0, nil) return "", "", queueErr } reserved = false diff --git a/control-plane/internal/handlers/mcp_test.go b/control-plane/internal/handlers/mcp_test.go index 4fa5926ce..07cd27c4d 100644 --- a/control-plane/internal/handlers/mcp_test.go +++ b/control-plane/internal/handlers/mcp_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "net/http" "net/http/httptest" @@ -24,6 +25,33 @@ type mcpTestStore struct { agents []*types.AgentNode } +type stopPoolOnCreateMCPStore struct { + *mcpTestStore + pool *asyncWorkerPool + cancel context.CancelFunc + updateErr error + updated bool +} + +func (s *stopPoolOnCreateMCPStore) CreateExecutionRecord(ctx context.Context, execution *types.Execution) error { + s.pool.mu.Lock() + s.pool.stopped = true + s.pool.mu.Unlock() + err := s.mcpTestStore.CreateExecutionRecord(ctx, execution) + if s.cancel != nil { + s.cancel() + } + return err +} + +func (s *stopPoolOnCreateMCPStore) UpdateExecutionRecord(ctx context.Context, executionID string, update func(*types.Execution) (*types.Execution, error)) (*types.Execution, error) { + s.updated = true + if s.updateErr != nil { + return nil, s.updateErr + } + return s.mcpTestStore.UpdateExecutionRecord(ctx, executionID, update) +} + func newMCPTestStore(agents ...*types.AgentNode) *mcpTestStore { var primary *types.AgentNode if len(agents) > 0 { @@ -362,6 +390,49 @@ func TestMCP_ExecuteReasonerConcurrencyRejectionHasNoPersistence(t *testing.T) { require.Empty(t, execs) } +func TestMCP_ExecuteReasonerPoolStoppedTerminatesPersistedRowsWithCancelledRequest(t *testing.T) { + pool := newAsyncWorkerPool(0, 2) + useAsyncPoolForTest(t, pool) + reqCtx, cancel := context.WithCancel(context.Background()) + base := newMCPTestStore(mcpActiveAgent()) + store := &stopPoolOnCreateMCPStore{mcpTestStore: base, pool: pool, cancel: cancel} + server := &mcpServer{store: store, payloads: services.NewFilePayloadStore(t.TempDir()), timeout: time.Second} + + _, _, err := server.startAsyncRun(reqCtx, "planner.plan", map[string]interface{}{"goal": "ship"}, executionHeaders{}, "", "") + require.Error(t, err) + var admissionErr *executionPreconditionError + require.ErrorAs(t, err, &admissionErr) + require.Equal(t, http.StatusServiceUnavailable, admissionErr.code) + + execs, err := base.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Len(t, execs, 1) + require.Equal(t, types.ExecutionStatusFailed, execs[0].Status) + require.NotNil(t, execs[0].StatusReason) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *execs[0].StatusReason) + workflows, err := base.QueryWorkflowExecutions(context.Background(), types.WorkflowExecutionFilters{}) + require.NoError(t, err) + require.Len(t, workflows, 1) + require.Equal(t, string(types.ExecutionStatusFailed), workflows[0].Status) + require.NotNil(t, workflows[0].StatusReason) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *workflows[0].StatusReason) +} + +func TestMCP_ExecuteReasonerPoolStoppedExercisesPersistenceFailure(t *testing.T) { + pool := newAsyncWorkerPool(0, 2) + useAsyncPoolForTest(t, pool) + base := newMCPTestStore(mcpActiveAgent()) + store := &stopPoolOnCreateMCPStore{mcpTestStore: base, pool: pool, updateErr: errors.New("update failed")} + server := &mcpServer{store: store, payloads: services.NewFilePayloadStore(t.TempDir()), timeout: time.Second} + + _, _, err := server.startAsyncRun(context.Background(), "planner.plan", map[string]interface{}{"goal": "ship"}, executionHeaders{}, "", "") + require.Error(t, err) + var admissionErr *executionPreconditionError + require.ErrorAs(t, err, &admissionErr) + require.Equal(t, http.StatusServiceUnavailable, admissionErr.code) + require.True(t, store.updated) +} + func TestMCP_ExecuteReasonerAuthorizesAndBindsRunToVerifiedCaller(t *testing.T) { store := newMCPTestStore(mcpActiveAgent()) var gotCaller, gotTarget string From 56076955eed3178147ede64008177031baac0ba8 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 14:46:17 -0400 Subject: [PATCH 6/8] style(control-plane): gofmt the admission fix Co-Authored-By: Claude Fable 5 --- .../coverage_handlers_90_additional_test.go | 18 +++++++++--------- .../handlers/coverage_raise_88_test.go | 18 ++++++++++-------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/control-plane/internal/handlers/coverage_handlers_90_additional_test.go b/control-plane/internal/handlers/coverage_handlers_90_additional_test.go index b2d4fe310..85497e5dc 100644 --- a/control-plane/internal/handlers/coverage_handlers_90_additional_test.go +++ b/control-plane/internal/handlers/coverage_handlers_90_additional_test.go @@ -142,20 +142,20 @@ func TestExecutionController_CompletionAndFailureCoverage(t *testing.T) { now := time.Now().UTC() require.NoError(t, store.CreateExecutionRecord(context.Background(), &types.Execution{ - ExecutionID: "exec-success", - RunID: "run-success", - AgentNodeID: agent.ID, - Status: types.ExecutionStatusRunning, + ExecutionID: "exec-success", + RunID: "run-success", + AgentNodeID: agent.ID, + Status: types.ExecutionStatusRunning, InputPayload: []byte(`{"prompt":"hello"}`), - CreatedAt: now, - StartedAt: now, - UpdatedAt: now, + CreatedAt: now, + StartedAt: now, + UpdatedAt: now, })) plan := &preparedExecution{ exec: &types.Execution{ - ExecutionID: "exec-success", - RunID: "run-success", + ExecutionID: "exec-success", + RunID: "run-success", InputPayload: []byte(`{"prompt":"hello"}`), }, agent: agent, diff --git a/control-plane/internal/handlers/coverage_raise_88_test.go b/control-plane/internal/handlers/coverage_raise_88_test.go index 5ab8b3b9c..31a5c4c8d 100644 --- a/control-plane/internal/handlers/coverage_raise_88_test.go +++ b/control-plane/internal/handlers/coverage_raise_88_test.go @@ -51,10 +51,10 @@ type statusManagerErrorStore struct { type registerCoverageStore struct { *nodeRESTStorageStub - versioned map[string]*types.AgentNode - deleteCalls []string - updateLifeErr error - registerErr error + versioned map[string]*types.AgentNode + deleteCalls []string + updateLifeErr error + registerErr error } func (s *statusManagerErrorStore) GetAgent(ctx context.Context, id string) (*types.AgentNode, error) { @@ -258,10 +258,12 @@ func TestExecuteReasonerAndSkillHandlers_TransportFailures(t *testing.T) { newHandle func(*reasonerHandlerStorage) gin.HandlerFunc }{ { - name: "reasoner", - route: "/reasoners/:reasoner_id", - target: "/reasoners/node-1.ping", - newStore: func() *reasonerHandlerStorage { return newReasonerHandlerStorage(newReasonerAgent("http://127.0.0.1:1")) }, + name: "reasoner", + route: "/reasoners/:reasoner_id", + target: "/reasoners/node-1.ping", + newStore: func() *reasonerHandlerStorage { + return newReasonerHandlerStorage(newReasonerAgent("http://127.0.0.1:1")) + }, newHandle: func(s *reasonerHandlerStorage) gin.HandlerFunc { return ExecuteReasonerHandler(s) }, }, { From 4a750825cb2dd6cdf339ad274cd4c314f5286bc6 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 17:40:47 -0400 Subject: [PATCH 7/8] fix(control-plane): balance admission ownership through shutdown --- control-plane/internal/handlers/execute.go | 27 +++- .../internal/handlers/execute_async_test.go | 137 ++++++++++++++++++ .../internal/handlers/execute_helpers.go | 108 ++++++++------ .../internal/handlers/execute_lifecycle.go | 13 ++ .../internal/handlers/execute_prepare.go | 13 +- .../internal/handlers/execute_restart.go | 23 ++- .../handlers/execute_sync_admission_test.go | 67 ++++++++- .../internal/handlers/execution_guards.go | 9 ++ control-plane/internal/handlers/mcp.go | 48 +++++- control-plane/internal/handlers/mcp_test.go | 74 ++++++++++ .../internal/services/llm_health_monitor.go | 74 ++++++++-- .../services/llm_health_monitor_retry_test.go | 32 ++-- docs/api/EXECUTE.md | 4 +- docs/api/EXECUTION_RESTART.md | 2 +- 14 files changed, 539 insertions(+), 92 deletions(-) diff --git a/control-plane/internal/handlers/execute.go b/control-plane/internal/handlers/execute.go index 1bf3d0b2d..4d30c0b9b 100644 --- a/control-plane/internal/handlers/execute.go +++ b/control-plane/internal/handlers/execute.go @@ -266,9 +266,7 @@ func (c *executionController) handleSync(ctx *gin.Context) { return } plan.executionMode = "sync" - if plan.slotHeld { - defer ReleaseExecutionSlot(plan.target.NodeID) - } + defer plan.releaseSlot() if plan.replayHit != nil { if err := c.completeReplayHit(reqCtx, plan); err != nil { @@ -433,7 +431,11 @@ func (c *executionController) handleAsync(ctx *gin.Context) { // is released when the worker job returns (early on an agent HTTP 202 ACK). // Workers plus queue capacity bound admitted work; a paused execution pins a // worker and reservation for up to 24 hours. - if !pool.reserve() { + if reserved, stopped := pool.reserveForAdmission(); !reserved { + if stopped { + writeExecutionError(ctx, newControlPlaneShutdownError("async execution queue stopped; retry later")) + return + } writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue is full; retry later") return } @@ -450,6 +452,10 @@ func (c *executionController) handleAsync(ctx *gin.Context) { return } plan.executionMode = "async" + // Preparation owns the slot until it is deliberately transferred to the + // worker job below. This also protects the recovered-panic interval between + // persistence and submission. + defer plan.releaseSlot() if plan.replayHit != nil { if err := c.completeReplayHit(reqCtx, plan); err != nil { @@ -490,6 +496,13 @@ func (c *executionController) handleAsync(ctx *gin.Context) { controller: c, plan: *plan, } + plan.slotHeld = false // ownership transferred to job + submitted := false + defer func() { + if !submitted { + job.plan.releaseSlot() + } + }() if ok := pool.submitReserved(job); !ok { // The pool only refuses a reserved submission once it has stopped, i.e. @@ -497,10 +510,12 @@ func (c *executionController) handleAsync(ctx *gin.Context) { // context: the request context is very likely being cancelled by the // same shutdown, and a cancelled write would strand this row in // "running" — exactly the orphan this branch exists to prevent. - job.terminateForControlPlaneShutdown() - writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue stopped; retry later") + shutdownErr := newControlPlaneShutdownError("async execution queue stopped; retry later") + job.terminateForControlPlaneShutdown(shutdownErr) + writeExecutionError(ctx, shutdownErr) return } + submitted = true reserved = false createdAt := plan.exec.CreatedAt.UTC().Format(time.RFC3339) diff --git a/control-plane/internal/handlers/execute_async_test.go b/control-plane/internal/handlers/execute_async_test.go index 7a0775009..eab76a5f0 100644 --- a/control-plane/internal/handlers/execute_async_test.go +++ b/control-plane/internal/handlers/execute_async_test.go @@ -59,6 +59,7 @@ func TestExecuteAsyncHandler_PoolStoppedTerminatesPersistedRow(t *testing.T) { require.Equal(t, "1", resp.Header().Get("Retry-After")) var body map[string]any require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), body["error_category"]) require.Equal(t, float64(1), body["retry_after"]) records, err := baseStore.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) require.NoError(t, err) @@ -71,15 +72,69 @@ func TestExecuteAsyncHandler_PoolStoppedTerminatesPersistedRow(t *testing.T) { require.Len(t, workflows, 1) require.NotNil(t, workflows[0].StatusReason) require.Equal(t, "control_plane_shutdown", *workflows[0].StatusReason) + require.Equal(t, string(records[0].Status), workflows[0].Status) require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) } +func TestExecuteAsyncHandler_AlreadyStoppedPoolReturnsShutdownWithoutPersistence(t *testing.T) { + pool := newAsyncWorkerPool(0, 2) + pool.mu.Lock() + pool.stopped = true + pool.mu.Unlock() + useAsyncPoolForTest(t, pool) + store := newTestExecutionStorage(testRestartAgent("http://agent.example")) + router := gin.New() + router.POST("/api/v1/execute/async/:target", ExecuteAsyncHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/async/node-1.reasoner-a", strings.NewReader(`{"input":{}}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + require.Equal(t, http.StatusServiceUnavailable, resp.Code) + require.Equal(t, "1", resp.Header().Get("Retry-After")) + var body map[string]interface{} + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), body["error_category"]) + records, err := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, err) + require.Empty(t, records) +} + type stopPoolOnCreateStorage struct { *testExecutionStorage pool *asyncWorkerPool cancel context.CancelFunc } +type cancelBeforeShutdownUpdateStore struct { + *testExecutionStorage + once sync.Once + interleaveErr error +} + +func (s *cancelBeforeShutdownUpdateStore) UpdateExecutionRecord(ctx context.Context, executionID string, update func(*types.Execution) (*types.Execution, error)) (*types.Execution, error) { + s.once.Do(func() { + reason := "cancelled_by_user" + _, s.interleaveErr = s.testExecutionStorage.UpdateExecutionRecord(ctx, executionID, func(current *types.Execution) (*types.Execution, error) { + current.Status = types.ExecutionStatusCancelled + current.StatusReason = &reason + return current, nil + }) + if s.interleaveErr != nil { + return + } + s.interleaveErr = s.testExecutionStorage.UpdateWorkflowExecution(ctx, executionID, func(current *types.WorkflowExecution) (*types.WorkflowExecution, error) { + current.Status = string(types.ExecutionStatusCancelled) + current.StatusReason = &reason + return current, nil + }) + }) + if s.interleaveErr != nil { + return nil, s.interleaveErr + } + return s.testExecutionStorage.UpdateExecutionRecord(ctx, executionID, update) +} + func (s *stopPoolOnCreateStorage) CreateExecutionRecord(ctx context.Context, execution *types.Execution) error { s.pool.mu.Lock() s.pool.stopped = true @@ -91,6 +146,88 @@ func (s *stopPoolOnCreateStorage) CreateExecutionRecord(ctx context.Context, exe return err } +func TestAsyncShutdownTerminalizationPreservesCancellationInterleaving(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 2} + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + + base := newTestExecutionStorage(testRestartAgent("http://agent.example")) + now := time.Now().UTC() + exec := &types.Execution{ + ExecutionID: "exec-cancel-race", RunID: "run-cancel-race", AgentNodeID: "node-1", NodeID: "node-1", + ReasonerID: "reasoner-a", Status: types.ExecutionStatusRunning, + CreatedAt: now, StartedAt: now, UpdatedAt: now, + } + require.NoError(t, base.CreateExecutionRecord(context.Background(), exec)) + require.NoError(t, base.StoreWorkflowExecution(context.Background(), &types.WorkflowExecution{ + ExecutionID: exec.ExecutionID, WorkflowID: exec.RunID, RunID: &exec.RunID, + AgentNodeID: "node-1", ReasonerID: "reasoner-a", Status: string(types.ExecutionStatusRunning), + CreatedAt: now, StartedAt: now, UpdatedAt: now, + })) + store := &cancelBeforeShutdownUpdateStore{testExecutionStorage: base} + target, err := parseTarget("node-1.reasoner-a") + require.NoError(t, err) + job := asyncExecutionJob{ + controller: newExecutionController(store, nil, nil, time.Second, ""), + plan: preparedExecution{ + exec: exec, target: target, slotHeld: true, + }, + } + + job.terminateForControlPlaneShutdown(newControlPlaneShutdownError("control plane stopped")) + + stored, err := base.GetExecutionRecord(context.Background(), exec.ExecutionID) + require.NoError(t, err) + workflow, err := base.GetWorkflowExecution(context.Background(), exec.ExecutionID) + require.NoError(t, err) + require.Equal(t, types.ExecutionStatusCancelled, stored.Status) + require.Equal(t, string(types.ExecutionStatusCancelled), workflow.Status) + require.Equal(t, "cancelled_by_user", *stored.StatusReason) + require.Equal(t, "cancelled_by_user", *workflow.StatusReason) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) +} + +func TestAsyncForcedWorkerCancellationReleasesOnlyOwnedSlot(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 2} + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + require.NoError(t, concurrencyLimiter.Acquire("node-1")) + + agent := testRestartAgent("http://agent.example") + store := newTestExecutionStorage(agent) + now := time.Now().UTC() + exec := &types.Execution{ + ExecutionID: "exec-forced-stop", RunID: "run-forced-stop", AgentNodeID: "node-1", NodeID: "node-1", + ReasonerID: "reasoner-a", Status: types.ExecutionStatusRunning, + CreatedAt: now, StartedAt: now, UpdatedAt: now, + } + require.NoError(t, store.CreateExecutionRecord(context.Background(), exec)) + require.NoError(t, store.StoreWorkflowExecution(context.Background(), &types.WorkflowExecution{ + ExecutionID: exec.ExecutionID, WorkflowID: exec.RunID, RunID: &exec.RunID, + AgentNodeID: "node-1", ReasonerID: "reasoner-a", Status: string(types.ExecutionStatusRunning), + CreatedAt: now, StartedAt: now, UpdatedAt: now, + })) + target, err := parseTarget("node-1.reasoner-a") + require.NoError(t, err) + job := asyncExecutionJob{ + controller: newExecutionController(store, nil, nil, time.Second, ""), + plan: preparedExecution{ + exec: exec, target: target, agent: agent, requestBody: []byte(`{}`), slotHeld: true, + }, + } + workerCtx, cancel := context.WithCancel(context.Background()) + cancel() + job.processWithContext(workerCtx) + + // This job owned one of two live slots. Forced shutdown must not consume the + // other execution's count through failForControlPlaneShutdown plus defer. + require.EqualValues(t, 1, concurrencyLimiter.GetRunningCount("node-1")) + ReleaseExecutionSlot("node-1") + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) +} + func TestExecuteAsyncHandler_QueueSaturation(t *testing.T) { gin.SetMode(gin.TestMode) useAsyncPoolForTest(t, newAsyncWorkerPool(1, 1)) diff --git a/control-plane/internal/handlers/execute_helpers.go b/control-plane/internal/handlers/execute_helpers.go index e469302fb..9b14fbc86 100644 --- a/control-plane/internal/handlers/execute_helpers.go +++ b/control-plane/internal/handlers/execute_helpers.go @@ -562,28 +562,9 @@ func writeExecutionError(ctx *gin.Context, err error) { var pe *executionPreconditionError if errors.As(err, &pe) { - body := gin.H{ - "error": pe.Error(), - "error_category": string(pe.Category()), - } - // When a stable machine code is set, promote it to `error` and move - // the human-readable text to `message` — matching the contract used - // by reasoners.go / skills.go / permission middleware. - if code := pe.ErrorCode(); code != "" { - body["error"] = code - body["message"] = pe.Error() - } - retryAfter := pe.retryAfter - if retryAfter <= 0 { - retryAfter = map[ErrorCategory]int{ - ErrorCategoryConcurrencyLimit: 1, - ErrorCategoryNodeUnavailable: 1, - ErrorCategoryLLMUnavailable: 30, - }[pe.Category()] - } + body, retryAfter := renderExecutionPreconditionError(pe) if retryAfter > 0 { ctx.Header("Retry-After", strconv.Itoa(retryAfter)) - body["retry_after"] = retryAfter } ctx.JSON(pe.HTTPStatusCode(), body) return @@ -601,6 +582,33 @@ func writeExecutionError(ctx *gin.Context, err error) { }) } +func renderExecutionPreconditionError(pe *executionPreconditionError) (gin.H, int) { + body := gin.H{ + "error": pe.Error(), + "error_category": string(pe.Category()), + } + // When a stable machine code is set, promote it to `error` and move the + // human-readable text to `message` — matching the contract used by sibling + // handlers. + if code := pe.ErrorCode(); code != "" { + body["error"] = code + body["message"] = pe.Error() + } + retryAfter := pe.retryAfter + if retryAfter <= 0 { + retryAfter = map[ErrorCategory]int{ + ErrorCategoryConcurrencyLimit: 1, + ErrorCategoryNodeUnavailable: 1, + ErrorCategoryLLMUnavailable: 30, + ErrorCategoryControlPlaneShutdown: 1, + }[pe.Category()] + } + if retryAfter > 0 { + body["retry_after"] = retryAfter + } + return body, retryAfter +} + // classifyExecutionError determines the error category from any execution error. func classifyExecutionError(err error) ErrorCategory { if err == nil { @@ -848,9 +856,7 @@ func (j asyncExecutionJob) process() { func (j asyncExecutionJob) processWithContext(workerCtx context.Context) { // Release the per-agent concurrency slot when this job finishes - if j.plan.slotHeld && j.plan.target != nil { - defer ReleaseExecutionSlot(j.plan.target.NodeID) - } + defer j.plan.releaseSlot() // Use a bounded context so that paused executions do not block goroutines // indefinitely if the resume/cancel event is never delivered (e.g. event bus @@ -880,7 +886,7 @@ func (j asyncExecutionJob) processWithContext(workerCtx context.Context) { resultBody, elapsed, asyncAccepted, callErr := j.controller.callAgent(bgCtx, &j.plan) if workerCtx.Err() != nil { persistCtx, persistCancel := shutdownPersistenceContext() - j.failForControlPlaneShutdown(persistCtx) + j.failForControlPlaneShutdown(persistCtx, newControlPlaneShutdownError("execution was interrupted because the control plane shut down")) persistCancel() return } @@ -951,7 +957,7 @@ func newAsyncWorkerPool(workerCount, queueCapacity int) *asyncWorkerPool { pool.mu.RUnlock() if stopped { persistCtx, cancel := shutdownPersistenceContext() - job.failForControlPlaneShutdown(persistCtx) + job.terminateForControlPlaneShutdownWithContext(persistCtx, newControlPlaneShutdownError("execution was not started before the control plane shut down")) cancel() } else { job.processWithContext(pool.workerCtx) @@ -981,16 +987,23 @@ func (p *asyncWorkerPool) submit(job asyncExecutionJob) bool { } func (p *asyncWorkerPool) reserve() bool { + reserved, _ := p.reserveForAdmission() + return reserved +} + +// reserveForAdmission distinguishes saturation from a pool that has stopped so +// callers can return the stable control_plane_shutdown contract. +func (p *asyncWorkerPool) reserveForAdmission() (reserved, stopped bool) { p.mu.RLock() defer p.mu.RUnlock() if p.stopped { - return false + return false, true } select { case p.reservations <- struct{}{}: - return true + return true, false default: - return false + return false, false } } @@ -1045,7 +1058,7 @@ func (p *asyncWorkerPool) Stop(ctx context.Context) { defer cancel() for job := range p.queue { p.releaseReservation() - job.failForControlPlaneShutdown(persistCtx) + job.terminateForControlPlaneShutdownWithContext(persistCtx, newControlPlaneShutdownError("execution was not started before the control plane shut down")) p.jobs.Done() } select { @@ -1058,38 +1071,45 @@ func shutdownPersistenceContext() (context.Context, context.CancelFunc) { return context.WithTimeout(context.Background(), 5*time.Second) } -func (j asyncExecutionJob) terminateForControlPlaneShutdown() { +func (j *asyncExecutionJob) terminateForControlPlaneShutdown(shutdownErr *executionPreconditionError) { persistCtx, cancel := shutdownPersistenceContext() defer cancel() - j.failForControlPlaneShutdown(persistCtx) + j.terminateForControlPlaneShutdownWithContext(persistCtx, shutdownErr) +} + +func (j *asyncExecutionJob) terminateForControlPlaneShutdownWithContext(ctx context.Context, shutdownErr *executionPreconditionError) { + defer j.plan.releaseSlot() + j.failForControlPlaneShutdown(ctx, shutdownErr) } -func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context) { +func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context, shutdownErr *executionPreconditionError) { nodeID := "" if j.plan.target != nil { nodeID = j.plan.target.NodeID } - if j.plan.slotHeld && j.plan.target != nil { - ReleaseExecutionSlot(nodeID) - } - shutdownErr := &executionPreconditionError{ - message: "execution was not started before the control plane shut down", - category: ErrorCategoryControlPlaneShutdown, + if shutdownErr == nil { + shutdownErr = newControlPlaneShutdownError("execution was not started before the control plane shut down") } if err := j.controller.failExecution(ctx, &j.plan, shutdownErr, 0, nil); err != nil { logger.Logger.Warn().Err(err).Str("node_id", nodeID).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to terminate queued execution during shutdown") return } - reason := string(ErrorCategoryControlPlaneShutdown) + // failExecution deliberately preserves cancellation/waiting. Re-read the + // execution after that atomic update and mirror the state that won the race + // instead of blindly overwriting workflow_executions with failed. + updatedExec, err := j.controller.store.GetExecutionRecord(ctx, j.plan.exec.ExecutionID) + if err != nil || updatedExec == nil { + logger.Logger.Warn().Err(err).Str("node_id", nodeID).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to load terminal execution during shutdown reconciliation") + return + } if err := j.controller.store.UpdateWorkflowExecution(ctx, j.plan.exec.ExecutionID, func(current *types.WorkflowExecution) (*types.WorkflowExecution, error) { if current == nil { return nil, fmt.Errorf("workflow execution %s not found", j.plan.exec.ExecutionID) } - now := time.Now().UTC() - current.Status = string(types.ExecutionStatusFailed) - current.StatusReason = &reason - current.CompletedAt = &now - current.UpdatedAt = now + current.Status = string(updatedExec.Status) + current.StatusReason = updatedExec.StatusReason + current.CompletedAt = updatedExec.CompletedAt + current.UpdatedAt = updatedExec.UpdatedAt return current, nil }); err != nil { logger.Logger.Warn().Err(err).Str("node_id", nodeID).Str("execution_id", j.plan.exec.ExecutionID).Msg("failed to record shutdown reason on queued workflow execution") diff --git a/control-plane/internal/handlers/execute_lifecycle.go b/control-plane/internal/handlers/execute_lifecycle.go index 29ffab754..bb053f55b 100644 --- a/control-plane/internal/handlers/execute_lifecycle.go +++ b/control-plane/internal/handlers/execute_lifecycle.go @@ -211,6 +211,19 @@ type preparedExecution struct { replayHit *replayHit } +// releaseSlot releases this plan's admission slot at most once. A plan owns +// its slot from successful preparation until a synchronous handler returns or +// ownership is transferred to an asyncExecutionJob. +func (p *preparedExecution) releaseSlot() { + if p == nil || !p.slotHeld { + return + } + p.slotHeld = false + if p.target != nil { + ReleaseExecutionSlot(p.target.NodeID) + } +} + func (c *executionController) callAgent(ctx context.Context, plan *preparedExecution) ([]byte, time.Duration, bool, error) { start := time.Now() diff --git a/control-plane/internal/handlers/execute_prepare.go b/control-plane/internal/handlers/execute_prepare.go index b98130e39..9cacea2c4 100644 --- a/control-plane/internal/handlers/execute_prepare.go +++ b/control-plane/internal/handlers/execute_prepare.go @@ -41,7 +41,7 @@ func (c *executionController) prepareExecutionWithAdmission(ctx context.Context, ) } -func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context.Context, targetParam string, req ExecuteRequest, headers executionHeaders, callerDID, targetDID string, acquireSlot bool) (_ *preparedExecution, retErr error) { +func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context.Context, targetParam string, req ExecuteRequest, headers executionHeaders, callerDID, targetDID string, acquireSlot bool) (*preparedExecution, error) { target, err := parseTarget(targetParam) if err != nil { return nil, fmt.Errorf("invalid target: %w", err) @@ -165,6 +165,7 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context llmEndpoint := extractRequestedLLMEndpoint(req) slotAcquired := false + slotTransferred := false if acquireSlot && hit == nil { if err := CheckExecutionPreconditions(target.NodeID, llmEndpoint); err != nil { logger.Logger.Warn(). @@ -176,7 +177,9 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context } slotAcquired = true defer func() { - if retErr != nil && slotAcquired { + // Gin recovers handler panics. Do not strand a slot when persistence or + // another preparation dependency panics before a plan takes ownership. + if slotAcquired && !slotTransferred { ReleaseExecutionSlot(target.NodeID) } }() @@ -258,7 +261,7 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context c.ensureWorkflowExecutionRecord(ctx, exec, target, storedPayload) - return &preparedExecution{ + plan := &preparedExecution{ exec: exec, requestBody: agentPayloadBytes, agent: agent, @@ -275,7 +278,9 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context replayBeforeExecutionID: headers.replayBeforeExecutionID, replayMode: headers.replayMode, replayHit: hit, - }, nil + } + slotTransferred = true + return plan, nil } // findReplayHit returns a previously-succeeded child output to reuse for the diff --git a/control-plane/internal/handlers/execute_restart.go b/control-plane/internal/handlers/execute_restart.go index a4bae714f..28291819c 100644 --- a/control-plane/internal/handlers/execute_restart.go +++ b/control-plane/internal/handlers/execute_restart.go @@ -152,7 +152,11 @@ func (c *executionController) handleRestart(ctx *gin.Context) { target := fmt.Sprintf("%s.%s", restartExec.NodeID, restartExec.ReasonerID) pool := getAsyncWorkerPool() - if !pool.reserve() { + if reserved, stopped := pool.reserveForAdmission(); !reserved { + if stopped { + writeExecutionError(ctx, newControlPlaneShutdownError("async execution queue stopped; retry later")) + return + } writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue is full; retry later") return } @@ -172,6 +176,9 @@ func (c *executionController) handleRestart(ctx *gin.Context) { writeExecutionError(ctx, err) return } + // Keep ownership in the handler until the job copy is ready. Gin recovery + // can then release the slot if metadata/event publication panics. + defer plan.releaseSlot() kind := "restart" if req.Fork || req.Input != nil || req.Context != nil { @@ -185,12 +192,20 @@ func (c *executionController) handleRestart(ctx *gin.Context) { controller: c, plan: *plan, } + plan.slotHeld = false // ownership transferred to job + submitted := false + defer func() { + if !submitted { + job.plan.releaseSlot() + } + }() if ok := pool.submitReserved(job); !ok { - job.terminateForControlPlaneShutdown() - queueErr := &executionPreconditionError{code: http.StatusServiceUnavailable, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} - writeExecutionError(ctx, queueErr) + shutdownErr := newControlPlaneShutdownError("async execution queue stopped; retry later") + job.terminateForControlPlaneShutdown(shutdownErr) + writeExecutionError(ctx, shutdownErr) return } + submitted = true reserved = false createdAt := plan.exec.CreatedAt.UTC().Format(time.RFC3339) diff --git a/control-plane/internal/handlers/execute_sync_admission_test.go b/control-plane/internal/handlers/execute_sync_admission_test.go index 5eb6edab5..f9cc0fe9a 100644 --- a/control-plane/internal/handlers/execute_sync_admission_test.go +++ b/control-plane/internal/handlers/execute_sync_admission_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "path/filepath" @@ -19,6 +20,62 @@ import ( "github.com/stretchr/testify/require" ) +type panicOnCreateExecutionStore struct { + *testExecutionStorage +} + +func (s *panicOnCreateExecutionStore) CreateExecutionRecord(context.Context, *types.Execution) error { + panic("injected execution persistence panic") +} + +func TestExecuteAdmission_RecoveredPersistencePanicReleasesSlot(t *testing.T) { + for _, test := range []struct { + name string + run func(t *testing.T, store *panicOnCreateExecutionStore) *httptest.ResponseRecorder + }{ + { + name: "sync", + run: func(t *testing.T, store *panicOnCreateExecutionStore) *httptest.ResponseRecorder { + router := gin.New() + router.Use(gin.CustomRecoveryWithWriter(io.Discard, func(c *gin.Context, _ interface{}) { c.AbortWithStatus(http.StatusInternalServerError) })) + router.POST("/api/v1/execute/:target", ExecuteHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/execute/node-1.reasoner-a", strings.NewReader(`{"input":{}}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + return resp + }, + }, + { + name: "restart", + run: func(t *testing.T, store *panicOnCreateExecutionStore) *httptest.ResponseRecorder { + useAsyncPoolForTest(t, newAsyncWorkerPool(0, 2)) + now := time.Now().UTC() + source := &types.Execution{ExecutionID: "source", RunID: "old-run", AgentNodeID: "node-1", NodeID: "node-1", ReasonerID: "reasoner-a", Status: types.ExecutionStatusFailed, InputPayload: json.RawMessage(`{"input":{}}`), StartedAt: now, CreatedAt: now, UpdatedAt: now} + require.NoError(t, store.testExecutionStorage.CreateExecutionRecord(context.Background(), source)) + router := gin.New() + router.Use(gin.CustomRecoveryWithWriter(io.Discard, func(c *gin.Context, _ interface{}) { c.AbortWithStatus(http.StatusInternalServerError) })) + router.POST("/api/v1/executions/:execution_id/restart", RestartExecutionHandler(store, services.NewFilePayloadStore(t.TempDir()), nil, time.Second, "")) + req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/source/restart", strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + return resp + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + store := &panicOnCreateExecutionStore{testExecutionStorage: newTestExecutionStorage(testRestartAgent("http://agent.example"))} + resp := test.run(t, store) + require.Equal(t, http.StatusInternalServerError, resp.Code) + require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) + }) + } +} + func TestExecuteHandler_ConcurrencyRejectionHasNoPersistence(t *testing.T) { gin.SetMode(gin.TestMode) oldLimiter := concurrencyLimiter @@ -179,6 +236,7 @@ func TestWriteExecutionError_RetryAfterPerCategory(t *testing.T) { {ErrorCategoryConcurrencyLimit, 1}, {ErrorCategoryNodeUnavailable, 1}, {ErrorCategoryLLMUnavailable, 17}, + {ErrorCategoryControlPlaneShutdown, 1}, } for _, test := range tests { t.Run(string(test.category), func(t *testing.T) { @@ -271,10 +329,8 @@ func TestRestartHandler_QueueFullCarriesRetryAfter(t *testing.T) { } // A restart admitted past reserve() but refused by submitReserved (the pool -// stopped in between) must persist a status_reason that matches the -// error_category it answers with — it used to answer concurrency_limit while -// recording internal_error, because the queue error was untyped. -func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { +// stopped in between) reports and persists one stable shutdown category. +func TestRestartHandler_PoolStoppedReturnsConsistentShutdownContract(t *testing.T) { gin.SetMode(gin.TestMode) pool := newAsyncWorkerPool(0, 4) useAsyncPoolForTest(t, pool) @@ -304,7 +360,7 @@ func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { require.Equal(t, "1", resp.Header().Get("Retry-After")) var body map[string]any require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &body)) - require.Equal(t, "concurrency_limit", body["error_category"]) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), body["error_category"]) require.Equal(t, float64(1), body["retry_after"]) records, err := base.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) @@ -326,5 +382,6 @@ func TestRestartHandler_PoolStoppedPersistsMatchingStatusReason(t *testing.T) { require.Equal(t, string(types.ExecutionStatusFailed), workflows[0].Status) require.NotNil(t, workflows[0].StatusReason) require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *workflows[0].StatusReason) + require.Equal(t, string(restarted.Status), workflows[0].Status) require.Zero(t, concurrencyLimiter.GetRunningCount("node-1")) } diff --git a/control-plane/internal/handlers/execution_guards.go b/control-plane/internal/handlers/execution_guards.go index 0a347601e..10b98e12e 100644 --- a/control-plane/internal/handlers/execution_guards.go +++ b/control-plane/internal/handlers/execution_guards.go @@ -113,6 +113,15 @@ func newLLMUnavailableError(message, lastErr string, retryAfter int) error { } } +func newControlPlaneShutdownError(message string) *executionPreconditionError { + return &executionPreconditionError{ + code: 503, + message: message, + category: ErrorCategoryControlPlaneShutdown, + retryAfter: 1, + } +} + // ReleaseExecutionSlot releases the concurrency slot for the given agent. // Safe to call even if concurrency limiting is disabled. func ReleaseExecutionSlot(agentNodeID string) { diff --git a/control-plane/internal/handlers/mcp.go b/control-plane/internal/handlers/mcp.go index e1b92772d..8df32c28b 100644 --- a/control-plane/internal/handlers/mcp.go +++ b/control-plane/internal/handlers/mcp.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -416,6 +417,17 @@ func (s *mcpServer) toolExecuteReasoner(c *gin.Context, rawArgs json.RawMessage) } runID, execID, err := s.startAsyncRun(ctx, target, input, headers, callerDID, targetDID) if err != nil { + var admissionErr *executionPreconditionError + if errors.As(err, &admissionErr) { + body, retryAfter := renderExecutionPreconditionError(admissionErr) + // Preserve the legacy human-readable MCP field while adding the stable + // category/retry contract for programmatic clients. + body["text"] = "failed to start execution: " + admissionErr.Error() + if retryAfter > 0 { + c.Header("Retry-After", fmt.Sprint(retryAfter)) + } + return mcpToolErrorValue(body), nil + } return mcpToolError("failed to start execution: " + err.Error()), nil } @@ -512,7 +524,10 @@ func (s *mcpServer) toolWaitRun(c *gin.Context, rawArgs json.RawMessage) (map[st func (s *mcpServer) startAsyncRun(ctx context.Context, target string, input map[string]interface{}, headers executionHeaders, callerDID, targetDID string) (runID, execID string, err error) { controller := newExecutionController(s.store, s.payloads, s.webhooks, s.timeout, s.internalToken) pool := getAsyncWorkerPool() - if !pool.reserve() { + if reserved, stopped := pool.reserveForAdmission(); !reserved { + if stopped { + return "", "", newControlPlaneShutdownError("async execution queue stopped; retry later") + } return "", "", &executionPreconditionError{code: 503, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} } reserved := true @@ -526,15 +541,24 @@ func (s *mcpServer) startAsyncRun(ctx context.Context, target string, input map[ if err != nil { return "", "", err } + defer plan.releaseSlot() controller.publishExecutionStartedEvent(plan) job := asyncExecutionJob{controller: controller, plan: *plan} + plan.slotHeld = false // ownership transferred to job + submitted := false + defer func() { + if !submitted { + job.plan.releaseSlot() + } + }() if ok := pool.submitReserved(job); !ok { - job.terminateForControlPlaneShutdown() - queueErr := &executionPreconditionError{code: 503, message: "async execution queue is full; retry later", category: ErrorCategoryConcurrencyLimit} - return "", "", queueErr + shutdownErr := newControlPlaneShutdownError("async execution queue stopped; retry later") + job.terminateForControlPlaneShutdown(shutdownErr) + return "", "", shutdownErr } + submitted = true reserved = false return plan.exec.RunID, plan.exec.ExecutionID, nil @@ -710,6 +734,22 @@ func mcpToolError(msg string) map[string]interface{} { } } +// mcpToolErrorValue preserves a structured execution rejection inside the MCP +// tool result while retaining the standard HTTP-200 JSON-RPC transport. +func mcpToolErrorValue(v interface{}) map[string]interface{} { + b, err := json.Marshal(v) + if err != nil { + return mcpToolError("failed to encode error: " + err.Error()) + } + return map[string]interface{}{ + "content": []map[string]interface{}{ + {"type": "text", "text": string(b)}, + }, + "structuredContent": v, + "isError": true, + } +} + // mcpToolCatalog returns the tools/list payload: the five AgentField tools with // their JSON Schemas. func mcpToolCatalog() []mcpTool { diff --git a/control-plane/internal/handlers/mcp_test.go b/control-plane/internal/handlers/mcp_test.go index 07cd27c4d..8e1080cde 100644 --- a/control-plane/internal/handlers/mcp_test.go +++ b/control-plane/internal/handlers/mcp_test.go @@ -33,6 +33,14 @@ type stopPoolOnCreateMCPStore struct { updated bool } +type panicOnCreateMCPStore struct { + *mcpTestStore +} + +func (s *panicOnCreateMCPStore) CreateExecutionRecord(context.Context, *types.Execution) error { + panic("injected MCP execution persistence panic") +} + func (s *stopPoolOnCreateMCPStore) CreateExecutionRecord(ctx context.Context, execution *types.Execution) error { s.pool.mu.Lock() s.pool.stopped = true @@ -390,6 +398,39 @@ func TestMCP_ExecuteReasonerConcurrencyRejectionHasNoPersistence(t *testing.T) { require.Empty(t, execs) } +func TestMCP_ExecuteReasonerAlreadyStoppedPoolReturnsShutdownWithoutPersistence(t *testing.T) { + pool := newAsyncWorkerPool(0, 2) + pool.mu.Lock() + pool.stopped = true + pool.mu.Unlock() + useAsyncPoolForTest(t, pool) + store := newMCPTestStore(mcpActiveAgent()) + server := &mcpServer{store: store, payloads: services.NewFilePayloadStore(t.TempDir()), timeout: time.Second} + + _, _, err := server.startAsyncRun(context.Background(), "planner.plan", map[string]interface{}{"goal": "ship"}, executionHeaders{}, "", "") + var admissionErr *executionPreconditionError + require.ErrorAs(t, err, &admissionErr) + require.Equal(t, ErrorCategoryControlPlaneShutdown, admissionErr.category) + require.Equal(t, 1, admissionErr.retryAfter) + execs, queryErr := store.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) + require.NoError(t, queryErr) + require.Empty(t, execs) +} + +func TestMCP_ExecuteReasonerPersistencePanicReleasesSlot(t *testing.T) { + oldLimiter := concurrencyLimiter + concurrencyLimiter = &AgentConcurrencyLimiter{maxPerAgent: 1} + t.Cleanup(func() { concurrencyLimiter = oldLimiter }) + useAsyncPoolForTest(t, newAsyncWorkerPool(0, 2)) + store := &panicOnCreateMCPStore{mcpTestStore: newMCPTestStore(mcpActiveAgent())} + server := &mcpServer{store: store, payloads: services.NewFilePayloadStore(t.TempDir()), timeout: time.Second} + + require.Panics(t, func() { + _, _, _ = server.startAsyncRun(context.Background(), "planner.plan", map[string]interface{}{"goal": "ship"}, executionHeaders{}, "", "") + }) + require.Zero(t, concurrencyLimiter.GetRunningCount("planner")) +} + func TestMCP_ExecuteReasonerPoolStoppedTerminatesPersistedRowsWithCancelledRequest(t *testing.T) { pool := newAsyncWorkerPool(0, 2) useAsyncPoolForTest(t, pool) @@ -403,6 +444,8 @@ func TestMCP_ExecuteReasonerPoolStoppedTerminatesPersistedRowsWithCancelledReque var admissionErr *executionPreconditionError require.ErrorAs(t, err, &admissionErr) require.Equal(t, http.StatusServiceUnavailable, admissionErr.code) + require.Equal(t, ErrorCategoryControlPlaneShutdown, admissionErr.category) + require.Equal(t, 1, admissionErr.retryAfter) execs, err := base.QueryExecutionRecords(context.Background(), types.ExecutionFilter{}) require.NoError(t, err) @@ -416,6 +459,37 @@ func TestMCP_ExecuteReasonerPoolStoppedTerminatesPersistedRowsWithCancelledReque require.Equal(t, string(types.ExecutionStatusFailed), workflows[0].Status) require.NotNil(t, workflows[0].StatusReason) require.Equal(t, string(ErrorCategoryControlPlaneShutdown), *workflows[0].StatusReason) + require.Equal(t, string(execs[0].Status), workflows[0].Status) +} + +func TestMCP_ExecuteReasonerPoolStoppedReturnsStructuredShutdownContract(t *testing.T) { + pool := newAsyncWorkerPool(0, 2) + useAsyncPoolForTest(t, pool) + base := newMCPTestStore(mcpActiveAgent()) + store := &stopPoolOnCreateMCPStore{mcpTestStore: base, pool: pool} + router := newMCPTestRouter(t, store) + + reqBody, err := json.Marshal(map[string]interface{}{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]interface{}{ + "name": "execute_reasoner", + "arguments": map[string]interface{}{ + "target": "planner.plan", + "input": map[string]interface{}{"goal": "ship"}, + }, + }, + }) + require.NoError(t, err) + w := mcpPost(t, router, string(reqBody)) + require.Equal(t, "1", w.Header().Get("Retry-After")) + resp := mcpDecode(t, w) + result := resp["result"].(map[string]interface{}) + require.Equal(t, true, result["isError"]) + structured := result["structuredContent"].(map[string]interface{}) + require.Equal(t, string(ErrorCategoryControlPlaneShutdown), structured["error_category"]) + require.Equal(t, float64(1), structured["retry_after"]) } func TestMCP_ExecuteReasonerPoolStoppedExercisesPersistenceFailure(t *testing.T) { diff --git a/control-plane/internal/services/llm_health_monitor.go b/control-plane/internal/services/llm_health_monitor.go index eee1c5740..33ef1ed10 100644 --- a/control-plane/internal/services/llm_health_monitor.go +++ b/control-plane/internal/services/llm_health_monitor.go @@ -36,18 +36,21 @@ type LLMEndpointStatus struct { TotalFailures int64 `json:"total_failures"` // Circuit breaker internals circuitOpenedAt time.Time + nextProbeAt time.Time halfOpenSuccesses int } // LLMHealthMonitor monitors LLM backend health using circuit breaker pattern. type LLMHealthMonitor struct { - config config.LLMHealthConfig - httpClient *http.Client - endpoints map[string]*LLMEndpointStatus - mu sync.RWMutex - stopCh chan struct{} - stopOnce sync.Once - uiService *UIService + config config.LLMHealthConfig + httpClient *http.Client + endpoints map[string]*LLMEndpointStatus + mu sync.RWMutex + stopCh chan struct{} + stopOnce sync.Once + uiService *UIService + now func() time.Time + nextCheckAt time.Time } // NewLLMHealthMonitor creates a new LLM health monitor. @@ -85,6 +88,7 @@ func NewLLMHealthMonitor(cfg config.LLMHealthConfig, uiService *UIService) *LLMH endpoints: endpoints, stopCh: make(chan struct{}), uiService: uiService, + now: time.Now, } } @@ -117,7 +121,14 @@ func (m *LLMHealthMonitor) RetryAfterSeconds(name string) int { } remaining := recovery if endpoint, ok := m.endpoints[normalizeLLMEndpointName(name)]; ok && endpoint.CircuitState == CircuitOpen && !endpoint.circuitOpenedAt.IsZero() { - remaining = recovery - time.Since(endpoint.circuitOpenedAt) + nextProbeAt := endpoint.nextProbeAt + if nextProbeAt.IsZero() { + // A circuit can transition only on a health-check tick. Without a + // recorded scheduler deadline, use the conservative end of the next + // check interval rather than promising recovery at the raw timeout. + nextProbeAt = endpoint.circuitOpenedAt.Add(recovery + m.config.CheckInterval) + } + remaining = nextProbeAt.Sub(m.currentTime()) } seconds := int(math.Ceil(remaining.Seconds())) if seconds < 1 { @@ -141,6 +152,7 @@ func (m *LLMHealthMonitor) Start() { ticker := time.NewTicker(m.config.CheckInterval) defer ticker.Stop() + m.setNextCheckAt(m.currentTime().Add(m.config.CheckInterval)) // Initial check m.checkAllEndpoints() @@ -150,7 +162,8 @@ func (m *LLMHealthMonitor) Start() { case <-m.stopCh: logger.Logger.Info().Msg("LLM health monitor stopped") return - case <-ticker.C: + case tickAt := <-ticker.C: + m.setNextCheckAt(tickAt.Add(m.config.CheckInterval)) m.checkAllEndpoints() } } @@ -240,8 +253,9 @@ func (m *LLMHealthMonitor) checkEndpoint(epCfg config.LLMEndpoint) { // If circuit is open, check if recovery timeout has elapsed if ep.CircuitState == CircuitOpen { - if time.Since(ep.circuitOpenedAt) >= m.config.RecoveryTimeout { + if m.currentTime().Sub(ep.circuitOpenedAt) >= m.config.RecoveryTimeout { ep.CircuitState = CircuitHalfOpen + ep.nextProbeAt = time.Time{} ep.halfOpenSuccesses = 0 logger.Logger.Info(). Str("endpoint", ep.Name). @@ -288,8 +302,39 @@ func normalizeLLMEndpointName(name string) string { return strings.TrimSpace(strings.ToLower(name)) } +func (m *LLMHealthMonitor) currentTime() time.Time { + if m != nil && m.now != nil { + return m.now() + } + return time.Now() +} + +func (m *LLMHealthMonitor) setNextCheckAt(next time.Time) { + m.mu.Lock() + m.nextCheckAt = next + m.mu.Unlock() +} + +// nextCircuitProbeAtLocked returns the first scheduled check at or after the +// recovery deadline. m.mu must be held by the caller. +func (m *LLMHealthMonitor) nextCircuitProbeAtLocked(openedAt time.Time) time.Time { + deadline := openedAt.Add(m.config.RecoveryTimeout) + interval := m.config.CheckInterval + if interval <= 0 { + interval = 15 * time.Second + } + candidate := m.nextCheckAt + if candidate.IsZero() || !candidate.After(openedAt) { + candidate = openedAt.Add(interval) + } + for candidate.Before(deadline) { + candidate = candidate.Add(interval) + } + return candidate +} + func (m *LLMHealthMonitor) handleSuccess(ep *LLMEndpointStatus) { - ep.LastSuccess = time.Now() + ep.LastSuccess = m.currentTime() ep.ConsecutiveFailures = 0 switch ep.CircuitState { @@ -297,6 +342,7 @@ func (m *LLMHealthMonitor) handleSuccess(ep *LLMEndpointStatus) { ep.halfOpenSuccesses++ if ep.halfOpenSuccesses >= m.config.HalfOpenMaxProbes { ep.CircuitState = CircuitClosed + ep.nextProbeAt = time.Time{} ep.Healthy = true logger.Logger.Info(). Str("endpoint", ep.Name). @@ -316,7 +362,8 @@ func (m *LLMHealthMonitor) handleFailure(ep *LLMEndpointStatus) { if ep.ConsecutiveFailures >= m.config.FailureThreshold { ep.CircuitState = CircuitOpen ep.Healthy = false - ep.circuitOpenedAt = time.Now() + ep.circuitOpenedAt = m.currentTime() + ep.nextProbeAt = m.nextCircuitProbeAtLocked(ep.circuitOpenedAt) logger.Logger.Error(). Str("endpoint", ep.Name). Int("consecutive_failures", ep.ConsecutiveFailures). @@ -327,7 +374,8 @@ func (m *LLMHealthMonitor) handleFailure(ep *LLMEndpointStatus) { // Any failure in half-open immediately re-opens the circuit ep.CircuitState = CircuitOpen ep.Healthy = false - ep.circuitOpenedAt = time.Now() + ep.circuitOpenedAt = m.currentTime() + ep.nextProbeAt = m.nextCircuitProbeAtLocked(ep.circuitOpenedAt) ep.halfOpenSuccesses = 0 logger.Logger.Warn(). Str("endpoint", ep.Name). diff --git a/control-plane/internal/services/llm_health_monitor_retry_test.go b/control-plane/internal/services/llm_health_monitor_retry_test.go index 6aeada44d..026f480b1 100644 --- a/control-plane/internal/services/llm_health_monitor_retry_test.go +++ b/control-plane/internal/services/llm_health_monitor_retry_test.go @@ -9,26 +9,40 @@ import ( ) func TestLLMHealthMonitor_RetryAfterSeconds(t *testing.T) { + base := time.Date(2026, time.August, 31, 12, 0, 0, 0, time.UTC) + now := base monitor := NewLLMHealthMonitor(config.LLMHealthConfig{ - RecoveryTimeout: 10 * time.Second, - Endpoints: []config.LLMEndpoint{{Name: "primary"}}, + CheckInterval: 6 * time.Second, + RecoveryTimeout: 10 * time.Second, + FailureThreshold: 1, + Endpoints: []config.LLMEndpoint{{Name: "primary"}}, }, nil) + monitor.now = func() time.Time { return now } assert.Equal(t, 10, monitor.RetryAfterSeconds("unknown")) assert.Equal(t, 10, monitor.RetryAfterSeconds("primary")) monitor.mu.Lock() - monitor.endpoints["primary"].CircuitState = CircuitOpen - monitor.endpoints["primary"].circuitOpenedAt = time.Now().Add(-4 * time.Second) + // Checks are scheduled at +6s, +12s, ...; a 10-second recovery timeout + // therefore cannot transition this circuit before the +12s tick. + monitor.nextCheckAt = base.Add(6 * time.Second) + monitor.handleFailure(monitor.endpoints["primary"]) monitor.mu.Unlock() - remaining := monitor.RetryAfterSeconds("primary") - assert.GreaterOrEqual(t, remaining, 5) - assert.LessOrEqual(t, remaining, 6) + assert.Equal(t, 12, monitor.RetryAfterSeconds("primary")) + now = base.Add(5*time.Second + 100*time.Millisecond) + assert.Equal(t, 7, monitor.RetryAfterSeconds("primary")) + now = base.Add(11*time.Second + 100*time.Millisecond) + assert.Equal(t, 1, monitor.RetryAfterSeconds("primary")) + + // If an open status predates scheduler bookkeeping, stay conservative by + // including one full check interval instead of reverting to the raw timeout. + now = base monitor.mu.Lock() - monitor.endpoints["primary"].circuitOpenedAt = time.Now().Add(-20 * time.Second) + monitor.endpoints["primary"].nextProbeAt = time.Time{} + monitor.endpoints["primary"].circuitOpenedAt = base monitor.mu.Unlock() - assert.Equal(t, 1, monitor.RetryAfterSeconds("primary")) + assert.Equal(t, 16, monitor.RetryAfterSeconds("primary")) defaults := NewLLMHealthMonitor(config.LLMHealthConfig{}, nil) assert.Equal(t, 30, defaults.RetryAfterSeconds("unknown")) diff --git a/docs/api/EXECUTE.md b/docs/api/EXECUTE.md index 810e3bd80..f14ddbb95 100644 --- a/docs/api/EXECUTE.md +++ b/docs/api/EXECUTE.md @@ -86,9 +86,9 @@ Execute requests can be rejected before dispatch: | --- | --- | --- | | `429` | Concurrency limit | `Retry-After: 1` and `{"error":"...","error_category":"concurrency_limit","retry_after":1}` | | `503` | Async dispatch queue full | `Retry-After: 1` and `{"error":"async execution queue is full; retry later","error_category":"concurrency_limit","retry_after":1}` | -| `503` | Control plane shutting down (async pool stopped) | `Retry-After: 1` and `{"error":"...","error_category":"concurrency_limit","retry_after":1}`; the `error` text distinguishes it from queue-full | +| `503` | Control plane shutting down (async pool stopped) | `Retry-After: 1` and `{"error":"...","error_category":"control_plane_shutdown","retry_after":1}`; an execution persisted before the pool stopped is terminalized with the same category | | `503` | Target node known to be down (after the drain hold expires) | `Retry-After: 1` and `{"error":"...","error_category":"node_unavailable","retry_after":1}` | -| `503` | Required LLM unavailable | `Retry-After: ` and `{"error":"...","error_category":"llm_unavailable","retry_after":}`; the window is the circuit breaker's remaining recovery timeout (default 30s, floor 1s) | +| `503` | Required LLM unavailable | `Retry-After: ` and `{"error":"...","error_category":"llm_unavailable","retry_after":}`; the window reaches the next scheduled health probe eligible to transition the circuit, including recovery timeout and check cadence (floor 1s) | | `413` | Body exceeds `AGENTFIELD_MAX_EXECUTE_BODY_BYTES` (default 32 MiB) | `{"error":"request body too large"}` | These pre-dispatch rejections persist no `executions` or `workflow_executions` row and no payload; the exception is a request rejected because the async pool has already stopped after preparation, which is terminated as `failed` with `status_reason` `control_plane_shutdown`. diff --git a/docs/api/EXECUTION_RESTART.md b/docs/api/EXECUTION_RESTART.md index 3a59b3aed..1145b495f 100644 --- a/docs/api/EXECUTION_RESTART.md +++ b/docs/api/EXECUTION_RESTART.md @@ -45,7 +45,7 @@ Operators polling execution state should branch on the stable category before an | `approval_rejected[: ...]` | Approval was rejected; an optional suffix contains feedback. | | `awaiting_child` | The parent is waiting for a child execution. | | `agent_client_error:` | The agent reported a client-facing HTTP 4xx failure. | -| `llm_unavailable`, `concurrency_limit`, `agent_timeout`, `agent_error`, `agent_unreachable`, `bad_response`, `internal_error`, `validation`, `permission_denied`, `node_unavailable`, `target_not_found` | Canonical failure categories used for operator routing and HTTP mapping. | +| `llm_unavailable`, `concurrency_limit`, `control_plane_shutdown`, `agent_timeout`, `agent_error`, `agent_unreachable`, `bad_response`, `internal_error`, `validation`, `permission_denied`, `node_unavailable`, `target_not_found` | Canonical failure categories used for operator routing and HTTP mapping. A pool that stops after restart persistence returns and stores `control_plane_shutdown`. | Concurrency and LLM-circuit admission checks run before restart persistence. A rejected restart creates no execution or workflow-execution row. Queue-full restart responses return `503` with matching `Retry-After` and `retry_after` values. From 294f60f2a4c844c2af4ed7721574af96ed6574e4 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Mon, 31 Aug 2026 17:50:12 -0400 Subject: [PATCH 8/8] fix(control-plane): preserve replay input envelopes --- .../handlers/workflow_execution_events.go | 7 ++++++- .../handlers/workflow_execution_events_test.go | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/control-plane/internal/handlers/workflow_execution_events.go b/control-plane/internal/handlers/workflow_execution_events.go index 3a41bdeed..c64626fe2 100644 --- a/control-plane/internal/handlers/workflow_execution_events.go +++ b/control-plane/internal/handlers/workflow_execution_events.go @@ -250,7 +250,12 @@ func applyEventToExecution(current *types.Execution, req *WorkflowExecutionEvent current.RunID = firstNonEmpty(req.RunID, req.WorkflowID, current.RunID) } - if payload := marshalJSON(req.InputData); len(payload) > 0 { + // Gateway-created executions already contain the authoritative persisted + // request envelope (input plus optional context). SDK lifecycle events carry + // only the reasoner's raw input, so replacing a non-empty payload here loses + // context and makes later restart/replay matching impossible. Event-created + // executions still receive their input in buildExecutionRecordFromEvent. + if payload := marshalJSON(req.InputData); len(current.InputPayload) == 0 && len(payload) > 0 { current.InputPayload = payload } if result := marshalJSON(req.Result); len(result) > 0 { diff --git a/control-plane/internal/handlers/workflow_execution_events_test.go b/control-plane/internal/handlers/workflow_execution_events_test.go index d63c125ea..a57584114 100644 --- a/control-plane/internal/handlers/workflow_execution_events_test.go +++ b/control-plane/internal/handlers/workflow_execution_events_test.go @@ -101,6 +101,22 @@ func TestWorkflowExecutionEventHandler_CreateAndUpdate(t *testing.T) { assert.Equal(t, duration, *wfExec.DurationMS) } +func TestWorkflowExecutionEventHandler_PreservesGatewayInputEnvelope(t *testing.T) { + gatewayPayload := json.RawMessage(`{"input":{"message":"same"},"context":{"tenant":"acme"}}`) + current := &types.Execution{ + ExecutionID: "exec_gateway", + Status: string(types.ExecutionStatusRunning), + InputPayload: append(json.RawMessage(nil), gatewayPayload...), + } + + applyEventToExecution(current, &WorkflowExecutionEventRequest{ + Status: string(types.ExecutionStatusRunning), + InputData: map[string]interface{}{"message": "same"}, + }, time.Now().UTC()) + + require.JSONEq(t, string(gatewayPayload), string(current.InputPayload)) +} + // TestWorkflowExecutionEventHandler_TerminalRegression covers the case where // fire-and-forget workflow events from the SDK arrive out of order — e.g. an // outer reasoner emits "failed" while an inner reasoner emits a delayed