rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add agent run supervisor worker

Browse Source 
Poll-based supervisor that claims PENDING agent runs with FOR UPDATE
SKIP LOCKED, runs them with lease-based heartbeat, and handles
graceful shutdown. On infrastructure stop the row stays RUNNING so
stale recovery resets it to PENDING on restart; Restore picks up
from the last checkpoint. Heartbeat loss cancels execution without
committing a terminal status.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
Aurélien Sibiril
2026-04-10 18:53:29 +02:00
parent 2f2b8fe6aa
commit e831ee3c63
4 changed files with 422 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pkg/agent/checkpoint.go
View File

@@ -35,7 +35,7 @@ type (
// Approval-interrupted checkpoints carry pending tool calls.
PendingToolCalls []llm.ToolCall `json:"pending_tool_calls,omitempty"`
PendingApprovals []llm.ToolCall `json:"pending_approvals,omitempty"`
ApprovalInput map[string]ApprovalResult `json:"approval_input,omitempty"`
ApprovalInput map[string]ApprovalResult `json:"approval_input,omitempty"` // keyed by tool call ID
// Nested agent-as-tool suspension: one entry per suspended inner agent.
AllToolCalls []llm.ToolCall `json:"all_tool_calls,omitempty"`

3
pkg/agent/restore.go
View File

@@ -60,9 +60,6 @@ func restoreCheckpoint(
runID string,
registry AgentRegistry,
) (*Result, error) {
if cp.Version != CheckpointVersion {
return nil, fmt.Errorf("cannot restore: unsupported checkpoint version %d", cp.Version)
}
emitHook(agent, func(h RunHooks) { h.OnRunRestore(ctx, agent, cp) })
switch cp.Status {

9
pkg/agent/run.go
View File

@@ -374,6 +374,10 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: %w", err))
}
if s.turns >= s.agent.maxTurns {
return s.finishRun(ctx, nil, &MaxTurnsExceededError{MaxTurns: s.agent.maxTurns})
}
if ch := stopSignalFrom(ctx); ch != nil {
select {
case <-ch:
@@ -383,6 +387,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
if s.opts.checkpointStore != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr))
se.Checkpoint = cp
}
} else {
se.Checkpoint = cp
@@ -393,10 +398,6 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
}
}
if s.turns >= s.agent.maxTurns {
return s.finishRun(ctx, nil, &MaxTurnsExceededError{MaxTurns: s.agent.maxTurns})
}
fullMessages := buildFullMessages(s.systemPrompt, s.messages)
var responseFormat *llm.ResponseFormat