218 lines
5.8 KiB
Go
218 lines
5.8 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/kamelion/ratatoskr-go/internal/opencode"
|
|
"github.com/kamelion/ratatoskr-go/internal/storage"
|
|
)
|
|
|
|
// OpenCodeRunner — интерфейс для opencode (подменяемый в тестах).
|
|
type OpenCodeRunner interface {
|
|
Run(ctx context.Context, prompt, cwd, agent, sessionID string) (*opencode.Result, error)
|
|
}
|
|
|
|
// PollTaskFunc — callback для обработки готовой задачи (подменяемый в тестах).
|
|
type PollTaskFunc func(ctx context.Context) error
|
|
|
|
// Worker — планировщик, запускающий готовые задачи (status=ready → running → success/failed/timeout).
|
|
type Worker struct {
|
|
Store *storage.Storage
|
|
Runner OpenCodeRunner
|
|
Worktree string // базовый путь, task.Repo — относительно него
|
|
Agent string // default "dev"
|
|
Interval time.Duration // интервал опроса БД
|
|
MaxJobs int // макс. параллельных задач
|
|
|
|
sem chan struct{} // семафор
|
|
cancel context.CancelFunc
|
|
|
|
// подменяемый poll для тестов
|
|
pollFn PollTaskFunc
|
|
}
|
|
|
|
// Start запускает цикл опроса в фоновой горутине.
|
|
func (w *Worker) Start(ctx context.Context) {
|
|
if w.Agent == "" {
|
|
w.Agent = "dev"
|
|
}
|
|
if w.Interval <= 0 {
|
|
w.Interval = 5 * time.Second
|
|
}
|
|
if w.MaxJobs <= 0 {
|
|
w.MaxJobs = 2
|
|
}
|
|
w.sem = make(chan struct{}, w.MaxJobs)
|
|
// заполняем семафор токенами
|
|
for i := 0; i < w.MaxJobs; i++ {
|
|
w.sem <- struct{}{}
|
|
}
|
|
|
|
ctx, w.cancel = context.WithCancel(ctx)
|
|
|
|
pollFn := w.pollFn
|
|
if pollFn == nil {
|
|
pollFn = w.pollAndDispatch
|
|
}
|
|
|
|
go func() {
|
|
// первый poll сразу
|
|
_ = pollFn(ctx)
|
|
|
|
ticker := time.NewTicker(w.Interval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
log.Print("worker: stopped")
|
|
return
|
|
case <-ticker.C:
|
|
_ = pollFn(ctx)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Stop останавливает воркер (отменяет контекст → убивает активные задачи).
|
|
func (w *Worker) Stop() {
|
|
if w.cancel != nil {
|
|
w.cancel()
|
|
}
|
|
}
|
|
|
|
// pollAndDispatch ищет готовые задачи, запускает их в пределах свободных слотов.
|
|
func (w *Worker) pollAndDispatch(ctx context.Context) error {
|
|
slots := len(w.sem)
|
|
if slots == 0 {
|
|
return nil
|
|
}
|
|
|
|
tasks, err := w.Store.ListTasks(ctx, storage.TaskFilter{
|
|
Status: storage.StatusReady,
|
|
Limit: slots,
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("%w: %v", ErrPoll, err)
|
|
}
|
|
|
|
for _, t := range tasks {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case slot := <-w.sem:
|
|
task := t
|
|
go func() {
|
|
defer func() { w.sem <- slot }()
|
|
if err := w.runTask(ctx, task); err != nil {
|
|
log.Printf("worker: task %d: %v", task.ID, err)
|
|
}
|
|
}()
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// runTask выполняет одну задачу: dev-агент через opencode.
|
|
func (w *Worker) runTask(ctx context.Context, task *storage.Task) (err error) {
|
|
// 1. проверяем статус
|
|
if task.Status != storage.StatusReady {
|
|
return fmt.Errorf("%w: task %d status=%q", ErrLaunch, task.ID, task.Status)
|
|
}
|
|
|
|
// 2. ставим running
|
|
task.Status = storage.StatusRunning
|
|
if err := w.Store.UpdateTask(ctx, task); err != nil {
|
|
return fmt.Errorf("%w: set running: %v", ErrUpdate, err)
|
|
}
|
|
|
|
// 3. рендерим промпт
|
|
prompt, err := RenderDevPrompt(DevPromptData{
|
|
Title: task.Title,
|
|
Goal: task.Goal,
|
|
Repo: task.Repo,
|
|
Why: task.Why,
|
|
AC: task.AC,
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("%w: render prompt: %v", ErrTrace, err)
|
|
}
|
|
|
|
// 4. создаём трассу
|
|
trace := &storage.Trace{
|
|
TaskID: task.ID,
|
|
Agent: w.Agent,
|
|
Prompt: prompt,
|
|
}
|
|
traceID, err := w.Store.AppendTrace(ctx, trace)
|
|
if err != nil {
|
|
return fmt.Errorf("%w: create: %v", ErrTrace, err)
|
|
}
|
|
|
|
// 5. вычисляем cwd
|
|
cwd := w.resolveCwd(task.Repo)
|
|
|
|
// 6. запускаем dev-агент
|
|
res, resErr := w.Runner.Run(ctx, prompt, cwd, w.Agent, "")
|
|
if resErr != nil {
|
|
// O1 ErrSpawn — не смог запустить бинарь
|
|
task.Status = storage.StatusFailed
|
|
if e := w.Store.UpdateTask(ctx, task); e != nil {
|
|
err = fmt.Errorf("%w: set failed: %v", ErrUpdate, e)
|
|
return
|
|
}
|
|
w.finalizeTrace(ctx, traceID, storage.TraceFailed, resErr.Error())
|
|
err = fmt.Errorf("%w: spawn: %v", ErrLaunch, resErr)
|
|
return
|
|
}
|
|
|
|
// 6b. сохраняем session_id из результата
|
|
if res.SessionID != "" {
|
|
_ = w.Store.UpdateTraceSessionID(ctx, traceID, res.SessionID)
|
|
}
|
|
|
|
// 7. определяем результат по RC
|
|
output := res.Stdout
|
|
var traceStatus storage.TraceStatus
|
|
|
|
switch {
|
|
case res.RC == 0:
|
|
task.Status = storage.StatusSuccess
|
|
traceStatus = storage.TraceSuccess
|
|
case res.RC == -1:
|
|
task.Status = storage.StatusTimeout
|
|
traceStatus = storage.TraceTimeout
|
|
default:
|
|
task.Status = storage.StatusFailed
|
|
traceStatus = storage.TraceFailed
|
|
}
|
|
|
|
// 8. сохраняем результат
|
|
if e := w.Store.UpdateTask(ctx, task); e != nil {
|
|
err = fmt.Errorf("%w: set %s: %v", ErrUpdate, task.Status, e)
|
|
return
|
|
}
|
|
w.finalizeTrace(ctx, traceID, traceStatus, output)
|
|
return nil
|
|
}
|
|
|
|
// finalizeTrace обновляет output и статус трассы.
|
|
func (w *Worker) finalizeTrace(ctx context.Context, traceID int64, status storage.TraceStatus, output string) {
|
|
if e := w.Store.UpdateTraceOutput(ctx, traceID, output); e != nil {
|
|
log.Printf("worker: update trace output %d: %v", traceID, e)
|
|
}
|
|
if e := w.Store.UpdateTraceStatus(ctx, traceID, status); e != nil {
|
|
log.Printf("worker: update trace status %d: %v", traceID, e)
|
|
}
|
|
}
|
|
|
|
func (w *Worker) resolveCwd(repo string) string {
|
|
if repo == "" {
|
|
return w.Worktree
|
|
}
|
|
return filepath.Join(w.Worktree, repo)
|
|
} |