Agent-System/internal/controlcenter/tasks/settings.go

package tasks

import (
	"context"
	"fmt"

	"chat-agent-system/internal/controlcenter/pipeline"
	"chat-agent-system/internal/controlcenter/store"
)

type StageSetting struct {
	Scope    string  `json:"scope"`
	TaskID   *string `json:"task_id,omitempty"`
	Stage    string  `json:"stage"`
	Model    string  `json:"model"`
	Executor string  `json:"executor"`
}

// ListStageSettings returns global rows, or — when taskID is given — the
// task-level overrides merged on top of the globals, keyed by stage.
func ListStageSettings(ctx context.Context, st *store.Store, taskID string) ([]StageSetting, error) {
	rows, err := st.Pool.Query(ctx,
		`SELECT scope, task_id, stage, model, executor FROM cc_stage_settings ORDER BY stage`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	global := map[string]StageSetting{}
	overrides := map[string]StageSetting{}
	for rows.Next() {
		var s StageSetting
		if err := rows.Scan(&s.Scope, &s.TaskID, &s.Stage, &s.Model, &s.Executor); err != nil {
			return nil, err
		}
		if s.Scope == "global" {
			global[s.Stage] = s
		} else if taskID != "" && s.TaskID != nil && *s.TaskID == taskID {
			overrides[s.Stage] = s
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	out := []StageSetting{}
	for _, stage := range pipeline.Stages {
		if s, ok := overrides[stage]; ok {
			out = append(out, s)
			continue
		}
		if s, ok := global[stage]; ok {
			out = append(out, s)
		}
	}
	return out, nil
}

// ResolveStageSetting picks the effective setting for a stage: task override
// first, global fallback.
func ResolveStageSetting(ctx context.Context, st *store.Store, taskID, stage string) (*StageSetting, error) {
	for _, s := range []struct {
		scope  string
		taskID *string
	}{
		{"task", &taskID},
		{"global", nil},
	} {
		var out StageSetting
		var err error
		if s.taskID == nil {
			err = st.Pool.QueryRow(ctx,
				`SELECT scope, task_id, stage, model, executor FROM cc_stage_settings
				 WHERE scope = 'global' AND stage = $1`, stage).
				Scan(&out.Scope, &out.TaskID, &out.Stage, &out.Model, &out.Executor)
		} else {
			err = st.Pool.QueryRow(ctx,
				`SELECT scope, task_id, stage, model, executor FROM cc_stage_settings
				 WHERE scope = 'task' AND task_id = $1 AND stage = $2`, *s.taskID, stage).
				Scan(&out.Scope, &out.TaskID, &out.Stage, &out.Model, &out.Executor)
		}
		if err == nil {
			return &out, nil
		}
	}
	return nil, nil
}

// PutStageSetting upserts one stage -> {model, executor} mapping, globally or
// for a single task.
func PutStageSetting(ctx context.Context, st *store.Store, taskID *string, stage, model, executor string) error {
	if !pipeline.ValidStage(stage) {
		return fmt.Errorf("%w: invalid stage %q", ErrBadInput, stage)
	}
	if executor == "" {
		executor = "agent"
	}
	switch executor {
	case "agent", "human", "external":
	default:
		return fmt.Errorf("%w: executor must be agent, human or external", ErrBadInput)
	}
	scope := "global"
	if taskID != nil && *taskID != "" {
		scope = "task"
		if _, err := GetTask(ctx, st, *taskID); err != nil {
			return err
		}
	} else {
		taskID = nil
	}
	res, err := st.Pool.Exec(ctx,
		`UPDATE cc_stage_settings SET model = $4, executor = $5, updated_at = now()
		 WHERE scope = $1 AND stage = $3 AND (task_id IS NOT DISTINCT FROM $2)`,
		scope, taskID, stage, model, executor)
	if err != nil {
		return err
	}
	if res.RowsAffected() > 0 {
		return nil
	}
	_, err = st.Pool.Exec(ctx,
		`INSERT INTO cc_stage_settings (scope, task_id, stage, model, executor) VALUES ($1,$2,$3,$4,$5)`,
		scope, taskID, stage, model, executor)
	return err
}