l-ai-chat-client-go/client.go

package chatclient

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"github.com/google/uuid"
	"google.golang.org/grpc"

	pb "gitea.fe.kaien.ru/kaien/l-ai-chat-client-go/proto/l-ai-chat-proto"
)

type Client struct {
	conn  *grpc.ClientConn
	own   bool
	api   pb.ChatServiceClient
	opts  options
	mu    sync.Mutex
	subs  map[*Subscription]struct{}
	close bool
}

func New(ctx context.Context, target string, opts ...Option) (*Client, error) {
	o := defaultOptions()
	for _, opt := range opts {
		opt(&o)
	}
	if err := validateOptions(o); err != nil {
		return nil, err
	}

	var conn *grpc.ClientConn
	own := false
	if o.conn != nil {
		conn = o.conn
	} else {
		dialOpts, err := buildDialOptions(o)
		if err != nil {
			return nil, err
		}
		conn, err = grpc.NewClient(target, dialOpts...)
		if err != nil {
			return nil, err
		}
		own = true
	}

	c := &Client{
		conn: conn,
		own:  own,
		api:  pb.NewChatServiceClient(conn),
		opts: o,
		subs: make(map[*Subscription]struct{}),
	}
	if ctx != nil {
		select {
		case <-ctx.Done():
			_ = c.Close()
			return nil, ctx.Err()
		default:
		}
	}
	return c, nil
}

func (c *Client) Close() error {
	c.mu.Lock()
	if c.close {
		c.mu.Unlock()
		return nil
	}
	c.close = true
	subs := make([]*Subscription, 0, len(c.subs))
	for sub := range c.subs {
		subs = append(subs, sub)
	}
	c.mu.Unlock()

	for _, sub := range subs {
		_ = sub.Close()
	}
	if c.own && c.conn != nil {
		return c.conn.Close()
	}
	return nil
}

func (c *Client) ListTopics(ctx context.Context) ([]*pb.Topic, error) {
	if err := c.checkOpen(); err != nil {
		return nil, err
	}
	resp, err := c.api.ListTopics(ctx, &pb.ListTopicsRequest{})
	if err != nil {
		return nil, err
	}
	return resp.GetTopics(), nil
}

func (c *Client) GetTopic(ctx context.Context, req GetTopicRequest) (*pb.Topic, error) {
	if err := c.checkOpen(); err != nil {
		return nil, err
	}
	if req.ID != "" && req.Slug != "" {
		return nil, fmt.Errorf("%w: only one of topic id or slug may be set", ErrInvalidConfig)
	}
	in := &pb.GetTopicRequest{}
	switch {
	case req.ID != "":
		in.Selector = &pb.GetTopicRequest_Id{Id: req.ID}
	case req.Slug != "":
		in.Selector = &pb.GetTopicRequest_Slug{Slug: req.Slug}
	default:
		return nil, fmt.Errorf("%w: topic id or slug is required", ErrInvalidConfig)
	}
	return c.api.GetTopic(ctx, in)
}

type GetTopicRequest struct {
	ID   string
	Slug string
}

func (c *Client) GetMessage(ctx context.Context, messageID string) (*pb.Message, error) {
	if err := c.checkOpen(); err != nil {
		return nil, err
	}
	return c.api.GetMessage(ctx, &pb.GetMessageRequest{MessageId: messageID})
}

type PostMessageRequest struct {
	TopicID          string
	Content          string
	ReplyToMessageID string
	IdempotencyKey   string
}

func (c *Client) PostMessage(ctx context.Context, req PostMessageRequest) (*pb.PostMessageResponse, error) {
	if err := c.checkOpen(); err != nil {
		return nil, err
	}
	in := &pb.AgentPostMessageRequest{
		TopicId:        req.TopicID,
		Content:        req.Content,
		IdempotencyKey: req.IdempotencyKey,
	}
	if req.ReplyToMessageID != "" {
		replyTo := req.ReplyToMessageID
		in.ReplyToMessageId = &replyTo
	}
	return c.api.AgentPostMessage(ctx, in)
}

func NewIdempotencyKey() string {
	return uuid.NewString()
}

func (c *Client) Subscribe(ctx context.Context, req SubscribeRequest) (*Subscription, error) {
	if err := c.checkOpen(); err != nil {
		return nil, err
	}
	if req.TopicID == "" {
		return nil, fmt.Errorf("%w: topic id is required", ErrInvalidConfig)
	}
	sub, err := newSubscription(ctx, c, req)
	if err != nil {
		return nil, err
	}
	c.mu.Lock()
	if c.close {
		c.mu.Unlock()
		_ = sub.Close()
		return nil, ErrClosed
	}
	c.subs[sub] = struct{}{}
	c.mu.Unlock()
	return sub, nil
}

type SubscribeRequest struct {
	TopicID string
}

func (c *Client) unregister(sub *Subscription) {
	c.mu.Lock()
	delete(c.subs, sub)
	c.mu.Unlock()
}

func (c *Client) checkOpen() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.close {
		return ErrClosed
	}
	return nil
}

func ignoreClosed(err error) error {
	if errors.Is(err, context.Canceled) || errors.Is(err, ErrClosed) {
		return nil
	}
	return err
}