alt/apps/cli/internal/operator/client.go

package operator

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"time"

	"google.golang.org/protobuf/proto"

	altv1 "git.toki-labs.com/toki/alt/packages/contracts/gen/go/alt/v1"
	protoSocket "git.toki-labs.com/toki/proto-socket/go"
)

// ErrTransport marks a connect/disconnect/timeout failure that happens before a
// typed ALT response is received. The runner maps it to the transport exit code
// so scripts can tell a dead socket apart from a typed application error.
var ErrTransport = errors.New("api transport error")

// APIClient is a thin proto-socket client for the ALT API control plane. It
// dials the same ws/wss rail the Flutter client and worker client use, with the
// CLI-local ParserMap so the operator runner can issue hello and market reads
// without importing the API's internal packages.
type APIClient struct {
	ws *protoSocket.WsClient
}

// Dial connects to the API at apiURL (ws:// or wss://). A connection failure is
// returned wrapped in ErrTransport so callers can distinguish it from a typed
// ErrorInfo response.
func Dial(ctx context.Context, apiURL string) (*APIClient, error) {
	host, port, path, secure, err := parseAPIURL(apiURL)
	if err != nil {
		return nil, err
	}

	var ws *protoSocket.WsClient
	if secure {
		ws, err = protoSocket.DialWss(ctx, host, port, path, nil, ParserMap())
	} else {
		ws, err = protoSocket.DialWs(ctx, host, port, path, ParserMap())
	}
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrTransport, err)
	}
	return &APIClient{ws: ws}, nil
}

// parseAPIURL splits an ws/wss URL into dial parameters, mirroring the worker
// client's parsing (services/api/internal/workerclient/client.go): default port
// 443 for wss and 80 for ws, and a "/" path when none is given.
func parseAPIURL(apiURL string) (host string, port int, path string, secure bool, err error) {
	u, err := url.Parse(apiURL)
	if err != nil {
		return "", 0, "", false, fmt.Errorf("invalid api URL %q: %w", apiURL, err)
	}
	if u.Scheme != "ws" && u.Scheme != "wss" {
		return "", 0, "", false, fmt.Errorf("api URL %q must use ws:// or wss://", apiURL)
	}
	secure = u.Scheme == "wss"

	host, portStr, splitErr := net.SplitHostPort(u.Host)
	if splitErr != nil {
		host = u.Host
		if secure {
			portStr = "443"
		} else {
			portStr = "80"
		}
	}
	if host == "" {
		return "", 0, "", false, fmt.Errorf("api URL %q is missing a host", apiURL)
	}

	port, err = strconv.Atoi(portStr)
	if err != nil {
		return "", 0, "", false, fmt.Errorf("invalid api port %q: %w", portStr, err)
	}

	path = u.Path
	if path == "" {
		path = "/"
	}
	return host, port, path, secure, nil
}

// Close releases the underlying socket. It is safe to call on a nil-socket
// client.
func (c *APIClient) Close() error {
	if c == nil || c.ws == nil {
		return nil
	}
	return c.ws.Close()
}

// Hello performs the API handshake.
func (c *APIClient) Hello(ctx context.Context, req *altv1.HelloRequest) (*altv1.HelloResponse, error) {
	return sendTyped[*altv1.HelloRequest, *altv1.HelloResponse](c, ctx, req)
}

// ListInstruments issues a market instrument list query.
func (c *APIClient) ListInstruments(ctx context.Context, req *altv1.ListInstrumentsRequest) (*altv1.ListInstrumentsResponse, error) {
	return sendTyped[*altv1.ListInstrumentsRequest, *altv1.ListInstrumentsResponse](c, ctx, req)
}

// ListBars issues a market bar list query.
func (c *APIClient) ListBars(ctx context.Context, req *altv1.ListBarsRequest) (*altv1.ListBarsResponse, error) {
	return sendTyped[*altv1.ListBarsRequest, *altv1.ListBarsResponse](c, ctx, req)
}

// StartBacktest starts a new backtest run.
func (c *APIClient) StartBacktest(ctx context.Context, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error) {
	return sendTyped[*altv1.StartBacktestRequest, *altv1.StartBacktestResponse](c, ctx, req)
}

// GetBacktestRun retrieves backtest run status/info.
func (c *APIClient) GetBacktestRun(ctx context.Context, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error) {
	return sendTyped[*altv1.GetBacktestRunRequest, *altv1.GetBacktestRunResponse](c, ctx, req)
}

// ListBacktestRuns lists backtest runs.
func (c *APIClient) ListBacktestRuns(ctx context.Context, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error) {
	return sendTyped[*altv1.ListBacktestRunsRequest, *altv1.ListBacktestRunsResponse](c, ctx, req)
}

// GetBacktestRunDetail gets detailed info (run + result) for a backtest run.
func (c *APIClient) GetBacktestRunDetail(ctx context.Context, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error) {
	return sendTyped[*altv1.GetBacktestRunDetailRequest, *altv1.GetBacktestRunDetailResponse](c, ctx, req)
}

// GetBacktestResult gets backtest summary, trades, etc.
func (c *APIClient) GetBacktestResult(ctx context.Context, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error) {
	return sendTyped[*altv1.GetBacktestResultRequest, *altv1.GetBacktestResultResponse](c, ctx, req)
}

// CompareBacktestRuns compares multiple backtest runs.
func (c *APIClient) CompareBacktestRuns(ctx context.Context, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error) {
	return sendTyped[*altv1.CompareBacktestRunsRequest, *altv1.CompareBacktestRunsResponse](c, ctx, req)
}

// ImportDailyBars requests the import of KIS daily bars into the worker storage.
func (c *APIClient) ImportDailyBars(ctx context.Context, req *altv1.ImportDailyBarsRequest) (*altv1.ImportDailyBarsResponse, error) {
	return sendTyped[*altv1.ImportDailyBarsRequest, *altv1.ImportDailyBarsResponse](c, ctx, req)
}

// AggregateMonthlyBars requests aggregation of stored daily bars into monthly bars.
func (c *APIClient) AggregateMonthlyBars(ctx context.Context, req *altv1.AggregateMonthlyBarsRequest) (*altv1.AggregateMonthlyBarsResponse, error) {
	return sendTyped[*altv1.AggregateMonthlyBarsRequest, *altv1.AggregateMonthlyBarsResponse](c, ctx, req)
}

// StartPaperTrading starts a paper trading account run.
func (c *APIClient) StartPaperTrading(ctx context.Context, req *altv1.StartPaperTradingRequest) (*altv1.StartPaperTradingResponse, error) {
	return sendTyped[*altv1.StartPaperTradingRequest, *altv1.StartPaperTradingResponse](c, ctx, req)
}

// GetPaperTradingState reads a paper trading account state.
func (c *APIClient) GetPaperTradingState(ctx context.Context, req *altv1.GetPaperTradingStateRequest) (*altv1.GetPaperTradingStateResponse, error) {
	return sendTyped[*altv1.GetPaperTradingStateRequest, *altv1.GetPaperTradingStateResponse](c, ctx, req)
}

// SubmitPaperOrder submits a virtual paper order.
func (c *APIClient) SubmitPaperOrder(ctx context.Context, req *altv1.SubmitPaperOrderRequest) (*altv1.SubmitPaperOrderResponse, error) {
	return sendTyped[*altv1.SubmitPaperOrderRequest, *altv1.SubmitPaperOrderResponse](c, ctx, req)
}

// CancelPaperOrder cancels a pending virtual paper order.
func (c *APIClient) CancelPaperOrder(ctx context.Context, req *altv1.CancelPaperOrderRequest) (*altv1.CancelPaperOrderResponse, error) {
	return sendTyped[*altv1.CancelPaperOrderRequest, *altv1.CancelPaperOrderResponse](c, ctx, req)
}

// FillPaperOrder simulates the fill of a pending virtual paper order.
func (c *APIClient) FillPaperOrder(ctx context.Context, req *altv1.FillPaperOrderRequest) (*altv1.FillPaperOrderResponse, error) {
	return sendTyped[*altv1.FillPaperOrderRequest, *altv1.FillPaperOrderResponse](c, ctx, req)
}

// SubmitLiveOrder submits a real live order through the broker.
func (c *APIClient) SubmitLiveOrder(ctx context.Context, req *altv1.SubmitLiveOrderRequest) (*altv1.SubmitLiveOrderResponse, error) {
	return sendTyped[*altv1.SubmitLiveOrderRequest, *altv1.SubmitLiveOrderResponse](c, ctx, req)
}

// CancelLiveOrder cancels a pending live order.
func (c *APIClient) CancelLiveOrder(ctx context.Context, req *altv1.CancelLiveOrderRequest) (*altv1.CancelLiveOrderResponse, error) {
	return sendTyped[*altv1.CancelLiveOrderRequest, *altv1.CancelLiveOrderResponse](c, ctx, req)
}

// GetLiveOrder retrieves the current status of a live order.
func (c *APIClient) GetLiveOrder(ctx context.Context, req *altv1.GetLiveOrderRequest) (*altv1.GetLiveOrderResponse, error) {
	return sendTyped[*altv1.GetLiveOrderRequest, *altv1.GetLiveOrderResponse](c, ctx, req)
}

// GetLiveRiskPolicy retrieves the current live trading risk policy.
func (c *APIClient) GetLiveRiskPolicy(ctx context.Context, req *altv1.GetLiveRiskPolicyRequest) (*altv1.GetLiveRiskPolicyResponse, error) {
	return sendTyped[*altv1.GetLiveRiskPolicyRequest, *altv1.GetLiveRiskPolicyResponse](c, ctx, req)
}

// GetLiveKillSwitch retrieves the current live kill switch state.
func (c *APIClient) GetLiveKillSwitch(ctx context.Context, req *altv1.GetLiveKillSwitchRequest) (*altv1.GetLiveKillSwitchResponse, error) {
	return sendTyped[*altv1.GetLiveKillSwitchRequest, *altv1.GetLiveKillSwitchResponse](c, ctx, req)
}

// SetLiveKillSwitch sets the live kill switch state.
func (c *APIClient) SetLiveKillSwitch(ctx context.Context, req *altv1.SetLiveKillSwitchRequest) (*altv1.SetLiveKillSwitchResponse, error) {
	return sendTyped[*altv1.SetLiveKillSwitchRequest, *altv1.SetLiveKillSwitchResponse](c, ctx, req)
}

// SyncLiveAccount fetches a fresh account snapshot from the broker via the worker.
func (c *APIClient) SyncLiveAccount(ctx context.Context, req *altv1.SyncLiveAccountRequest) (*altv1.SyncLiveAccountResponse, error) {
	return sendTyped[*altv1.SyncLiveAccountRequest, *altv1.SyncLiveAccountResponse](c, ctx, req)
}

// GetLiveAccountSnapshot retrieves the last synced account snapshot from the worker.
func (c *APIClient) GetLiveAccountSnapshot(ctx context.Context, req *altv1.GetLiveAccountSnapshotRequest) (*altv1.GetLiveAccountSnapshotResponse, error) {
	return sendTyped[*altv1.GetLiveAccountSnapshotRequest, *altv1.GetLiveAccountSnapshotResponse](c, ctx, req)
}

// ListLiveAuditEvents queries the durable live operation audit trail.
func (c *APIClient) ListLiveAuditEvents(ctx context.Context, req *altv1.ListLiveAuditEventsRequest) (*altv1.ListLiveAuditEventsResponse, error) {
	return sendTyped[*altv1.ListLiveAuditEventsRequest, *altv1.ListLiveAuditEventsResponse](c, ctx, req)
}

// SchedulerRefreshStatus queries the runtime refresh status of one or all
// scheduler schedules.
func (c *APIClient) SchedulerRefreshStatus(ctx context.Context, req *altv1.SchedulerRefreshStatusRequest) (*altv1.SchedulerRefreshStatusResponse, error) {
	return sendTyped[*altv1.SchedulerRefreshStatusRequest, *altv1.SchedulerRefreshStatusResponse](c, ctx, req)
}

// sendTyped is the shared request path for every API call. It derives the
// request timeout from the context deadline and maps a dropped connection to
// ErrTransport so each method stays a one-line forwarder and the transport vs.
// typed-error distinction cannot drift between them.
func sendTyped[Req proto.Message, Res proto.Message](c *APIClient, ctx context.Context, req Req) (Res, error) {
	var zero Res

	if err := ctx.Err(); err != nil {
		return zero, mapContextError(err)
	}
	if c == nil || c.ws == nil || !c.ws.IsAlive() {
		return zero, ErrTransport
	}

	timeout := 5 * time.Second
	if dl, ok := ctx.Deadline(); ok {
		timeout = time.Until(dl)
		if timeout <= 0 {
			return zero, fmt.Errorf("%w: deadline already exceeded", ErrTransport)
		}
	}

	type result struct {
		res Res
		err error
	}
	ch := make(chan result, 1)
	go func() {
		res, err := protoSocket.SendRequestTyped[Req, Res](&c.ws.Communicator, req, timeout)
		ch <- result{res: res, err: err}
	}()

	select {
	case <-ctx.Done():
		return zero, mapContextError(ctx.Err())
	case r := <-ch:
		if r.err != nil {
			return zero, fmt.Errorf("%w: %v", ErrTransport, r.err)
		}
		return r.res, nil
	}
}

// mapContextError normalises context errors: cancellation propagates as-is,
// deadlines surface as ErrTransport because the typed response never arrived.
func mapContextError(err error) error {
	if errors.Is(err, context.DeadlineExceeded) {
		return fmt.Errorf("%w: %v", ErrTransport, err)
	}
	return err
}