ai-workflow-skill/internal/store/council.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"
)

var councilReviewerRoles = []string{
	"architecture-reviewer",
	"implementation-reviewer",
	"risk-reviewer",
}

type CouncilRun struct {
	RunID         string `json:"run_id"`
	Mode          string `json:"mode"`
	TargetType    string `json:"target_type"`
	OutputMode    string `json:"output_mode"`
	OnlyUnanimous bool   `json:"only_unanimous"`
}

type CouncilInput struct {
	RunID        string `json:"run_id"`
	Prompt       string `json:"prompt,omitempty"`
	TargetFile   string `json:"target_file,omitempty"`
	RepoPath     string `json:"repo_path,omitempty"`
	TargetTaskID string `json:"task_id,omitempty"`
}

type CouncilReviewer struct {
	ReviewerRole string `json:"reviewer_role"`
	TaskID       string `json:"task_id"`
	Status       string `json:"status"`
}

type CouncilStartInput struct {
	RunID         string
	Target        string
	TargetFile    string
	RepoPath      string
	TargetTaskID  string
	TargetType    string
	Mode          string
	OutputMode    string
	OnlyUnanimous bool
}

type CouncilStartResult struct {
	Run       CouncilRun        `json:"run"`
	Input     CouncilInput      `json:"input"`
	Reviewers []CouncilReviewer `json:"reviewers"`
}

type CouncilWaitInput struct {
	RunID   string
	Timeout time.Duration
}

type CouncilWaitResult struct {
	Woke             bool              `json:"woke"`
	RunID            string            `json:"run_id"`
	AllComplete      bool              `json:"all_complete"`
	ReviewerStatuses []CouncilReviewer `json:"reviewers"`
}

func (s *OrchStore) StartCouncil(ctx context.Context, input CouncilStartInput) (CouncilStartResult, error) {
	runID := strings.TrimSpace(input.RunID)
	if runID == "" {
		return CouncilStartResult{}, fmt.Errorf("%w: run id is required", ErrInvalidInput)
	}

	councilInput, err := normalizeCouncilInput(input)
	if err != nil {
		return CouncilStartResult{}, err
	}

	councilRun, err := normalizeCouncilRun(input)
	if err != nil {
		return CouncilStartResult{}, err
	}

	now := nowUTC()
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return CouncilStartResult{}, fmt.Errorf("begin council start transaction: %w", err)
	}
	defer tx.Rollback()

	if _, err := selectRun(ctx, tx, runID); err == nil {
		return CouncilStartResult{}, fmt.Errorf("%w: run %s already exists", ErrInvalidState, runID)
	} else if !errors.Is(err, ErrRunNotFound) {
		return CouncilStartResult{}, err
	}

	goal := buildCouncilRunGoal(councilInput)
	summary := buildCouncilRunSummary(councilRun, councilInput)

	_, err = tx.ExecContext(
		ctx,
		`INSERT INTO runs (run_id, goal, summary, status, created_at, updated_at)
		 VALUES (?, ?, ?, ?, ?, ?)`,
		runID,
		goal,
		summary,
		"active",
		formatTime(now),
		formatTime(now),
	)
	if err != nil {
		return CouncilStartResult{}, fmt.Errorf("insert council run into runs: %w", err)
	}

	if err := insertEvent(ctx, tx, eventInput{
		RunID:       runID,
		Source:      "orch",
		EventType:   "run_initialized",
		Summary:     summary,
		PayloadJSON: marshalJSON(map[string]any{"goal": goal, "summary": summary}),
		CreatedAt:   now,
	}); err != nil {
		return CouncilStartResult{}, err
	}

	_, err = tx.ExecContext(
		ctx,
		`INSERT INTO council_runs (
			run_id, mode, target_type, output_mode, only_unanimous, created_at, updated_at
		) VALUES (?, ?, ?, ?, ?, ?, ?)`,
		runID,
		councilRun.Mode,
		councilRun.TargetType,
		councilRun.OutputMode,
		boolToInt(councilRun.OnlyUnanimous),
		formatTime(now),
		formatTime(now),
	)
	if err != nil {
		return CouncilStartResult{}, fmt.Errorf("insert council run metadata: %w", err)
	}

	_, err = tx.ExecContext(
		ctx,
		`INSERT INTO council_inputs (
			run_id, prompt, target_file, repo_path, target_task_id, created_at, updated_at
		) VALUES (?, ?, ?, ?, ?, ?, ?)`,
		runID,
		councilInput.Prompt,
		councilInput.TargetFile,
		councilInput.RepoPath,
		councilInput.TargetTaskID,
		formatTime(now),
		formatTime(now),
	)
	if err != nil {
		return CouncilStartResult{}, fmt.Errorf("insert council input metadata: %w", err)
	}

	reviewers := make([]CouncilReviewer, 0, len(councilReviewerRoles))
	for i, reviewerRole := range councilReviewerRoles {
		taskID := fmt.Sprintf("CR%d", i+1)
		task := Task{
			RunID:          runID,
			TaskID:         taskID,
			Title:          buildCouncilTaskTitle(reviewerRole),
			Summary:        buildCouncilTaskSummary(reviewerRole),
			Status:         "ready",
			DefaultTo:      reviewerRole,
			Priority:       "normal",
			AcceptanceJSON: []byte(buildCouncilTaskAcceptanceJSON(councilRun, councilInput, reviewerRole)),
			CreatedAt:      now,
			UpdatedAt:      now,
		}

		_, err = tx.ExecContext(
			ctx,
			`INSERT INTO tasks (
				run_id, task_id, title, summary, status, default_to, priority,
				acceptance_json, latest_attempt_no, created_at, updated_at
			) VALUES (?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?)`,
			task.RunID,
			task.TaskID,
			task.Title,
			task.Summary,
			task.Status,
			nullIfEmpty(task.DefaultTo),
			task.Priority,
			string(task.AcceptanceJSON),
			formatTime(task.CreatedAt),
			formatTime(task.UpdatedAt),
		)
		if err != nil {
			return CouncilStartResult{}, fmt.Errorf("insert council reviewer task: %w", err)
		}

		if err := insertEvent(ctx, tx, eventInput{
			RunID:       runID,
			TaskID:      taskID,
			Source:      "orch",
			EventType:   "task_added",
			Summary:     task.Title,
			PayloadJSON: marshalJSON(map[string]any{"title": task.Title, "priority": task.Priority}),
			CreatedAt:   now,
		}); err != nil {
			return CouncilStartResult{}, err
		}
		if err := insertEvent(ctx, tx, eventInput{
			RunID:       runID,
			TaskID:      taskID,
			Source:      "orch",
			EventType:   "task_ready",
			Summary:     task.Title,
			PayloadJSON: marshalJSON(map[string]any{"task_id": taskID}),
			CreatedAt:   now,
		}); err != nil {
			return CouncilStartResult{}, err
		}

		dispatchResult, finalizeWorkspace, err := s.dispatchTaskTx(
			ctx,
			tx,
			task,
			reviewerRole,
			buildCouncilTaskBody(councilRun, councilInput, reviewerRole),
			"",
			nil,
			now,
		)
		if err != nil {
			return CouncilStartResult{}, err
		}
		defer finalizeWorkspace(false)

		_, err = tx.ExecContext(
			ctx,
			`INSERT INTO council_reviewers (run_id, reviewer_role, task_id, status)
			 VALUES (?, ?, ?, ?)`,
			runID,
			reviewerRole,
			taskID,
			dispatchResult.Task.Status,
		)
		if err != nil {
			return CouncilStartResult{}, fmt.Errorf("insert council reviewer row: %w", err)
		}

		reviewers = append(reviewers, CouncilReviewer{
			ReviewerRole: reviewerRole,
			TaskID:       taskID,
			Status:       dispatchResult.Task.Status,
		})
	}

	if err := insertEvent(ctx, tx, eventInput{
		RunID:     runID,
		Source:    "orch",
		EventType: "council_started",
		Summary:   "council reviewers dispatched",
		PayloadJSON: marshalJSON(map[string]any{
			"mode":           councilRun.Mode,
			"target_type":    councilRun.TargetType,
			"output_mode":    councilRun.OutputMode,
			"only_unanimous": councilRun.OnlyUnanimous,
			"reviewers":      reviewers,
		}),
		CreatedAt: now,
	}); err != nil {
		return CouncilStartResult{}, err
	}

	if err := updateRunAggregateStatus(ctx, tx, runID, now); err != nil {
		return CouncilStartResult{}, err
	}

	if err := tx.Commit(); err != nil {
		return CouncilStartResult{}, fmt.Errorf("commit council start transaction: %w", err)
	}

	return CouncilStartResult{
		Run:       councilRun,
		Input:     councilInput,
		Reviewers: reviewers,
	}, nil
}

func normalizeCouncilRun(input CouncilStartInput) (CouncilRun, error) {
	mode := defaultString(strings.TrimSpace(input.Mode), "brainstorm")
	switch mode {
	case "brainstorm", "review":
	default:
		return CouncilRun{}, fmt.Errorf("%w: mode must be brainstorm or review", ErrInvalidInput)
	}

	targetType := defaultString(strings.TrimSpace(input.TargetType), "mixed")
	switch targetType {
	case "text", "repo", "mixed":
	default:
		return CouncilRun{}, fmt.Errorf("%w: target-type must be text, repo, or mixed", ErrInvalidInput)
	}

	outputMode := defaultString(strings.TrimSpace(input.OutputMode), "both")
	switch outputMode {
	case "markdown", "json", "both":
	default:
		return CouncilRun{}, fmt.Errorf("%w: output must be markdown, json, or both", ErrInvalidInput)
	}

	return CouncilRun{
		RunID:         strings.TrimSpace(input.RunID),
		Mode:          mode,
		TargetType:    targetType,
		OutputMode:    outputMode,
		OnlyUnanimous: input.OnlyUnanimous,
	}, nil
}

func normalizeCouncilInput(input CouncilStartInput) (CouncilInput, error) {
	result := CouncilInput{
		RunID:        strings.TrimSpace(input.RunID),
		Prompt:       strings.TrimSpace(input.Target),
		TargetFile:   strings.TrimSpace(input.TargetFile),
		RepoPath:     strings.TrimSpace(input.RepoPath),
		TargetTaskID: strings.TrimSpace(input.TargetTaskID),
	}

	if result.Prompt == "" && result.TargetFile == "" && result.RepoPath == "" && result.TargetTaskID == "" {
		return CouncilInput{}, fmt.Errorf("%w: at least one of target, target-file, repo-path, or task-id is required", ErrInvalidInput)
	}

	return result, nil
}

func buildCouncilRunGoal(input CouncilInput) string {
	switch {
	case input.Prompt != "":
		return "Council review: " + truncateSingleLine(input.Prompt, 80)
	case input.TargetTaskID != "":
		return "Council review for task " + input.TargetTaskID
	case input.TargetFile != "":
		return "Council review for " + input.TargetFile
	case input.RepoPath != "":
		return "Council review for repo " + input.RepoPath
	default:
		return "Council review"
	}
}

func buildCouncilRunSummary(run CouncilRun, input CouncilInput) string {
	return fmt.Sprintf("%s council (%s)", run.Mode, run.TargetType)
}

func buildCouncilTaskTitle(reviewerRole string) string {
	switch reviewerRole {
	case "architecture-reviewer":
		return "Council architecture review"
	case "implementation-reviewer":
		return "Council implementation review"
	case "risk-reviewer":
		return "Council risk review"
	default:
		return "Council review"
	}
}

func buildCouncilTaskSummary(reviewerRole string) string {
	switch reviewerRole {
	case "architecture-reviewer":
		return "Review the target for architecture, boundaries, and interfaces"
	case "implementation-reviewer":
		return "Review the target for simplicity, maintainability, and practicality"
	case "risk-reviewer":
		return "Review the target for regressions, correctness, and operability risks"
	default:
		return "Review the target"
	}
}

func buildCouncilTaskAcceptanceJSON(run CouncilRun, input CouncilInput, reviewerRole string) string {
	return marshalJSON(map[string]any{
		"mode": "analysis",
		"council": map[string]any{
			"reviewer_role":  reviewerRole,
			"council_mode":   run.Mode,
			"target_type":    run.TargetType,
			"output_mode":    run.OutputMode,
			"only_unanimous": run.OnlyUnanimous,
			"target": map[string]any{
				"prompt":      input.Prompt,
				"target_file": input.TargetFile,
				"repo_path":   input.RepoPath,
				"task_id":     input.TargetTaskID,
			},
			"response_format": map[string]any{
				"reviewer_role": reviewerRole,
				"findings": []map[string]any{
					{
						"title":       "string",
						"summary":     "string",
						"proposal":    "string",
						"rationale":   "string",
						"confidence":  "low|medium|high",
						"tags":        []string{},
						"target_refs": map[string]any{},
					},
				},
			},
		},
	})
}

func buildCouncilTaskBody(run CouncilRun, input CouncilInput, reviewerRole string) string {
	parts := []string{
		fmt.Sprintf("Reviewer role: %s", reviewerRole),
		fmt.Sprintf("Council mode: %s", run.Mode),
		fmt.Sprintf("Target type: %s", run.TargetType),
		"Analyze the target from your assigned reviewer perspective.",
		"Return structured findings with title, summary, proposal, rationale, confidence, tags, and optional target references.",
	}

	if input.Prompt != "" {
		parts = append(parts, "", "Prompt:", input.Prompt)
	}
	if input.TargetFile != "" {
		parts = append(parts, "", "Target file:", input.TargetFile)
	}
	if input.RepoPath != "" {
		parts = append(parts, "", "Repo path:", input.RepoPath)
	}
	if input.TargetTaskID != "" {
		parts = append(parts, "", "Related task id:", input.TargetTaskID)
	}

	return strings.Join(parts, "\n")
}

func truncateSingleLine(value string, maxLen int) string {
	value = strings.TrimSpace(value)
	value = strings.ReplaceAll(value, "\n", " ")
	value = strings.ReplaceAll(value, "\r", " ")
	value = strings.Join(strings.Fields(value), " ")
	if maxLen <= 0 || len(value) <= maxLen {
		return value
	}
	if maxLen <= 3 {
		return value[:maxLen]
	}
	return value[:maxLen-3] + "..."
}

func boolToInt(value bool) int {
	if value {
		return 1
	}
	return 0
}

func (s *OrchStore) WaitForCouncil(ctx context.Context, input CouncilWaitInput) (CouncilWaitResult, error) {
	runID := strings.TrimSpace(input.RunID)
	if runID == "" {
		return CouncilWaitResult{}, fmt.Errorf("%w: run id is required", ErrInvalidInput)
	}

	if _, err := s.GetCouncilRun(ctx, runID); err != nil {
		return CouncilWaitResult{}, err
	}

	waitCtx := ctx
	cancel := func() {}
	if input.Timeout > 0 {
		waitCtx, cancel = context.WithTimeout(ctx, input.Timeout)
	}
	defer cancel()

	for {
		reviewers, allComplete, err := s.GetCouncilReviewerStatuses(waitCtx, runID)
		if err != nil {
			if isDeadlineExceeded(waitCtx) {
				return CouncilWaitResult{
					Woke:             false,
					RunID:            runID,
					AllComplete:      false,
					ReviewerStatuses: reviewers,
				}, nil
			}
			return CouncilWaitResult{}, err
		}
		if allComplete {
			return CouncilWaitResult{
				Woke:             true,
				RunID:            runID,
				AllComplete:      true,
				ReviewerStatuses: reviewers,
			}, nil
		}

		if _, err := s.ReconcileRun(waitCtx, runID); err != nil {
			if isSQLiteBusyError(err) {
				ok, waitErr := waitForNextPoll(waitCtx, 25*time.Millisecond)
				if waitErr != nil {
					if errors.Is(waitErr, context.DeadlineExceeded) {
						reviewers, _, _ := s.GetCouncilReviewerStatuses(ctx, runID)
						return CouncilWaitResult{
							Woke:             false,
							RunID:            runID,
							AllComplete:      false,
							ReviewerStatuses: reviewers,
						}, nil
					}
					return CouncilWaitResult{}, waitErr
				}
				if !ok {
					reviewers, _, _ := s.GetCouncilReviewerStatuses(ctx, runID)
					return CouncilWaitResult{
						Woke:             false,
						RunID:            runID,
						AllComplete:      false,
						ReviewerStatuses: reviewers,
					}, nil
				}
				continue
			}
			if isDeadlineExceeded(waitCtx) {
				reviewers, _, _ := s.GetCouncilReviewerStatuses(ctx, runID)
				return CouncilWaitResult{
					Woke:             false,
					RunID:            runID,
					AllComplete:      false,
					ReviewerStatuses: reviewers,
				}, nil
			}
			return CouncilWaitResult{}, err
		}

		ok, err := waitForNextPoll(waitCtx, 200*time.Millisecond)
		if err != nil {
			if errors.Is(err, context.DeadlineExceeded) {
				reviewers, _, _ := s.GetCouncilReviewerStatuses(ctx, runID)
				return CouncilWaitResult{
					Woke:             false,
					RunID:            runID,
					AllComplete:      false,
					ReviewerStatuses: reviewers,
				}, nil
			}
			return CouncilWaitResult{}, err
		}
		if !ok {
			reviewers, _, _ := s.GetCouncilReviewerStatuses(ctx, runID)
			return CouncilWaitResult{
				Woke:             false,
				RunID:            runID,
				AllComplete:      false,
				ReviewerStatuses: reviewers,
			}, nil
		}
	}
}

func (s *OrchStore) GetCouncilRun(ctx context.Context, runID string) (CouncilRun, error) {
	row := s.db.QueryRowContext(
		ctx,
		`SELECT run_id, mode, target_type, output_mode, only_unanimous
		 FROM council_runs
		 WHERE run_id = ?`,
		runID,
	)

	var (
		run           CouncilRun
		onlyUnanimous int
	)
	err := row.Scan(&run.RunID, &run.Mode, &run.TargetType, &run.OutputMode, &onlyUnanimous)
	if errors.Is(err, sql.ErrNoRows) {
		return CouncilRun{}, fmt.Errorf("%w: council run %s not found", ErrRunNotFound, runID)
	}
	if err != nil {
		return CouncilRun{}, fmt.Errorf("scan council run: %w", err)
	}
	run.OnlyUnanimous = onlyUnanimous != 0
	return run, nil
}

func (s *OrchStore) GetCouncilReviewerStatuses(ctx context.Context, runID string) ([]CouncilReviewer, bool, error) {
	rows, err := s.db.QueryContext(
		ctx,
		`SELECT cr.reviewer_role, cr.task_id, t.status
		 FROM council_reviewers cr
		 JOIN tasks t
		   ON t.run_id = cr.run_id
		  AND t.task_id = cr.task_id
		 WHERE cr.run_id = ?
		 ORDER BY cr.reviewer_role ASC`,
		runID,
	)
	if err != nil {
		return nil, false, fmt.Errorf("query council reviewer statuses: %w", err)
	}
	defer rows.Close()

	reviewers := make([]CouncilReviewer, 0, len(councilReviewerRoles))
	allComplete := true
	for rows.Next() {
		var reviewer CouncilReviewer
		if err := rows.Scan(&reviewer.ReviewerRole, &reviewer.TaskID, &reviewer.Status); err != nil {
			return nil, false, fmt.Errorf("scan council reviewer status: %w", err)
		}
		if reviewer.Status != "done" && reviewer.Status != "failed" && reviewer.Status != "cancelled" {
			allComplete = false
		}
		reviewers = append(reviewers, reviewer)
	}
	if err := rows.Err(); err != nil {
		return nil, false, fmt.Errorf("iterate council reviewer statuses: %w", err)
	}
	if len(reviewers) == 0 {
		return nil, false, fmt.Errorf("%w: council reviewers for run %s not found", ErrRunNotFound, runID)
	}

	return reviewers, allComplete, nil
}