Agent-System-Dev/tests/ops/drill_driver/main.go

// 演练驱动：对隔离实例发真实 HTTP，记录每次请求的结果。
// 连接失败和超时记为 status=0，不把它们当成业务成功。
package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/cookiejar"
	"os"
	"strconv"
	"strings"
	"time"
)

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(2)
	}
	var err error
	switch os.Args[1] {
	case "prepare":
		err = cmdPrepare(os.Args[2:])
	case "timeline":
		err = cmdTimeline(os.Args[2:])
	case "idem-verify":
		err = cmdIdem(os.Args[2:])
	case "sse-check":
		err = cmdSSE(os.Args[2:])
	case "markers":
		err = cmdMarkers(os.Args[2:])
	case "artifacts-put":
		err = cmdArtifactsPut(os.Args[2:])
	case "artifacts-verify":
		err = cmdArtifactsVerify(os.Args[2:])
	case "burst":
		err = cmdBurst(os.Args[2:])
	case "perm-check":
		err = cmdPerm(os.Args[2:])
	case "writes":
		err = cmdWrites(os.Args[2:])
	case "-h", "--help", "help":
		usage()
		return
	default:
		usage()
		os.Exit(2)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprintln(os.Stderr, "drill_driver prepare|timeline|idem-verify|sse-check|markers|artifacts-put|artifacts-verify|burst|perm-check|writes")
}

type session struct {
	base   string
	client *http.Client
}

func newSession(base, owner, password string) (*session, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}
	s := &session{
		base:   strings.TrimRight(base, "/"),
		client: &http.Client{Timeout: 15 * time.Second, Jar: jar},
	}
	_, err = s.do("POST", "/api/auth/login", map[string]string{"owner": owner, "password": password})
	if err != nil {
		return nil, fmt.Errorf("login %s: %w", base, err)
	}
	return s, nil
}

type callResult struct {
	Status    int             `json:"status"`
	Elapsed   time.Duration   `json:"-"`
	ElapsedMs float64         `json:"elapsedMs"`
	Body      json.RawMessage `json:"body,omitempty"`
	Error     string          `json:"error,omitempty"`
}

func (s *session) do(method, path string, body any) (callResult, error) {
	var rd io.Reader
	if body != nil {
		buf, err := json.Marshal(body)
		if err != nil {
			return callResult{}, err
		}
		rd = bytes.NewReader(buf)
	}
	req, err := http.NewRequest(method, s.base+path, rd)
	if err != nil {
		return callResult{}, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	start := time.Now()
	resp, err := s.client.Do(req)
	elapsed := time.Since(start)
	if err != nil {
		return callResult{Elapsed: elapsed, ElapsedMs: ms(elapsed), Error: err.Error()}, nil
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	out := callResult{Status: resp.StatusCode, Elapsed: elapsed, ElapsedMs: ms(elapsed), Body: raw}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return out, fmt.Errorf("%s %s: status %d", method, path, resp.StatusCode)
	}
	return out, nil
}

func ms(d time.Duration) float64 {
	return float64(d) / float64(time.Millisecond)
}

func writeJSON(path string, v any) error {
	data, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	data = append(data, '\n')
	if path == "" || path == "-" {
		_, err = os.Stdout.Write(data)
		return err
	}
	return os.WriteFile(path, data, 0o644)
}

func cmdPrepare(args []string) error {
	fs := flag.NewFlagSet("prepare", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	out := fs.String("out", "-", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	proj, err := s.do("POST", "/api/projects", map[string]any{"name": "drill"})
	if err != nil {
		return err
	}
	var projBody struct {
		ID string `json:"id"`
	}
	if err := json.Unmarshal(proj.Body, &projBody); err != nil {
		return err
	}
	conv, err := s.do("POST", "/api/projects/"+projBody.ID+"/conversations", map[string]any{"title": "drill"})
	if err != nil {
		return err
	}
	var convBody struct {
		ID string `json:"id"`
	}
	if err := json.Unmarshal(conv.Body, &convBody); err != nil {
		return err
	}
	return writeJSON(*out, map[string]string{"projectId": projBody.ID, "conversationId": convBody.ID})
}

type attempt struct {
	RequestID string  `json:"requestId"`
	Base      string  `json:"base"`
	Started   string  `json:"started"`
	ElapsedMs float64 `json:"elapsedMs"`
	Status    int     `json:"status"`
	Error     string  `json:"error,omitempty"`
	MessageID string  `json:"messageId,omitempty"`
}

func cmdTimeline(args []string) error {
	fs := flag.NewFlagSet("timeline", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	conv := fs.String("conv", "", "conversation id")
	every := fs.Duration("every", 500*time.Millisecond, "write interval")
	dur := fs.Duration("dur", 90*time.Second, "how long to write")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	var attempts []attempt
	deadline := time.Now().Add(*dur)
	for i := 0; time.Now().Before(deadline); i++ {
		reqID := fmt.Sprintf("tl-%d-%d", time.Now().UnixNano(), i)
		started := time.Now().UTC()
		res, _ := s.do("POST", "/api/conversations/"+*conv+"/messages", map[string]any{
			"text": "drill " + reqID, "requestId": reqID,
		})
		attempts = append(attempts, attempt{
			RequestID: reqID, Base: *base, Started: started.Format(time.RFC3339Nano),
			ElapsedMs: res.ElapsedMs, Status: res.Status, Error: res.Error, MessageID: messageID(res.Body),
		})
		time.Sleep(*every)
	}
	return writeJSON(*out, map[string]any{"attempts": attempts})
}

func messageID(raw json.RawMessage) string {
	var body struct {
		Message struct {
			ID string `json:"id"`
		} `json:"message"`
	}
	_ = json.Unmarshal(raw, &body)
	return body.Message.ID
}

func cmdIdem(args []string) error {
	fs := flag.NewFlagSet("idem-verify", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	conv := fs.String("conv", "", "conversation id")
	in := fs.String("in", "", "timeline json")
	out := fs.String("out", "", "json output")
	since := fs.String("since", "", "only replay attempts at or after this time")
	until := fs.String("until", "", "only replay attempts at or before this time")
	retry := fs.Duration("retry", 2*time.Minute, "retry 409 until the conversation is free")
	if err := fs.Parse(args); err != nil {
		return err
	}
	data, err := os.ReadFile(*in)
	if err != nil {
		return err
	}
	var doc struct {
		Attempts []attempt `json:"attempts"`
	}
	if err := json.Unmarshal(data, &doc); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	var pending []attempt
	for _, a := range doc.Attempts {
		if a.Status >= 200 && a.Status < 300 {
			continue
		}
		if !attemptInWindow(a.Started, *since, *until) {
			continue
		}
		pending = append(pending, a)
	}
	var replays []attempt
	deadline := time.Now().Add(*retry)
	for len(pending) > 0 {
		var again []attempt
		for _, a := range pending {
			res, _ := s.do("POST", "/api/conversations/"+*conv+"/messages", map[string]any{
				"text": "drill " + a.RequestID, "requestId": a.RequestID,
			})
			// 409 发生在写入之前：会话仍被 mock run 占用，requestId 还没落库。
			if res.Status == http.StatusConflict && time.Now().Before(deadline) {
				again = append(again, a)
				continue
			}
			replays = append(replays, attempt{
				RequestID: a.RequestID, Base: *base, Started: time.Now().UTC().Format(time.RFC3339Nano),
				ElapsedMs: res.ElapsedMs, Status: res.Status, Error: res.Error, MessageID: messageID(res.Body),
			})
		}
		pending = again
		if len(pending) > 0 {
			time.Sleep(time.Second)
		}
	}
	return writeJSON(*out, map[string]any{"replays": replays})
}

func parseStamp(s string) (time.Time, error) {
	if t, err := time.Parse(time.RFC3339Nano, s); err == nil {
		return t, nil
	}
	return time.Parse(time.RFC3339, s)
}

func attemptInWindow(started, since, until string) bool {
	t, err := parseStamp(started)
	if err != nil {
		return false
	}
	if since != "" {
		s, err := parseStamp(since)
		if err != nil || t.Before(s) {
			return false
		}
	}
	if until != "" {
		u, err := parseStamp(until)
		if err != nil || t.After(u) {
			return false
		}
	}
	return true
}

func cmdSSE(args []string) error {
	fs := flag.NewFlagSet("sse-check", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	conv := fs.String("conv", "", "conversation id")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.base+"/api/events?conversationId="+*conv, nil)
	if err != nil {
		return err
	}
	stream := &http.Client{Jar: s.client.Jar}
	resp, err := stream.Do(req)
	var seqs []int64
	status := 0
	errText := ""
	if err != nil && resp == nil {
		errText = err.Error()
	} else if resp != nil {
		status = resp.StatusCode
		defer resp.Body.Close()
		body, readErr := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
		if readErr != nil {
			errText = readErr.Error()
		}
		for _, line := range strings.Split(string(body), "\n") {
			if !strings.HasPrefix(line, "id: ") {
				continue
			}
			n, convErr := strconv.ParseInt(strings.TrimPrefix(line, "id: "), 10, 64)
			if convErr == nil {
				seqs = append(seqs, n)
			}
		}
	}
	return writeJSON(*out, map[string]any{
		"status": status, "error": errText, "seqs": seqs, "holes": seqHoles(seqs),
	})
}

func cmdMarkers(args []string) error {
	fs := flag.NewFlagSet("markers", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	conv := fs.String("conv", "", "conversation id")
	every := fs.Duration("every", 5*time.Second, "marker interval")
	dur := fs.Duration("dur", 5*time.Minute, "how long to write")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	var marks []map[string]any
	deadline := time.Now().Add(*dur)
	for i := 0; time.Now().Before(deadline); i++ {
		reqID := fmt.Sprintf("mk-%d", i)
		acked := time.Now().UTC()
		res, _ := s.do("POST", "/api/conversations/"+*conv+"/messages", map[string]any{
			"text": "marker " + reqID, "requestId": reqID,
		})
		marks = append(marks, map[string]any{
			"requestId": reqID, "acked": acked.Format(time.RFC3339Nano),
			"status": res.Status, "error": res.Error, "messageId": messageID(res.Body),
		})
		time.Sleep(*every)
	}
	return writeJSON(*out, map[string]any{"markers": marks})
}

func cmdArtifactsPut(args []string) error {
	fs := flag.NewFlagSet("artifacts-put", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	project := fs.String("project", "", "project id")
	n := fs.Int("n", 5, "how many files")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	var items []map[string]any
	for i := 0; i < *n; i++ {
		payload := []byte(fmt.Sprintf("drill-artifact-%d-%d", i, time.Now().UnixNano()))
		sum := sha256.Sum256(payload)
		id, status, errText, acked := upload(s, *project, fmt.Sprintf("drill-%d.txt", i), payload)
		items = append(items, map[string]any{
			"name": fmt.Sprintf("drill-%d.txt", i), "sha256": hex.EncodeToString(sum[:]),
			"id": id, "status": status, "error": errText, "acked": acked.Format(time.RFC3339Nano),
		})
	}
	return writeJSON(*out, map[string]any{"artifacts": items})
}

func upload(s *session, project, name string, payload []byte) (id string, status int, errText string, acked time.Time) {
	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	part, err := w.CreateFormFile("file", name)
	if err != nil {
		return "", 0, err.Error(), time.Now().UTC()
	}
	_, _ = part.Write(payload)
	_ = w.Close()
	req, err := http.NewRequest(http.MethodPost, s.base+"/api/projects/"+project+"/uploads", &buf)
	if err != nil {
		return "", 0, err.Error(), time.Now().UTC()
	}
	req.Header.Set("Content-Type", w.FormDataContentType())
	acked = time.Now().UTC()
	resp, err := s.client.Do(req)
	if err != nil {
		return "", 0, err.Error(), acked
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return artifactIDFromUpload(raw), resp.StatusCode, string(raw), acked
	}
	return artifactIDFromUpload(raw), resp.StatusCode, "", acked
}

func artifactIDFromUpload(raw []byte) string {
	var body struct {
		ID         string `json:"id"`
		ArtifactID string `json:"artifactId"`
	}
	_ = json.Unmarshal(raw, &body)
	if body.ID != "" {
		return body.ID
	}
	return body.ArtifactID
}

func cmdArtifactsVerify(args []string) error {
	fs := flag.NewFlagSet("artifacts-verify", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	in := fs.String("in", "", "artifacts json")
	before := fs.String("before", "", "only artifacts acked before this RFC3339 time")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	data, err := os.ReadFile(*in)
	if err != nil {
		return err
	}
	var doc struct {
		Artifacts []struct {
			ID     string `json:"id"`
			SHA256 string `json:"sha256"`
			Acked  string `json:"acked"`
			Status int    `json:"status"`
		} `json:"artifacts"`
	}
	if err := json.Unmarshal(data, &doc); err != nil {
		return err
	}
	cutoff, err := time.Parse(time.RFC3339Nano, *before)
	if err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	missing := 0
	mismatch := 0
	checked := 0
	for _, item := range doc.Artifacts {
		acked, err := time.Parse(time.RFC3339Nano, item.Acked)
		if err != nil || item.Status < 200 || item.Status >= 300 || !acked.Before(cutoff) {
			continue
		}
		checked++
		req, err := http.NewRequest(http.MethodGet, s.base+"/api/artifacts/"+item.ID+"/content", nil)
		if err != nil {
			missing++
			continue
		}
		resp, err := s.client.Do(req)
		if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 {
			if resp != nil {
				resp.Body.Close()
			}
			missing++
			continue
		}
		raw, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		sum := sha256.Sum256(raw)
		if hex.EncodeToString(sum[:]) != item.SHA256 {
			mismatch++
		}
	}
	return writeJSON(*out, map[string]any{"checked": checked, "missing": missing, "mismatch": mismatch})
}

func cmdBurst(args []string) error {
	fs := flag.NewFlagSet("burst", flag.ContinueOnError)
	bases := fs.String("bases", "", "comma-separated base URLs")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	n := fs.Int("n", 40, "requests per base")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	counts := map[string]int{}
	for _, base := range strings.Split(*bases, ",") {
		base = strings.TrimSpace(base)
		if base == "" {
			continue
		}
		s, err := newSession(base, *owner, *password)
		if err != nil {
			return err
		}
		for i := 0; i < *n; i++ {
			res, _ := s.do("GET", "/api/projects", nil)
			key := fmt.Sprintf("%s:%d", base, res.Status)
			if res.Status == 0 {
				key = base + ":error"
			}
			counts[key]++
		}
	}
	return writeJSON(*out, map[string]any{"counts": counts})
}

func cmdWrites(args []string) error {
	fs := flag.NewFlagSet("writes", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "login owner")
	password := fs.String("password", "", "login password")
	conv := fs.String("conv", "", "conversation id")
	project := fs.String("project", "", "project id used when -fresh creates a conversation per write")
	fresh := fs.Bool("fresh", false, "one new conversation per write, so a live run cannot 409 the probe")
	n := fs.Int("n", 5, "consecutive successful writes required")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if *fresh && *project == "" {
		return fmt.Errorf("writes -fresh requires -project")
	}
	if !*fresh && *conv == "" {
		return fmt.Errorf("writes requires -conv")
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	for i := 0; i < *n; i++ {
		target := *conv
		if *fresh {
			created, err := s.do("POST", "/api/projects/"+*project+"/conversations", map[string]any{
				"title": fmt.Sprintf("recover-%d", i),
			})
			if err != nil {
				return fmt.Errorf("write %d conversation: %w", i, err)
			}
			var body struct {
				ID string `json:"id"`
			}
			if err := json.Unmarshal(created.Body, &body); err != nil || body.ID == "" {
				return fmt.Errorf("write %d conversation id missing", i)
			}
			target = body.ID
		}
		reqID := fmt.Sprintf("wr-%d-%d", time.Now().UnixNano(), i)
		res, err := s.do("POST", "/api/conversations/"+target+"/messages", map[string]any{
			"text": "recover " + reqID, "requestId": reqID,
		})
		if err != nil || res.Status < 200 || res.Status >= 300 {
			return fmt.Errorf("write %d status %d: %s", i, res.Status, res.Error)
		}
	}
	fmt.Printf("writes ok n=%d\n", *n)
	return nil
}

func cmdPerm(args []string) error {
	fs := flag.NewFlagSet("perm-check", flag.ContinueOnError)
	base := fs.String("base", "", "instance base URL")
	owner := fs.String("owner", "", "other user")
	password := fs.String("password", "", "other user password")
	project := fs.String("project", "", "someone else's project")
	artifact := fs.String("artifact", "", "someone else's artifact")
	out := fs.String("out", "", "json output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	s, err := newSession(*base, *owner, *password)
	if err != nil {
		return err
	}
	proj, _ := s.do("GET", "/api/projects/"+*project+"/conversations", nil)
	art, _ := s.do("GET", "/api/artifacts/"+*artifact+"/content", nil)
	denied := func(status int) bool { return status == http.StatusForbidden || status == http.StatusNotFound }
	return writeJSON(*out, map[string]any{
		"projectStatus": proj.Status, "artifactStatus": art.Status,
		"projectDenied": denied(proj.Status), "artifactDenied": denied(art.Status),
	})
}