openteam/server/internal/channel/ops.go

package channel

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"
	"gorm.io/gorm"

	"openteam/server/internal/pkg/crypto"
	"openteam/server/internal/store"
)

// EncryptKey encrypts an upstream API key for storage.
func (s *Service) EncryptKey(plain string) (string, error) {
	return crypto.Encrypt(plain, s.master)
}

// ImportModels fetches the channel's model list (GET /v1/models) and creates
// global Model records plus bindings. Returns the imported model names.
func (s *Service) ImportModels(ch *store.Channel) ([]string, error) {
	key, err := s.DecryptKey(ch.APIKeyEnc)
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	base := strings.TrimSuffix(ch.BaseURL, "/")
	url := base + "/v1/models"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Accept", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("channel returned %d: %s", resp.StatusCode, truncate(string(raw), 200))
	}

	var payload struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.Unmarshal(raw, &payload); err != nil {
		return nil, err
	}
	if len(payload.Data) == 0 {
		return nil, errors.New("no models returned")
	}

	imported := []string{}
	err = s.db.Transaction(func(tx *gorm.DB) error {
		for _, m := range payload.Data {
			if m.ID == "" {
				continue
			}
			var model store.Model
			err := tx.Where("name = ?", m.ID).First(&model).Error
			if errors.Is(err, gorm.ErrRecordNotFound) {
				model = store.Model{Name: m.ID, DisplayName: m.ID}
				if err := tx.Create(&model).Error; err != nil {
					return err
				}
			} else if err != nil {
				return err
			}
			var binding store.ChannelModelBinding
			err = tx.Where("channel_id = ? AND model_id = ?", ch.ID, model.ID).First(&binding).Error
			if errors.Is(err, gorm.ErrRecordNotFound) {
				binding = store.ChannelModelBinding{
					ChannelID: ch.ID, ModelID: model.ID, UpstreamModel: m.ID, Weight: 1,
				}
				if err := tx.Create(&binding).Error; err != nil {
					return err
				}
			}
			imported = append(imported, m.ID)
		}
		return nil
	})
	return imported, err
}

// Test sends a cheap test request to the channel and reports success + latency.
func (s *Service) Test(ch *store.Channel) (bool, int, string, error) {
	key, err := s.DecryptKey(ch.APIKeyEnc)
	if err != nil {
		return false, 0, "", err
	}
	model := s.cfg.HealthCheck.TestModel
	if model == "" {
		// Pick the first bound model, if any.
		var b store.ChannelModelBinding
		if err := s.db.Where("channel_id = ?", ch.ID).First(&b).Error; err == nil {
			model = b.UpstreamModel
		}
	}
	if model == "" {
		return false, 0, "", errors.New("no test model configured; set HEALTHCHECK_TEST_MODEL or bind a model")
	}

	timeout := time.Duration(s.cfg.HealthCheck.TimeoutMs) * time.Millisecond
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// Test the first endpoint the channel actually serves, so channels that
	// omit chat completions (e.g. Anthropic-only) still get a valid probe.
	testFmt := pickTestFormat(ch)
	var body []byte
	switch testFmt {
	case store.FormatAnthropic:
		body, _ = json.Marshal(map[string]any{
			"model": model, "max_tokens": 8,
			"messages": []any{map[string]any{"role": "user", "content": "ping"}},
		})
	case store.FormatOpenAIResponses:
		body, _ = json.Marshal(map[string]any{"model": model, "input": "ping"})
	default:
		body, _ = json.Marshal(map[string]any{
			"model": model,
			"messages": []any{map[string]any{"role": "user", "content": "ping"}},
			"max_tokens": 8,
		})
	}

	url := strings.TrimSuffix(ch.BaseURL, "/") + healthCheckPath(testFmt)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(body)))
	if err != nil {
		return false, 0, "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	if testFmt == store.FormatAnthropic {
		req.Header.Set("anthropic-version", "2023-06-01")
	}

	start := time.Now()
	resp, err := http.DefaultClient.Do(req)
	latency := int(time.Since(start).Milliseconds())
	if err != nil {
		return false, latency, "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return false, latency, "", fmt.Errorf("test request failed: %d %s", resp.StatusCode, truncate(string(raw), 160))
	}
	return true, latency, model, nil
}

// MarkHealthy updates health status and resets the failure counter.
func (s *Service) MarkHealthy(chID int64, status string) {
	s.db.Model(&store.Channel{}).Where("id = ?", chID).
		Updates(map[string]any{"health_status": status, "health_failures": 0})
}

// MarkFailure increments the failure counter and sets cooldown when maxed.
func (s *Service) MarkFailure(chID int64) {
	var ch store.Channel
	if err := s.db.First(&ch, chID).Error; err != nil {
		return
	}
	failures := ch.HealthFailures + 1
	status := "degraded"
	if failures >= s.cfg.HealthCheck.MaxFailures {
		status = "cooldown"
	}
	s.db.Model(&ch).Updates(map[string]any{"health_failures": failures, "health_status": status})
	s.log.Info("channel health failure", zap.Int64("channel_id", chID),
		zap.Int("failures", failures), zap.String("status", status))
}

// pickTestFormat chooses a supported format to probe, preferring the cheapest
// endpoint. Falls back to chat completions so legacy channels keep working.
func pickTestFormat(ch *store.Channel) string {
	for _, want := range []string{store.FormatOpenAIChat, store.FormatAnthropic, store.FormatOpenAIResponses} {
		for _, f := range ch.FormatsResolved() {
			if f == want {
				return want
			}
		}
	}
	return store.FormatOpenAIChat
}

func healthCheckPath(f string) string {
	switch f {
	case store.FormatAnthropic:
		return "/v1/messages"
	case store.FormatOpenAIResponses:
		return "/v1/responses"
	default:
		return "/v1/chat/completions"
	}
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}