Provide a standalone agent client for l-ai-chat with service-token auth, safe at-least-once subscriptions, ordered cumulative ACK tracking, and self-contained gRPC fake tests. Co-authored-by: Cursor <cursoragent@cursor.com>
154 lines
4.0 KiB
Go
154 lines
4.0 KiB
Go
package chatclient
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"fmt"
|
|
"time"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
)
|
|
|
|
const (
|
|
defaultSubscriptionBufferSize = 100
|
|
defaultReconnectInitial = 200 * time.Millisecond
|
|
defaultReconnectMax = 5 * time.Second
|
|
defaultCloseTimeout = 5 * time.Second
|
|
)
|
|
|
|
type Sleeper func(ctx context.Context, d time.Duration) error
|
|
type Jitter func(d time.Duration) time.Duration
|
|
|
|
type Option func(*options)
|
|
|
|
type options struct {
|
|
tokenProvider TokenProvider
|
|
insecureTransport bool
|
|
dialOptions []grpc.DialOption
|
|
conn *grpc.ClientConn
|
|
subscriptionBufferSize int
|
|
reconnectInitial time.Duration
|
|
reconnectMax time.Duration
|
|
closeTimeout time.Duration
|
|
sleeper Sleeper
|
|
jitter Jitter
|
|
}
|
|
|
|
func defaultOptions() options {
|
|
return options{
|
|
subscriptionBufferSize: defaultSubscriptionBufferSize,
|
|
reconnectInitial: defaultReconnectInitial,
|
|
reconnectMax: defaultReconnectMax,
|
|
closeTimeout: defaultCloseTimeout,
|
|
sleeper: sleepContext,
|
|
jitter: defaultJitter,
|
|
}
|
|
}
|
|
|
|
func WithTokenProvider(provider TokenProvider) Option {
|
|
return func(o *options) { o.tokenProvider = provider }
|
|
}
|
|
|
|
func WithStaticToken(token string) Option {
|
|
return WithTokenProvider(StaticToken(token))
|
|
}
|
|
|
|
func WithInsecureTransport() Option {
|
|
return func(o *options) { o.insecureTransport = true }
|
|
}
|
|
|
|
func WithGRPCDialOptions(opts ...grpc.DialOption) Option {
|
|
return func(o *options) { o.dialOptions = append(o.dialOptions, opts...) }
|
|
}
|
|
|
|
func WithClientConn(conn *grpc.ClientConn) Option {
|
|
return func(o *options) { o.conn = conn }
|
|
}
|
|
|
|
func WithSubscriptionBufferSize(size int) Option {
|
|
return func(o *options) { o.subscriptionBufferSize = size }
|
|
}
|
|
|
|
func WithReconnectBackoff(initial, max time.Duration) Option {
|
|
return func(o *options) {
|
|
o.reconnectInitial = initial
|
|
o.reconnectMax = max
|
|
}
|
|
}
|
|
|
|
func withSleeperForTest(s Sleeper) Option {
|
|
return func(o *options) { o.sleeper = s }
|
|
}
|
|
|
|
func withJitterForTest(j Jitter) Option {
|
|
return func(o *options) { o.jitter = j }
|
|
}
|
|
|
|
func buildDialOptions(o options) ([]grpc.DialOption, error) {
|
|
if o.tokenProvider == nil {
|
|
return nil, fmt.Errorf("%w: token provider is required", ErrInvalidConfig)
|
|
}
|
|
|
|
var transport credentials.TransportCredentials
|
|
if o.insecureTransport {
|
|
transport = insecure.NewCredentials()
|
|
} else {
|
|
transport = credentials.NewTLS(&tls.Config{MinVersion: tls.VersionTLS12})
|
|
}
|
|
|
|
creds := tokenCredentials{
|
|
provider: o.tokenProvider,
|
|
requireTransportSecurity: !o.insecureTransport,
|
|
}
|
|
out := []grpc.DialOption{
|
|
grpc.WithTransportCredentials(transport),
|
|
grpc.WithPerRPCCredentials(creds),
|
|
}
|
|
out = append(out, o.dialOptions...)
|
|
return out, nil
|
|
}
|
|
|
|
func validateOptions(o options) error {
|
|
if o.tokenProvider == nil {
|
|
return fmt.Errorf("%w: token provider is required", ErrInvalidConfig)
|
|
}
|
|
if o.subscriptionBufferSize <= 0 {
|
|
return fmt.Errorf("%w: subscription buffer size must be positive", ErrInvalidConfig)
|
|
}
|
|
if o.reconnectInitial <= 0 {
|
|
return fmt.Errorf("%w: reconnect initial backoff must be positive", ErrInvalidConfig)
|
|
}
|
|
if o.reconnectMax < o.reconnectInitial {
|
|
return fmt.Errorf("%w: reconnect max backoff must be >= initial backoff", ErrInvalidConfig)
|
|
}
|
|
if o.closeTimeout <= 0 {
|
|
return fmt.Errorf("%w: close timeout must be positive", ErrInvalidConfig)
|
|
}
|
|
if o.sleeper == nil || o.jitter == nil {
|
|
return fmt.Errorf("%w: backoff hooks must be configured", ErrInvalidConfig)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func sleepContext(ctx context.Context, d time.Duration) error {
|
|
timer := time.NewTimer(d)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-timer.C:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
func defaultJitter(d time.Duration) time.Duration {
|
|
// Deterministic half-jitter keeps reconnects bounded without pulling in a
|
|
// retry framework. Tests replace this hook.
|
|
if d <= 1 {
|
|
return d
|
|
}
|
|
return d - d/4
|
|
}
|