teabot/internal/dispatch/jobs.go

package dispatch

import (
	"context"
	"fmt"

	"git.unkin.net/unkin/teabot/internal/config"
	"git.unkin.net/unkin/teabot/internal/docker"
	"git.unkin.net/unkin/teabot/internal/gitea"
	"git.unkin.net/unkin/teabot/internal/prompt"
)

// baseJob fills the personality/repo/runtime fields shared by every job kind.
func (d *Dispatcher) baseJob(repo string, p config.Personality, label, promptText string) docker.Job {
	return docker.Job{
		Label:            label,
		Image:            d.cfg.JobImage,
		ContainerHome:    d.cfg.ContainerHome,
		Prompt:           promptText,
		CloneURL:         d.cloneURL(repo),
		GitHost:          d.gitHost,
		GitName:          p.GitName,
		GitEmail:         p.GitEmail,
		GitUser:          p.Login,
		Token:            p.Token,
		TeaConfigPath:    p.TeaConfig,
		ClaudeConfigDir:  d.cfg.ClaudeConfigDir,
		AnthropicAPIKey:  d.cfg.AnthropicAPIKey,
		AnthropicBaseURL: d.cfg.AnthropicBaseURL,
		Timeout:          d.cfg.JobTimeout,
	}
}

func (d *Dispatcher) dispatchIssue(ctx context.Context, repo string, p config.Personality, issue gitea.Issue, comments []gitea.Comment) {
	text := prompt.Issue(prompt.IssueContext{
		Repo:            repo,
		PersonalityName: p.Name,
		Issue:           issue,
		Comments:        comments,
	})
	label := fmt.Sprintf("%s#issue-%d", repo, issue.Index)
	d.runJob(ctx, d.baseJob(repo, p, label, text))
}

func (d *Dispatcher) dispatchPull(ctx context.Context, repo string, p config.Personality, pull gitea.PullRequest, diff string, comments []gitea.Comment) {
	text := prompt.Pull(prompt.PullContext{
		Repo:            repo,
		PersonalityName: p.Name,
		Pull:            pull,
		Diff:            diff,
		Comments:        comments,
	})
	label := fmt.Sprintf("%s#pull-%d", repo, pull.Index)
	d.runJob(ctx, d.baseJob(repo, p, label, text))
}

func (d *Dispatcher) dispatchIssueFollowUp(ctx context.Context, repo string, p config.Personality, issue gitea.Issue, thread []gitea.Comment, trigger gitea.Comment) {
	text := prompt.FollowUp(prompt.FollowUpContext{
		Repo:            repo,
		PersonalityName: p.Name,
		Kind:            prompt.FollowUpIssue,
		Index:           issue.Index,
		Title:           issue.Title,
		URL:             issue.HTMLURL,
		Comments:        thread,
		NewComment:      trigger,
	})
	label := fmt.Sprintf("%s#issue-%d-followup-%d", repo, issue.Index, trigger.ID)
	d.runJob(ctx, d.baseJob(repo, p, label, text))
}

func (d *Dispatcher) dispatchPullFollowUp(ctx context.Context, repo string, p config.Personality, pull gitea.PullRequest, thread []gitea.Comment, trigger gitea.Comment) {
	text := prompt.FollowUp(prompt.FollowUpContext{
		Repo:            repo,
		PersonalityName: p.Name,
		Kind:            prompt.FollowUpPull,
		Index:           pull.Index,
		Title:           pull.Title,
		URL:             pull.HTMLURL,
		Comments:        thread,
		NewComment:      trigger,
	})
	label := fmt.Sprintf("%s#pull-%d-followup-%d", repo, pull.Index, trigger.ID)
	d.runJob(ctx, d.baseJob(repo, p, label, text))
}

// runJob launches a job in a bounded goroutine so at most MaxConcurrent
// containers run at once.
func (d *Dispatcher) runJob(ctx context.Context, job docker.Job) {
	d.wg.Add(1)
	go func() {
		defer d.wg.Done()
		select {
		case d.sem <- struct{}{}:
			defer func() { <-d.sem }()
		case <-ctx.Done():
			d.log.Warn("cancelled before start", "job", job.Label)
			return
		}
		d.log.Info("dispatching job", "job", job.Label, "image", job.Image)
		res, err := d.runner.Run(ctx, job)
		if err != nil {
			d.log.Error("job failed", "job", job.Label, "err", err, "output", tail(res.Output))
			return
		}
		if res.ExitCode != 0 {
			d.log.Warn("job exited non-zero",
				"job", job.Label, "exit", res.ExitCode, "duration", res.Duration, "output", tail(res.Output))
			return
		}
		d.log.Info("job completed", "job", job.Label, "duration", res.Duration)
	}()
}

// tail returns the last chunk of output for concise error logging.
func tail(s string) string {
	const max = 2000
	if len(s) <= max {
		return s
	}
	return "..." + s[len(s)-max:]
}