iop/apps/edge/internal/edgecmd/smoke_openai.go

package edgecmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

func smokeCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "smoke",
		Short: "Run verification and smoke tests",
		Long: `Run diagnostic smoke tests to verify the connectivity and functionality of IOP Edge components.

Typical local flow:
  1. iop-edge config init
  2. iop-edge node register <node-id>
  3. iop-edge serve
  4. iop-edge smoke openai --model <model>`,
		Example: `  iop-edge smoke openai --model gemma4:26b
  iop-edge smoke openai --model gemma4:26b --base-url http://127.0.0.1:18081`,
	}
	c.AddCommand(smokeOpenAICmd())
	return c
}

var (
	smokeModel          string
	smokeBaseURL        string
	smokePrompt         string
	smokeTimeout        string
	smokeWorkspace      string
	smokeExpectFile     string
	smokeExpectContains string
	smokeAPIKey         string
	smokeAPIKeyFile     string
)

type openAIModel struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

type openAIModelsResponse struct {
	Object string        `json:"object"`
	Data   []openAIModel `json:"data"`
}

type responsesContentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

type responsesOutputItem struct {
	Type    string                 `json:"type"`
	Role    string                 `json:"role"`
	Content []responsesContentItem `json:"content"`
}

type responsesResponse struct {
	ID         string                `json:"id"`
	Object     string                `json:"object"`
	CreatedAt  int64                 `json:"created_at"`
	Model      string                `json:"model"`
	OutputText string                `json:"output_text"`
	Output     []responsesOutputItem `json:"output"`
}

func (r responsesResponse) text() string {
	if strings.TrimSpace(r.OutputText) != "" {
		return r.OutputText
	}
	for _, item := range r.Output {
		for _, content := range item.Content {
			if strings.TrimSpace(content.Text) != "" {
				return content.Text
			}
		}
	}
	return ""
}

func resolveSmokeBaseURL(cmd *cobra.Command) (string, error) {
	if smokeBaseURL != "" {
		return strings.TrimSuffix(smokeBaseURL, "/"), nil
	}
	cfg, _, err := LoadEdgeConfig(cmd)
	if err != nil {
		return "", fmt.Errorf("failed to load edge config to resolve base URL (pass --base-url to override): %w", err)
	}
	listen := cfg.OpenAI.Listen
	if listen == "" {
		listen = "127.0.0.1:18081"
	}
	host, port, err := net.SplitHostPort(listen)
	if err != nil {
		return "http://127.0.0.1:18081", nil
	}
	if host == "0.0.0.0" || host == "" {
		host = "127.0.0.1"
	}
	return "http://" + net.JoinHostPort(host, port), nil
}

func resolveSmokeAPIKey() (string, error) {
	if strings.TrimSpace(smokeAPIKey) != "" {
		return strings.TrimSpace(smokeAPIKey), nil
	}
	if strings.TrimSpace(smokeAPIKeyFile) != "" {
		data, err := os.ReadFile(smokeAPIKeyFile)
		if err != nil {
			return "", fmt.Errorf("read --api-key-file: %w", err)
		}
		return strings.TrimSpace(string(data)), nil
	}
	return strings.TrimSpace(os.Getenv("OPENAI_API_KEY")), nil
}

// openAISmokeRun holds the resolved endpoint context shared by the sequential
// smoke stages (health -> models -> responses -> expected file). Extracting the
// stages keeps smokeOpenAICmd to command construction while preserving the
// original output text, request order, and early-return behavior.
type openAISmokeRun struct {
	out     io.Writer
	client  *http.Client
	baseURL string
	model   string
	prompt  string
	apiKey  string
}

func smokeOpenAICmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "openai",
		Short: "Perform OpenAI-compatible endpoint smoke test",
		Long: `Validate the IOP Edge OpenAI-compatible HTTP API surface.
	It runs a three-step diagnostics check:
  1. Checks GET /healthz (server status).
  2. Checks GET /v1/models (advertised model compatibility).
  3. Checks POST /v1/responses (non-streaming Responses API request processing).

If --base-url is not provided, it auto-discovers the OpenAI URL from the effective config's listen address.`,
		Example: `  iop-edge smoke openai --model gemma4:26b
  iop-edge smoke openai --model gemma4:26b --base-url http://127.0.0.1:18081 --prompt "say hello" --timeout 15s`,
		RunE: func(cmd *cobra.Command, _ []string) error {
			return runOpenAISmoke(cmd)
		},
	}

	c.Flags().StringVar(&smokeModel, "model", "", "model identifier to check (required)")
	c.Flags().StringVar(&smokeBaseURL, "base-url", "", "override OpenAI base URL (e.g. http://127.0.0.1:18081)")
	c.Flags().StringVar(&smokePrompt, "prompt", "ping", "prompt text to send")
	c.Flags().StringVar(&smokeTimeout, "timeout", "10s", "client timeout duration")
	c.Flags().StringVar(&smokeWorkspace, "workspace", "", "path to the temporary workspace directory")
	c.Flags().StringVar(&smokeExpectFile, "expect-file", "", "relative path under workspace to check for existence after smoke")
	c.Flags().StringVar(&smokeExpectContains, "expect-contains", "", "substring expected to be found in the expected file")
	c.Flags().StringVar(&smokeAPIKey, "api-key", "", "bearer token for authenticated OpenAI-compatible endpoints (defaults to OPENAI_API_KEY)")
	c.Flags().StringVar(&smokeAPIKeyFile, "api-key-file", "", "file containing bearer token for authenticated OpenAI-compatible endpoints")
	c.MarkFlagRequired("model")

	return c
}

// runOpenAISmoke resolves the endpoint context, prints the header, and runs the
// smoke stages in their fixed order, stopping at the first stage error.
func runOpenAISmoke(cmd *cobra.Command) error {
	run, err := newOpenAISmokeRun(cmd)
	if err != nil {
		return err
	}
	if err := run.checkHealth(); err != nil {
		return err
	}
	if err := run.checkModels(); err != nil {
		return err
	}
	if err := run.checkResponses(); err != nil {
		return err
	}
	return run.verifyExpectedFile()
}

// newOpenAISmokeRun resolves base URL, timeout, API key, and prompt, then prints
// the smoke header before the sequential stages run.
func newOpenAISmokeRun(cmd *cobra.Command) (*openAISmokeRun, error) {
	baseURL, err := resolveSmokeBaseURL(cmd)
	if err != nil {
		return nil, err
	}

	timeoutDur := 10 * time.Second
	if smokeTimeout != "" {
		d, err := time.ParseDuration(smokeTimeout)
		if err != nil {
			return nil, fmt.Errorf("invalid timeout duration: %w", err)
		}
		timeoutDur = d
	}

	apiKey, err := resolveSmokeAPIKey()
	if err != nil {
		return nil, err
	}

	prompt := smokePrompt
	if prompt == "" {
		prompt = "ping"
	}

	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "IOP Edge OpenAI Smoke Test\n")
	fmt.Fprintf(out, "========================================\n")
	fmt.Fprintf(out, "Target Endpoint: %s\n", baseURL)
	fmt.Fprintf(out, "Target Model:    %s\n", smokeModel)
	fmt.Fprintf(out, "Timeout:         %s\n\n", timeoutDur)

	return &openAISmokeRun{
		out:     out,
		client:  &http.Client{Timeout: timeoutDur},
		baseURL: baseURL,
		model:   smokeModel,
		prompt:  prompt,
		apiKey:  apiKey,
	}, nil
}

// do issues an HTTP request, attaching the bearer token when one was resolved.
func (r *openAISmokeRun) do(req *http.Request) (*http.Response, error) {
	if r.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+r.apiKey)
	}
	return r.client.Do(req)
}

// checkHealth performs Step 1: GET /healthz and parse the JSON status body.
func (r *openAISmokeRun) checkHealth() error {
	fmt.Fprintf(r.out, "Step 1: Checking /healthz ... ")
	healthURL := r.baseURL + "/healthz"
	req, err := http.NewRequest(http.MethodGet, healthURL, nil)
	if err != nil {
		return fmt.Errorf("build GET %s: %w", healthURL, err)
	}
	resp, err := r.do(req)
	if err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("GET %s failed: %w (make sure the edge server is running)", healthURL, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("GET %s returned non-200 status %d: %s", healthURL, resp.StatusCode, string(body))
	}
	var healthData map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&healthData); err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("failed to parse /healthz JSON response: %w", err)
	}
	fmt.Fprintln(r.out, "[OK]")
	return nil
}

// checkModels performs Step 2: GET /v1/models and warn when the target model is
// not advertised.
func (r *openAISmokeRun) checkModels() error {
	fmt.Fprintf(r.out, "Step 2: Checking /v1/models ... ")
	modelsURL := r.baseURL + "/v1/models"
	req, err := http.NewRequest(http.MethodGet, modelsURL, nil)
	if err != nil {
		return fmt.Errorf("build GET %s: %w", modelsURL, err)
	}
	resp, err := r.do(req)
	if err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("GET %s failed: %w", modelsURL, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("GET %s returned non-200 status %d: %s", modelsURL, resp.StatusCode, string(body))
	}
	var modelsResp openAIModelsResponse
	if err := json.NewDecoder(resp.Body).Decode(&modelsResp); err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("failed to parse /v1/models JSON response: %w", err)
	}
	modelIDs := make([]string, 0, len(modelsResp.Data))
	modelMatched := false
	for _, m := range modelsResp.Data {
		modelIDs = append(modelIDs, m.ID)
		if m.ID == r.model {
			modelMatched = true
		}
	}
	fmt.Fprintf(r.out, "[OK] (found models: %s)\n", strings.Join(modelIDs, ", "))
	if !modelMatched {
		fmt.Fprintf(r.out, "WARNING: Target model %q is not advertised by /v1/models (advertised: %s)\n", r.model, strings.Join(modelIDs, ", "))
	}
	return nil
}

// checkResponses performs Step 3: POST /v1/responses and confirm non-empty text.
func (r *openAISmokeRun) checkResponses() error {
	fmt.Fprintf(r.out, "Step 3: Checking /v1/responses ... ")
	metadata := map[string]interface{}{
		"request_id": "iop-edge-smoke",
		"task_id":    "smoke",
	}
	if smokeWorkspace != "" {
		metadata["workspace"] = smokeWorkspace
	}

	responsesPayload := map[string]interface{}{
		"model":             r.model,
		"input":             r.prompt,
		"stream":            false,
		"metadata":          metadata,
		"max_output_tokens": 32,
		"temperature":       0,
	}
	payloadBytes, err := json.Marshal(responsesPayload)
	if err != nil {
		return fmt.Errorf("failed to marshal responses request payload: %w", err)
	}

	responsesURL := r.baseURL + "/v1/responses"
	req, err := http.NewRequest(http.MethodPost, responsesURL, bytes.NewReader(payloadBytes))
	if err != nil {
		return fmt.Errorf("build POST %s: %w", responsesURL, err)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := r.do(req)
	if err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("POST %s failed: %w (make sure active nodes are online)", responsesURL, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("POST %s returned non-200 status %d: %s", responsesURL, resp.StatusCode, string(body))
	}
	var responsesResp responsesResponse
	if err := json.NewDecoder(resp.Body).Decode(&responsesResp); err != nil {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("failed to parse /v1/responses JSON response: %w", err)
	}
	outputText := responsesResp.text()
	if strings.TrimSpace(outputText) == "" {
		fmt.Fprintln(r.out, "[FAILED]")
		return fmt.Errorf("POST %s returned empty response text", responsesURL)
	}
	fmt.Fprintf(r.out, "[OK]\n")
	fmt.Fprintf(r.out, "Responses Output Text: %q\n\n", outputText)
	return nil
}

// verifyExpectedFile checks an optional workspace-relative expected file and
// prints the final success line.
func (r *openAISmokeRun) verifyExpectedFile() error {
	if smokeExpectFile != "" {
		if smokeWorkspace == "" {
			return fmt.Errorf("--expect-file requires --workspace to be set")
		}
		expectFilePath := filepath.Join(smokeWorkspace, smokeExpectFile)
		rel, err := filepath.Rel(smokeWorkspace, expectFilePath)
		if err != nil || strings.HasPrefix(rel, "..") {
			return fmt.Errorf("expect-file %q must be within the workspace", smokeExpectFile)
		}
		data, err := os.ReadFile(expectFilePath)
		if err != nil {
			return fmt.Errorf("failed to read expected file %q: %w", expectFilePath, err)
		}
		if smokeExpectContains != "" {
			if !strings.Contains(string(data), smokeExpectContains) {
				return fmt.Errorf("expected file %q does not contain expected substring %q", smokeExpectFile, smokeExpectContains)
			}
		}
	}

	fmt.Fprintf(r.out, "IOP Edge OpenAI Smoke Test SUCCESS!\n")
	return nil
}