promptkit/internal/usecase/smoke_authentication.go

package usecase

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/generation"
	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
)

type authenticationProbe struct {
	client smokeHTTPGetter
}

type authenticationProbeOutcome struct {
	supported           bool
	applicable          *bool
	reason              domain.SmokeTestReason
	evidence            llm.HTTPExecutionEvidence
	credentialsAccepted *bool
	failure             domain.ExecutionFailureKind
}

func (p authenticationProbe) run(
	ctx context.Context,
	backend domain.Backend,
	target domain.SmokeTestTarget,
) (authenticationProbeOutcome, error) {
	outcome, err := p.support(backend)
	if err != nil || !outcome.supported || (outcome.applicable != nil && !*outcome.applicable) {
		return outcome, err
	}
	return p.runOpenRouterKey(ctx, target)
}

func (p authenticationProbe) support(backend domain.Backend) (authenticationProbeOutcome, error) {
	if backend.SmokeTests == nil || backend.SmokeTests.Authentication == nil {
		return authenticationProbeOutcome{reason: domain.SmokeTestReasonProbeNotConfigured}, nil
	}
	authentication := backend.SmokeTests.Authentication
	if authentication.Kind == nil {
		return authenticationProbeOutcome{}, fmt.Errorf("%w: authentication probe kind is missing", ErrInvalidRequest)
	}
	switch *authentication.Kind {
	case "none":
		return authenticationProbeOutcome{
			supported:  true,
			applicable: boolPointer(false),
			reason:     domain.SmokeTestReasonAuthenticationNotApplicable,
		}, nil
	case "openrouter_key":
		if p.client == nil {
			return authenticationProbeOutcome{reason: domain.SmokeTestReasonHTTPDiagnosticsUnavailable}, nil
		}
		return authenticationProbeOutcome{supported: true, applicable: boolPointer(true)}, nil
	default:
		return authenticationProbeOutcome{}, fmt.Errorf("%w: unsupported authentication probe kind %q", ErrInvalidRequest, *authentication.Kind)
	}
}

func (p authenticationProbe) runOpenRouterKey(
	ctx context.Context,
	target domain.SmokeTestTarget,
) (authenticationProbeOutcome, error) {
	endpoint, err := endpointWithoutUserInfo(target.Endpoint)
	if err != nil {
		return authenticationProbeOutcome{}, fmt.Errorf("prepare authentication endpoint: %w", err)
	}
	response, evidence, requestErr := p.client.GetWithAuthenticationEvidence(
		ctx,
		endpoint,
		target.EffectiveSettings.TimeoutSeconds,
		target.EffectiveSettings,
	)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}
	outcome := authenticationProbeOutcome{
		supported:  true,
		applicable: boolPointer(true),
		evidence:   evidence,
	}
	if evidence.ResponseReceived || response != nil {
		status := evidence.HTTPStatus
		if status == nil && response != nil {
			status = &response.StatusCode
		}
		if status == nil {
			outcome.failure = domain.ExecutionFailureMalformedResponse
			outcome.reason = domain.SmokeTestReasonAuthenticationResponseUnrecognized
			return outcome, nil
		}
		switch *status {
		case http.StatusUnauthorized:
			outcome.credentialsAccepted = boolPointer(false)
		case http.StatusOK:
			if response == nil || response.Body == nil {
				outcome.reason = domain.SmokeTestReasonAuthenticationResponseUnrecognized
				outcome.failure = domain.ExecutionFailureMalformedResponse
				return outcome, nil
			}
			recognized, bodyErr := recognizedOpenRouterKeyResponse(response.Body)
			if bodyErr != nil {
				outcome.failure = generation.ClassifyFailure(fmt.Errorf("%w: %w", llm.ErrTransport, bodyErr))
				return outcome, nil
			}
			if recognized {
				outcome.credentialsAccepted = boolPointer(true)
			} else {
				outcome.reason = domain.SmokeTestReasonAuthenticationResponseUnrecognized
				outcome.failure = domain.ExecutionFailureMalformedResponse
			}
		default:
			outcome.failure = generation.ClassifyFailure(llm.ErrUnexpectedStatus)
		}
		return outcome, nil
	}

	if evidence.DispatchInvoked {
		if requestErr != nil {
			outcome.failure = generation.ClassifyFailure(fmt.Errorf("%w: %w", llm.ErrTransport, requestErr))
		} else {
			outcome.failure = generation.ClassifyFailure(llm.ErrTransport)
		}
		return outcome, nil
	}
	if requestErr != nil {
		return outcome, fmt.Errorf("execute authentication request: %w", requestErr)
	}
	outcome.failure = generation.ClassifyFailure(llm.ErrTransport)
	return outcome, nil
}

func recognizedOpenRouterKeyResponse(body io.Reader) (bool, error) {
	contents, err := io.ReadAll(body)
	if err != nil {
		return false, err
	}
	var envelope map[string]json.RawMessage
	if err := json.Unmarshal(contents, &envelope); err != nil || envelope == nil {
		return false, nil
	}

	dataJSON, ok := envelope["data"]
	if !ok || !isJSONObject(dataJSON) {
		return false, nil
	}
	var data map[string]json.RawMessage
	if err := json.Unmarshal(dataJSON, &data); err != nil || data == nil {
		return false, nil
	}
	label, labelOK := data["label"]
	freeTier, freeTierOK := data["is_free_tier"]
	usage, usageOK := data["usage"]
	if !labelOK || !isJSONString(label) || !freeTierOK || !isJSONBoolean(freeTier) || !usageOK {
		return false, nil
	}
	var number any
	usageDecoder := json.NewDecoder(bytes.NewReader(usage))
	usageDecoder.UseNumber()
	if usageDecoder.Decode(&number) != nil {
		return false, nil
	}
	_, ok = number.(json.Number)
	return ok, nil
}

func isJSONObject(value json.RawMessage) bool {
	value = bytes.TrimSpace(value)
	return len(value) > 0 && value[0] == '{'
}

func isJSONString(value json.RawMessage) bool {
	value = bytes.TrimSpace(value)
	if len(value) == 0 || value[0] != '"' {
		return false
	}
	var parsed string
	return json.Unmarshal(value, &parsed) == nil
}

func isJSONBoolean(value json.RawMessage) bool {
	value = bytes.TrimSpace(value)
	return bytes.Equal(value, []byte("true")) || bytes.Equal(value, []byte("false"))
}