diff --git a/README.md b/README.md index fa43e6bd4..1bb706937 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. | @@ -319,7 +319,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/control-plane/internal/handlers/coverage_handlers_90_additional_test.go b/control-plane/internal/handlers/coverage_handlers_90_additional_test.go index 8559384f1..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, @@ -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..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) }, }, { @@ -934,6 +936,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.go b/control-plane/internal/handlers/execute.go index 119b74949..afd91b8ac 100644 --- a/control-plane/internal/handlers/execute.go +++ b/control-plane/internal/handlers/execute.go @@ -269,6 +269,7 @@ func (c *executionController) handleSync(ctx *gin.Context) { return } plan.executionMode = "sync" + defer plan.releaseSlot() if plan.replayHit != nil { if err := c.completeReplayHit(reqCtx, plan); err != nil { @@ -290,14 +291,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) @@ -437,9 +430,15 @@ 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. - if !pool.reserve() { + // 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 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 } @@ -456,9 +455,12 @@ 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 { - ReleaseExecutionSlot(plan.target.NodeID) if err := c.completeReplayHit(reqCtx, plan); err != nil { writeExecutionError(ctx, err) return @@ -497,12 +499,26 @@ 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 { - ReleaseExecutionSlot(plan.target.NodeID) // Release since process() won't run - writeAsyncAdmissionError(ctx, http.StatusServiceUnavailable, "async execution queue stopped; retry later") + // 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. + 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 a94fffd3b..eab76a5f0 100644 --- a/control-plane/internal/handlers/execute_async_test.go +++ b/control-plane/internal/handlers/execute_async_test.go @@ -36,6 +36,198 @@ 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, 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) + 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.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 + s.pool.mu.Unlock() + err := s.testExecutionStorage.CreateExecutionRecord(ctx, execution) + if s.cancel != nil { + s.cancel() + } + 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 955788821..9b14fbc86 100644 --- a/control-plane/internal/handlers/execute_helpers.go +++ b/control-plane/internal/handlers/execute_helpers.go @@ -562,20 +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() - } - if pe.Category() == ErrorCategoryConcurrencyLimit || pe.Category() == ErrorCategoryNodeUnavailable { - ctx.Header("Retry-After", "1") - body["retry_after"] = 1 + body, retryAfter := renderExecutionPreconditionError(pe) + if retryAfter > 0 { + ctx.Header("Retry-After", strconv.Itoa(retryAfter)) } ctx.JSON(pe.HTTPStatusCode(), body) return @@ -593,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 { @@ -840,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.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 @@ -872,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 } @@ -943,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) @@ -973,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 } } @@ -1037,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 { @@ -1050,27 +1071,48 @@ func shutdownPersistenceContext() (context.Context, context.CancelFunc) { return context.WithTimeout(context.Background(), 5*time.Second) } -func (j asyncExecutionJob) failForControlPlaneShutdown(ctx context.Context) { +func (j *asyncExecutionJob) terminateForControlPlaneShutdown(shutdownErr *executionPreconditionError) { + persistCtx, cancel := shutdownPersistenceContext() + defer cancel() + 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, shutdownErr *executionPreconditionError) { + nodeID := "" if j.plan.target != nil { - ReleaseExecutionSlot(j.plan.target.NodeID) + nodeID = j.plan.target.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.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 + } + // 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 } - reason := string(ErrorCategoryControlPlaneShutdown) 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) } - current.StatusReason = &reason + current.Status = string(updatedExec.Status) + current.StatusReason = updatedExec.StatusReason + current.CompletedAt = updatedExec.CompletedAt + current.UpdatedAt = updatedExec.UpdatedAt 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_lifecycle.go b/control-plane/internal/handlers/execute_lifecycle.go index ae4639fea..bb053f55b 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. @@ -210,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 4f6df4140..cf2d976bf 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,11 +41,7 @@ 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) { +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) @@ -153,33 +149,46 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context } target.TargetType = targetType + storedPayload, err := json.Marshal(buildClientPayload(req)) + 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 { + slotTransferred := false + 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 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) } }() } - runID := headers.runID - if runID == "" { - runID = utils.GenerateRunID() - } - executionID := utils.GenerateExecutionID() now := time.Now().UTC() - storedPayload, err := json.Marshal(buildClientPayload(req)) - if err != nil { - return nil, fmt.Errorf("encode execution payload: %w", err) - } - exec := &types.Execution{ ExecutionID: executionID, RunID: runID, @@ -256,18 +265,14 @@ 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{ + plan := &preparedExecution{ exec: exec, requestBody: agentPayloadBytes, agent: agent, target: target, targetType: targetType, llmEndpoint: llmEndpoint, + slotHeld: slotAcquired, webhookRegistered: webhookRegistered, webhookError: webhookError, callerDID: callerDID, @@ -277,7 +282,9 @@ func (c *executionController) prepareExecutionForTargetWithAdmission(ctx context replayBeforeExecutionID: headers.replayBeforeExecutionID, replayMode: headers.replayMode, replayHit: hit, - }, nil + } + slotTransferred = true + return plan, nil } // buildClientPayload builds the blob persisted as executions.input_payload, diff --git a/control-plane/internal/handlers/execute_restart.go b/control-plane/internal/handlers/execute_restart.go index eb1d67bee..a5365c55f 100644 --- a/control-plane/internal/handlers/execute_restart.go +++ b/control-plane/internal/handlers/execute_restart.go @@ -151,22 +151,35 @@ func (c *executionController) handleRestart(ctx *gin.Context) { } target := fmt.Sprintf("%s.%s", restartExec.NodeID, restartExec.ReasonerID) + pool := getAsyncWorkerPool() + 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 + } + reserved := true + defer func() { + if reserved { + pool.releaseReservation() + } + }() + // Restarts mint a new run identity and deliberately leave RunMetadata nil. - plan, err := c.prepareExecutionForTarget(reqCtx, target, ExecuteRequest{ + 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 - } + // 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 { @@ -176,20 +189,25 @@ 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 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") + plan.slotHeld = false // ownership transferred to job + submitted := false + defer func() { + if !submitted { + job.plan.releaseSlot() } - ctx.JSON(http.StatusServiceUnavailable, gin.H{"error": queueErr.Error(), "error_category": "concurrency_limit"}) + }() + if ok := pool.submitReserved(job); !ok { + 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) var replayBefore *string 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..f9cc0fe9a --- /dev/null +++ b/control-plane/internal/handlers/execute_sync_admission_test.go @@ -0,0 +1,387 @@ +package handlers + +import ( + "context" + "encoding/json" + "fmt" + "io" + "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" +) + +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 + 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}, + {ErrorCategoryControlPlaneShutdown, 1}, + } + 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) +} + +// A restart admitted past reserve() but refused by submitReserved (the pool +// 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) + 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, + }) + 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(`{}`)).WithContext(reqCtx) + 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, string(ErrorCategoryControlPlaneShutdown), 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(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.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 ff1661e71..10b98e12e 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,24 @@ 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, + } +} + +func newControlPlaneShutdownError(message string) *executionPreconditionError { + return &executionPreconditionError{ + code: 503, + message: message, + category: ErrorCategoryControlPlaneShutdown, + retryAfter: 1, } } @@ -152,10 +163,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 b64ba1e97..1dbd980ed 100644 --- a/control-plane/internal/handlers/mcp.go +++ b/control-plane/internal/handlers/mcp.go @@ -417,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,27 +523,44 @@ 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) - // MCP creates an ordinary execution and deliberately leaves RunMetadata nil. - plan, err := controller.prepareExecutionForTarget(ctx, target, ExecuteRequest{Input: input}, headers, callerDID, targetDID) - if err != nil { - return "", "", err + pool := getAsyncWorkerPool() + 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 + defer func() { + if reserved { + pool.releaseReservation() + } + }() - if err := CheckExecutionPreconditions(plan.target.NodeID, plan.llmEndpoint); err != nil { - _ = controller.failExecution(ctx, plan, err, 0, nil) + // MCP creates an ordinary execution and deliberately leaves RunMetadata nil. + plan, err := controller.prepareExecutionForTargetWithAdmission(ctx, target, ExecuteRequest{Input: input}, headers, callerDID, targetDID, true) + if err != nil { return "", "", err } + defer plan.releaseSlot() 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") - _ = controller.failExecution(ctx, plan, queueErr, 0, nil) - return "", "", queueErr + plan.slotHeld = false // ownership transferred to job + submitted := false + defer func() { + if !submitted { + job.plan.releaseSlot() + } + }() + if ok := pool.submitReserved(job); !ok { + 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 } @@ -707,6 +735,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 d5e002385..8e1080cde 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,41 @@ type mcpTestStore struct { agents []*types.AgentNode } +type stopPoolOnCreateMCPStore struct { + *mcpTestStore + pool *asyncWorkerPool + cancel context.CancelFunc + updateErr error + 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 + 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 { @@ -342,6 +378,135 @@ 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_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) + 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) + require.Equal(t, ErrorCategoryControlPlaneShutdown, admissionErr.category) + require.Equal(t, 1, admissionErr.retryAfter) + + 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) + 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) { + 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 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 diff --git a/control-plane/internal/services/llm_health_monitor.go b/control-plane/internal/services/llm_health_monitor.go index dc653f1c5..33ef1ed10 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" @@ -35,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. @@ -84,6 +88,7 @@ func NewLLMHealthMonitor(cfg config.LLMHealthConfig, uiService *UIService) *LLMH endpoints: endpoints, stopCh: make(chan struct{}), uiService: uiService, + now: time.Now, } } @@ -102,6 +107,36 @@ 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() { + 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 { + return 1 + } + return seconds +} + // Start begins the health monitoring loop. func (m *LLMHealthMonitor) Start() { if !m.config.Enabled || len(m.config.Endpoints) == 0 { @@ -117,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() @@ -126,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() } } @@ -216,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). @@ -264,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 { @@ -273,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). @@ -292,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). @@ -303,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 new file mode 100644 index 000000000..026f480b1 --- /dev/null +++ b/control-plane/internal/services/llm_health_monitor_retry_test.go @@ -0,0 +1,50 @@ +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) { + base := time.Date(2026, time.August, 31, 12, 0, 0, 0, time.UTC) + now := base + monitor := NewLLMHealthMonitor(config.LLMHealthConfig{ + 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() + // 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() + 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"].nextProbeAt = time.Time{} + monitor.endpoints["primary"].circuitOpenedAt = base + monitor.mu.Unlock() + assert.Equal(t, 16, monitor.RetryAfterSeconds("primary")) + + defaults := NewLLMHealthMonitor(config.LLMHealthConfig{}, nil) + assert.Equal(t, 30, defaults.RetryAfterSeconds("unknown")) + assert.Equal(t, 30, (*LLMHealthMonitor)(nil).RetryAfterSeconds("unknown")) +} diff --git a/docs/ENVIRONMENT_VARIABLES.md b/docs/ENVIRONMENT_VARIABLES.md index 127e803e6..73fc32e61 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. Accepts bare seconds (`30`) and Go duration strings (`30s`, `5m`). The same budget is shared with `StopAsyncWorkerPool`, which is guaranteed a fresh budget of at least 5s, so total shutdown can exceed this value. See the [Kubernetes shutdown and drain recipe](deploying-on-kubernetes.md). diff --git a/docs/api/EXECUTE.md b/docs/api/EXECUTE.md index 18e2babd2..84af99859 100644 --- a/docs/api/EXECUTE.md +++ b/docs/api/EXECUTE.md @@ -90,9 +90,11 @@ 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 | `{"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 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`. + 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..1145b495f 100644 --- a/docs/api/EXECUTION_RESTART.md +++ b/docs/api/EXECUTION_RESTART.md @@ -45,6 +45,8 @@ 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. Do not emulate restart by re-submitting `/execute`: execute has no idempotency key, creates unrelated executions, and cannot establish restart lineage or replay boundaries.