internal/agent/stream.go - master - t3/cli

package agent

import (
	"fmt"
	"io"
	"log/slog"

	"gitverse.ru/t3/cli/internal/lifecycle"
	"gitverse.ru/t3/cli/pkg/pb"
)

func (ag *Agent) Stream(stream pb.AgentService_StreamServer) error {
	// Start the metrics reporter for this stream connection
	startMetricsReporter(stream.Context(), stream, ag.registry, slog.Default().With("hostname", ag.hostname))

	for {
		req, err := stream.Recv()
		if err == io.EOF {
			// client disconnected — clean up running test if any
			ag.launch.Stop()
			ag.clearRunID()
			return nil
		}
		if err != nil {
			// If a test is running, stop it and clear run_id before returning
			ag.launch.Stop()
			ag.clearRunID()
			return fmt.Errorf("client communication error: %w", err)
		}

		switch payload := req.Payload.(type) {
		case *pb.AgentRequest_Start:
			ag.executeStep(stream, payload.Start)

		case *pb.AgentRequest_Stop:
			err := ag.handleStop(stream, payload.Stop)
			if err != nil {
				slog.Error("stop failed", "error", err)
			}
			return nil
		}
	}
}

// handleStop validates the stop request and initiates test stop.
func (ag *Agent) handleStop(stream pb.AgentService_StreamServer, req *pb.StopRequest) error {
	if req.RunId == "" {
		return fmt.Errorf("missing run_id")
	}

	ag.runIDMu.Lock()
	if ag.currentRunID == "" {
		ag.runIDMu.Unlock()
		return fmt.Errorf("no running test to stop")
	}
	if ag.currentRunID != req.RunId {
		runID := ag.currentRunID
		ag.runIDMu.Unlock()
		return fmt.Errorf("run_id mismatch: expected %s, got %s", runID, req.RunId)
	}
	ag.runIDMu.Unlock()

	ag.launch.Stop()
	ag.clearRunID()

	slog.Info("test stopped", "run_id", req.RunId)
	_ = ag.SendStatus(stream, lifecycle.SUCCESS, "test stopped by request")
	return nil
}