Adds pkg/ratelimit for Gemini API rate limiting with sliding window
(RPM/TPM/RPD), persistent state, and token counting. Replaces the
bash agent-runner.sh with a native Go implementation under
`core ai dispatch {run,watch,status}` for local queue processing.
Rate limiting:
- Per-model quotas (RPM, TPM, RPD) with 1-minute sliding window
- WaitForCapacity blocks until capacity available or context cancelled
- Persistent state in ~/.core/ratelimits.yaml
- Default quotas for Gemini 3 Pro/Flash, 2.5 Pro, 2.0 Flash/Lite
- CountTokens helper calls Google tokenizer API
- CLI: core ai ratelimits {show,reset,count,config,check}
Dispatch runner:
- core ai dispatch run — process single ticket from queue
- core ai dispatch watch — daemon mode with configurable interval
- core ai dispatch status — show queue/active/done counts
- Supports claude/codex/gemini runners with rate-limited Gemini
- File-based locking with stale PID detection
- Completion handler updates issue labels on success/failure
Closes #42
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
87 lines
2.7 KiB
Go
87 lines
2.7 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/host-uk/core/pkg/forge"
|
|
"github.com/host-uk/core/pkg/jobrunner"
|
|
)
|
|
|
|
const (
|
|
ColorAgentComplete = "#0e8a16" // Green
|
|
)
|
|
|
|
// CompletionHandler manages issue state when an agent finishes work.
|
|
type CompletionHandler struct {
|
|
forge *forge.Client
|
|
}
|
|
|
|
// NewCompletionHandler creates a handler for agent completion events.
|
|
func NewCompletionHandler(client *forge.Client) *CompletionHandler {
|
|
return &CompletionHandler{
|
|
forge: client,
|
|
}
|
|
}
|
|
|
|
// Name returns the handler identifier.
|
|
func (h *CompletionHandler) Name() string {
|
|
return "completion"
|
|
}
|
|
|
|
// Match returns true if the signal indicates an agent has finished a task.
|
|
func (h *CompletionHandler) Match(signal *jobrunner.PipelineSignal) bool {
|
|
return signal.Type == "agent_completion"
|
|
}
|
|
|
|
// Execute updates the issue labels based on the completion status.
|
|
func (h *CompletionHandler) Execute(ctx context.Context, signal *jobrunner.PipelineSignal) (*jobrunner.ActionResult, error) {
|
|
start := time.Now()
|
|
|
|
// Remove in-progress label.
|
|
if inProgressLabel, err := h.forge.GetLabelByName(signal.RepoOwner, signal.RepoName, LabelInProgress); err == nil {
|
|
_ = h.forge.RemoveIssueLabel(signal.RepoOwner, signal.RepoName, int64(signal.ChildNumber), inProgressLabel.ID)
|
|
}
|
|
|
|
if signal.Success {
|
|
completeLabel, err := h.forge.EnsureLabel(signal.RepoOwner, signal.RepoName, LabelAgentComplete, ColorAgentComplete)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("ensure label %s: %w", LabelAgentComplete, err)
|
|
}
|
|
|
|
if err := h.forge.AddIssueLabels(signal.RepoOwner, signal.RepoName, int64(signal.ChildNumber), []int64{completeLabel.ID}); err != nil {
|
|
return nil, fmt.Errorf("add completed label: %w", err)
|
|
}
|
|
|
|
if signal.Message != "" {
|
|
_ = h.forge.CreateIssueComment(signal.RepoOwner, signal.RepoName, int64(signal.ChildNumber), signal.Message)
|
|
}
|
|
} else {
|
|
failedLabel, err := h.forge.EnsureLabel(signal.RepoOwner, signal.RepoName, LabelAgentFailed, ColorAgentFailed)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("ensure label %s: %w", LabelAgentFailed, err)
|
|
}
|
|
|
|
if err := h.forge.AddIssueLabels(signal.RepoOwner, signal.RepoName, int64(signal.ChildNumber), []int64{failedLabel.ID}); err != nil {
|
|
return nil, fmt.Errorf("add failed label: %w", err)
|
|
}
|
|
|
|
msg := "Agent reported failure."
|
|
if signal.Error != "" {
|
|
msg += fmt.Sprintf("\n\nError: %s", signal.Error)
|
|
}
|
|
_ = h.forge.CreateIssueComment(signal.RepoOwner, signal.RepoName, int64(signal.ChildNumber), msg)
|
|
}
|
|
|
|
return &jobrunner.ActionResult{
|
|
Action: "completion",
|
|
RepoOwner: signal.RepoOwner,
|
|
RepoName: signal.RepoName,
|
|
EpicNumber: signal.EpicNumber,
|
|
ChildNumber: signal.ChildNumber,
|
|
Success: true,
|
|
Timestamp: time.Now(),
|
|
Duration: time.Since(start),
|
|
}, nil
|
|
}
|