robinhood-agentic-mcp/equity/write.go

package equity

import (
	"context"
	"encoding/json"

	decimal "github.com/alpacahq/alpacadecimal"
	"s1d3sw1ped/robinhood-agentic-mcp/client"
	"s1d3sw1ped/robinhood-agentic-mcp/internal/wire"
)

const (
	toolReview = "review_equity_order"
	toolPlace  = "place_equity_order"
	toolCancel = "cancel_equity_order"
)

// PlaceOrderRequest is the argument set for review_equity_order and place_equity_order.
type PlaceOrderRequest struct {
	AccountNumber string
	Symbol        string
	Side          client.Side
	Type          client.OrderType
	Qty           *decimal.Decimal
	DollarAmount  *decimal.Decimal
	LimitPrice    *decimal.Decimal
	StopPrice     *decimal.Decimal
	TimeInForce   client.TimeInForce // empty → omit (Robinhood defaults gfd)
	MarketHours   client.MarketHours
	TaxLots       []TaxLot
	RefID         string
}

// TaxLot is a specified-lot selection for a sell order.
type TaxLot struct {
	OpenLotID string
	Quantity  decimal.Decimal
}

// ReviewResult is the pre-trade check from review_equity_order.
type ReviewResult struct {
	Errors   []string `json:"errors"`
	Warnings []string `json:"warnings"`
}

// Order is a placed or listed equity order.
type Order struct {
	ID             string          `json:"id"`
	IdempotencyKey string          `json:"idempotency_key"`
	Symbol         string          `json:"symbol"`
	Side           string          `json:"side"`
	Qty            decimal.Decimal `json:"-"`
	FilledQty      decimal.Decimal `json:"-"`
	Price          decimal.Decimal `json:"-"`
	State          string          `json:"state"`
}

func (o *Order) UnmarshalJSON(b []byte) error {
	var row struct {
		ID             string `json:"id"`
		IdempotencyKey string `json:"idempotency_key"`
		Symbol         string `json:"symbol"`
		Side           string `json:"side"`
		Quantity       any    `json:"quantity"`
		Qty            any    `json:"qty"`
		FilledQty      any    `json:"filled_quantity"`
		CumulativeQty  any    `json:"cumulative_quantity"`
		Price          any    `json:"price"`
		AveragePrice   any    `json:"average_price"`
		State          string `json:"state"`
		Status         string `json:"status"`
	}
	if err := json.Unmarshal(b, &row); err != nil {
		return err
	}
	o.ID = row.ID
	o.IdempotencyKey = row.IdempotencyKey
	o.Symbol = row.Symbol
	o.Side = row.Side
	o.State = row.State
	if o.State == "" {
		o.State = row.Status
	}
	var err error
	if o.Qty, err = firstDec(row.Quantity, row.Qty); err != nil {
		return err
	}
	if o.FilledQty, err = firstDec(row.FilledQty, row.CumulativeQty); err != nil {
		return err
	}
	if o.Price, err = firstDec(row.Price, row.AveragePrice); err != nil {
		return err
	}
	return nil
}

func parseOrders(raw json.RawMessage) ([]Order, error) {
	rows, _, err := wire.UnmarshalRows[Order](raw, "orders", "results")
	if err != nil {
		return nil, err
	}
	if len(rows) > 0 {
		return rows, nil
	}
	var one Order
	if err := json.Unmarshal(wire.Unwrap(raw), &one); err != nil {
		return nil, err
	}
	if one.ID != "" {
		return []Order{one}, nil
	}
	return nil, nil
}

// CancelOrderRequest is the argument set for cancel_equity_order.
type CancelOrderRequest struct {
	AccountNumber string
	OrderID       string
}

// ReviewOrder calls review_equity_order. RefID is not sent.
func (c *Client) ReviewOrder(ctx context.Context, req PlaceOrderRequest) (ReviewResult, error) {
	var out ReviewResult
	if err := c.parse(ctx, toolReview, placeArgs(req, false), &out); err != nil {
		return ReviewResult{}, err
	}
	return out, nil
}

// PlaceOrder calls place_equity_order. RefID is sent as ref_id only.
func (c *Client) PlaceOrder(ctx context.Context, req PlaceOrderRequest) (Order, error) {
	raw, err := c.c.Call(ctx, toolPlace, placeArgs(req, true))
	if err != nil {
		return Order{}, err
	}
	ords, err := parseOrders(raw)
	if err != nil {
		return Order{}, client.ToolErrorf(toolPlace, "parse: %w", err)
	}
	if len(ords) == 1 {
		return ords[0], nil
	}
	if len(ords) > 1 {
		return ords[0], nil
	}
	return Order{}, nil
}

// CancelOrder calls cancel_equity_order.
func (c *Client) CancelOrder(ctx context.Context, req CancelOrderRequest) error {
	args := map[string]any{}
	if req.AccountNumber != "" {
		args["account_number"] = req.AccountNumber
	}
	if req.OrderID != "" {
		args["order_id"] = req.OrderID
	}
	_, err := c.c.Call(ctx, toolCancel, args)
	return err
}

func placeArgs(req PlaceOrderRequest, withRef bool) map[string]any {
	args := map[string]any{}
	if req.AccountNumber != "" {
		args["account_number"] = req.AccountNumber
	}
	if req.Symbol != "" {
		args["symbol"] = req.Symbol
	}
	if req.Side != "" {
		args["side"] = string(req.Side)
	}
	if req.Type != "" {
		args["type"] = string(req.Type)
	}
	if req.Qty != nil {
		args["quantity"] = wire.Encode(*req.Qty)
	}
	if req.DollarAmount != nil {
		args["dollar_amount"] = wire.Encode(*req.DollarAmount)
	}
	if req.LimitPrice != nil {
		args["limit_price"] = wire.Encode(*req.LimitPrice)
	}
	if req.StopPrice != nil {
		args["stop_price"] = wire.Encode(*req.StopPrice)
	}
	if req.TimeInForce != "" {
		args["time_in_force"] = string(req.TimeInForce)
	}
	if req.MarketHours != "" {
		args["market_hours"] = string(req.MarketHours)
	}
	if len(req.TaxLots) > 0 {
		lots := make([]map[string]any, len(req.TaxLots))
		for i, lot := range req.TaxLots {
			lots[i] = map[string]any{
				"open_lot_id": lot.OpenLotID,
				"quantity":    wire.Encode(lot.Quantity),
			}
		}
		args["tax_lots"] = lots
	}
	if withRef && req.RefID != "" {
		args["ref_id"] = req.RefID
	}
	return args
}