From 38fd2936bf950a97203a0a8d124eb4e2f5ac505e Mon Sep 17 00:00:00 2001 From: Danielle Mustillo Date: Thu, 5 Nov 2026 12:15:04 -0500 Subject: [PATCH 1/3] fix(grpcserver): report database connect failures as Unavailable When the controller could not log in to PostgreSQL at runtime, the RPC returned Internal "internal server error" and the cause never reached the logs. - A *pgconn.ConnectError in the error chain now maps to Unavailable "database unavailable", so clients know to retry. Internal errors log their cause. - A login PostgreSQL rejects (28P01, 28000, 3D000, 42501) is a config error, not an outage, so it stays Internal. - A database failure during runtime auth returns Unavailable instead of Unauthenticated "invalid credentials". Signed-off-by: Danielle Mustillo --- go/core/internal/database/errors.go | 20 ++ go/core/internal/database/errors_test.go | 86 ++++++++ go/core/internal/dbtest/dbtest.go | 34 +++- go/core/internal/grpcserver/interceptors.go | 30 ++- .../internal/grpcserver/interceptors_test.go | 188 ++++++++++++++++++ 5 files changed, 353 insertions(+), 5 deletions(-) create mode 100644 go/core/internal/database/errors_test.go diff --git a/go/core/internal/database/errors.go b/go/core/internal/database/errors.go index 3184e03884..5b4f828cec 100644 --- a/go/core/internal/database/errors.go +++ b/go/core/internal/database/errors.go @@ -3,6 +3,8 @@ package database import ( "errors" "fmt" + + "github.com/jackc/pgx/v5/pgconn" ) // Store errors describe resource-independent categories. Wrap them with resource @@ -21,3 +23,21 @@ var ( // ErrObjectDeleting reports that deletion has started. The object is unavailable // to new callers, so it also matches ErrNotFound. var ErrObjectDeleting = fmt.Errorf("object is being deleted: %w", ErrNotFound) + +// IsUnavailable reports whether err is a failed login that may succeed on retry. +// No SQL was sent, so a retry is safe. Rejected credentials are not unavailable. +func IsUnavailable(err error) bool { + var connectErr *pgconn.ConnectError + if !errors.As(err, &connectErr) { + return false + } + var pgErr *pgconn.PgError + if !errors.As(connectErr, &pgErr) { + return true + } + switch pgErr.Code { + case "28P01", "28000", "3D000", "42501": + return false + } + return true +} diff --git a/go/core/internal/database/errors_test.go b/go/core/internal/database/errors_test.go new file mode 100644 index 0000000000..8ba9e0f72b --- /dev/null +++ b/go/core/internal/database/errors_test.go @@ -0,0 +1,86 @@ +package database + +import ( + "errors" + "fmt" + "net" + "testing" + + "github.com/jackc/pgx/v5/pgconn" + + "github.com/kagent-dev/kagent/go/core/internal/dbtest" +) + +func TestIsUnavailable(t *testing.T) { + scenarios := []struct { + name string + err func(t *testing.T) error + retry bool + }{ + {"Given PostgreSQL refuses connections, when a login fails, assert that the database is unavailable", func(t *testing.T) error { + return loginError(t, refusingAddress(t)) + }, true}, + {"Given PostgreSQL is starting up, when a login is rejected with 57P03, assert that the database is unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "57P03")) + }, true}, + {"Given a proxy cannot reach PostgreSQL, when a login is rejected with 08001, assert that the database is unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "08001")) + }, true}, + {"Given a connect failure, when a caller wraps it, assert that the database is unavailable", func(t *testing.T) error { + return fmt.Errorf("list sessions: %w", loginError(t, refusingAddress(t))) + }, true}, + {"Given a wrong password, when a login is rejected with 28P01, assert that the database is not unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "28P01")) + }, false}, + {"Given an unknown login, when a login is rejected with 28000, assert that the database is not unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "28000")) + }, false}, + {"Given a missing database, when a login is rejected with 3D000, assert that the database is not unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "3D000")) + }, false}, + {"Given a login without CONNECT privilege, when a login is rejected with 42501, assert that the database is not unavailable", func(t *testing.T) error { + return loginError(t, dbtest.RejectLogins(t, "42501")) + }, false}, + {"Given an open connection, when a statement fails, assert that the database is not unavailable", func(*testing.T) error { + return &pgconn.PgError{Severity: "ERROR", Code: "23505"} + }, false}, + {"Given a non-database failure, when it is checked, assert that the database is not unavailable", func(*testing.T) error { + return errors.New("row decode failed") + }, false}, + } + for _, scenario := range scenarios { + t.Run(scenario.name, func(t *testing.T) { + err := scenario.err(t) + if got := IsUnavailable(err); got != scenario.retry { + t.Fatalf("IsUnavailable(%v) = %v, want %v", err, got, scenario.retry) + } + }) + } +} + +func refusingAddress(t *testing.T) string { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + address := listener.Addr().String() + if err := listener.Close(); err != nil { + t.Fatal(err) + } + return address +} + +func loginError(t *testing.T, address string) error { + t.Helper() + conn, err := pgconn.Connect(t.Context(), "postgres://kagent@"+address+"/kagent?sslmode=disable") + if err == nil { + _ = conn.Close(t.Context()) + t.Fatal("login succeeded") + } + var connectErr *pgconn.ConnectError + if !errors.As(err, &connectErr) { + t.Fatalf("got %T, want *pgconn.ConnectError", err) + } + return err +} diff --git a/go/core/internal/dbtest/dbtest.go b/go/core/internal/dbtest/dbtest.go index 492b9d0a2c..a7df6cca14 100644 --- a/go/core/internal/dbtest/dbtest.go +++ b/go/core/internal/dbtest/dbtest.go @@ -1,10 +1,12 @@ -// Package dbtest provides test helpers for spinning up a Postgres container. +// Package dbtest provides PostgreSQL test helpers: a real Postgres container and +// a fake server that rejects logins. package dbtest import ( "context" "database/sql" "fmt" + "net" "net/url" "strings" "testing" @@ -12,6 +14,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgproto3" "github.com/jackc/pgx/v5/pgxpool" testcontainers "github.com/testcontainers/testcontainers-go" tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" @@ -135,6 +138,35 @@ func PoolConfig(connStr string) *pgxpool.Config { return config } +// RejectLogins starts a server that answers every login with a FATAL error +// carrying sqlstate, such as 28P01. It stops when t ends. +func RejectLogins(t *testing.T, sqlstate string) string { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = listener.Close() }) + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func() { + defer conn.Close() + backend := pgproto3.NewBackend(conn, conn) + if _, err := backend.ReceiveStartupMessage(); err != nil { + return + } + backend.Send(&pgproto3.ErrorResponse{Severity: "FATAL", Code: sqlstate, Message: "login rejected"}) + _ = backend.Flush() + }() + } + }() + return listener.Addr().String() +} + // MigrateT runs the embedded migrations against connStr and calls t.Fatal on error. // If vectorEnabled is true the vector pass is also applied. func MigrateT(t *testing.T, connStr string, vectorEnabled bool) { diff --git a/go/core/internal/grpcserver/interceptors.go b/go/core/internal/grpcserver/interceptors.go index 924ed8da21..020f4ee1dd 100644 --- a/go/core/internal/grpcserver/interceptors.go +++ b/go/core/internal/grpcserver/interceptors.go @@ -20,6 +20,7 @@ import ( apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" runtimev1 "github.com/kagent-dev/kagent/go/api/gen/kagent/runtime/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" "github.com/kagent-dev/kagent/go/core/internal/service/agentidentity" "github.com/kagent-dev/kagent/go/core/internal/service/bootstrap" "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" @@ -76,7 +77,7 @@ func authenticate(ctx context.Context, fullMethod string, authenticator auth.Aut return ctx, status.Error(codes.Unauthenticated, "Agent identity authentication is not configured") } authenticatedContext, err := identity.Authenticate(ctx, headers) - return authenticatedContext, mapError(err) + return authenticatedContext, mapAndLogError(ctx, err) } if access == auth.AccessPublic && len(headers.Values("Authorization")) == 0 { return ctx, nil @@ -85,6 +86,9 @@ func authenticate(ctx context.Context, fullMethod string, authenticator auth.Aut return ctx, status.Error(codes.Unauthenticated, "authentication is not configured") } session, err := authenticator.Authenticate(ctx, headers, url.Values{}) + if database.IsUnavailable(err) { + return ctx, mapAndLogError(ctx, err) + } if err != nil || session == nil { return ctx, status.Error(codes.Unauthenticated, "invalid credentials") } @@ -94,7 +98,7 @@ func authenticate(ctx context.Context, fullMethod string, authenticator auth.Aut authenticatedContext := auth.AuthSessionTo(ctx, session) share, err := sessionsvc.ResolveShare(authenticatedContext, shareStore, headers.Get("X-Share-Token")) if err != nil { - return ctx, mapError(err) + return ctx, mapAndLogError(ctx, err) } if share == nil { return authenticatedContext, nil @@ -171,11 +175,24 @@ func requestLogger(ctx context.Context, method, rpcType string) *slog.Logger { func errorMappingUnaryInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { response, err := handler(ctx, req) - return response, mapError(err) + return response, mapAndLogError(ctx, err) } func errorMappingStreamInterceptor(srv any, stream grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error { - return mapError(handler(srv, stream)) + return mapAndLogError(stream.Context(), handler(srv, stream)) +} + +// mapAndLogError logs the cause of server-side failures, because mapError +// redacts it from the client response. +func mapAndLogError(ctx context.Context, err error) error { + mapped := mapError(err) + switch code := status.Code(mapped); code { + case codes.Internal: + logging.FromContext(ctx).ErrorContext(ctx, "rpc failed", "grpc_code", code.String(), "error", err) + case codes.Unavailable: + logging.FromContext(ctx).WarnContext(ctx, "rpc failed", "grpc_code", code.String(), "error", err) + } + return mapped } func mapError(err error) error { @@ -185,6 +202,11 @@ func mapError(err error) error { if errors.Is(err, context.Canceled) { return status.Error(codes.Canceled, "request canceled") } + // A connect timeout wraps DeadlineExceeded, but it reports an unreachable + // database rather than an expired caller deadline. + if database.IsUnavailable(err) { + return status.Error(codes.Unavailable, "database unavailable") + } if errors.Is(err, context.DeadlineExceeded) { return status.Error(codes.DeadlineExceeded, "request deadline exceeded") } diff --git a/go/core/internal/grpcserver/interceptors_test.go b/go/core/internal/grpcserver/interceptors_test.go index 029b087762..fdbd4890b8 100644 --- a/go/core/internal/grpcserver/interceptors_test.go +++ b/go/core/internal/grpcserver/interceptors_test.go @@ -1,14 +1,21 @@ package grpcserver import ( + "bytes" "context" "errors" + "fmt" + "log/slog" + "net" "net/http" "net/url" + "strings" "testing" + "time" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" @@ -17,8 +24,10 @@ import ( apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" runtimev1 "github.com/kagent-dev/kagent/go/api/gen/kagent/runtime/v1" "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/dbtest" "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" pkgauth "github.com/kagent-dev/kagent/go/core/pkg/auth" + "github.com/kagent-dev/kagent/go/pkg/logging" ) const ( @@ -225,6 +234,33 @@ func TestCallerWithoutRuntimeBindingCannotAccessRuntimeStorage(t *testing.T) { } } +func TestRuntimeAuthenticationReportsDatabaseConnectFailure(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + provider := &testAuthenticator{err: fmt.Errorf("resolve actor session: %w", connectErrorFromRefusedDial(t))} + + _, err := authenticate(ctx, runtimev1.TaskStoreService_GetTask_FullMethodName, provider, nil, DefaultMethodPolicies(), nil, nil) + + if got := status.Code(err); got != codes.Unavailable { + t.Fatalf("code = %v, want Unavailable: %v", got, err) + } + if got := status.Convert(err).Message(); got != "database unavailable" { + t.Fatalf("message = %q", got) + } + if !strings.Contains(logs.String(), "connection refused") { + t.Fatalf("cause missing from logs: %s", logs.String()) + } +} + +func TestRuntimeAuthenticationRejectsInvalidCredentials(t *testing.T) { + provider := &testAuthenticator{err: errors.New("token signature is invalid")} + + _, err := authenticate(t.Context(), runtimev1.TaskStoreService_GetTask_FullMethodName, provider, nil, DefaultMethodPolicies(), nil, nil) + + if got := status.Code(err); got != codes.Unauthenticated { + t.Fatalf("code = %v, want Unauthenticated: %v", got, err) + } +} + func TestMapError(t *testing.T) { tests := []struct { name string @@ -273,3 +309,155 @@ func TestRecoverUnaryInterceptor(t *testing.T) { t.Fatalf("message = %q", got) } } + +func TestErrorMappingReportsDatabaseConnectFailure(t *testing.T) { + refused := connectErrorFromRefusedDial(t) + timedOut := connectErrorFromUnansweredLogin(t) + scenarios := []struct { + name string + err error + }{ + {"Given PostgreSQL refuses connections, when a handler returns the connect error, assert that the RPC is Unavailable and the cause is logged", refused}, + {"Given PostgreSQL refuses connections, when a service wraps the connect error as internal, assert that the RPC is Unavailable and the cause is logged", serviceerrors.NewInternal("Failed to reserve Session", refused)}, + {"Given PostgreSQL never answers the login, when the connect deadline expires, assert that the RPC is Unavailable and the cause is logged", timedOut}, + } + for _, scenario := range scenarios { + t.Run(scenario.name, func(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + + _, err := errorMappingUnaryInterceptor(ctx, nil, &grpc.UnaryServerInfo{FullMethod: createMethod}, func(context.Context, any) (any, error) { + return nil, scenario.err + }) + + if got := status.Code(err); got != codes.Unavailable { + t.Fatalf("code = %v, want Unavailable", got) + } + if got := status.Convert(err).Message(); got != "database unavailable" { + t.Fatalf("message = %q", got) + } + if !strings.Contains(logs.String(), refusedOrTimedOutCause(scenario.err)) { + t.Fatalf("cause missing from logs: %s", logs.String()) + } + }) + } + + t.Run("Given a non-database failure, when a handler returns it, assert that the RPC is Internal with a generic message and the cause is logged", func(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + + _, err := errorMappingUnaryInterceptor(ctx, nil, &grpc.UnaryServerInfo{FullMethod: createMethod}, func(context.Context, any) (any, error) { + return nil, serviceerrors.NewInternal("Failed to reserve Session", errors.New("row decode failed")) + }) + + if got := status.Code(err); got != codes.Internal { + t.Fatalf("code = %v, want Internal", got) + } + if got := status.Convert(err).Message(); got != "internal server error" { + t.Fatalf("message = %q", got) + } + if !strings.Contains(logs.String(), "row decode failed") { + t.Fatalf("cause missing from logs: %s", logs.String()) + } + }) + + t.Run("Given PostgreSQL rejects the login password, when a handler returns the connect error, assert that the RPC is Internal and the cause is logged", func(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + rejected := connectError(t, t.Context(), dbtest.RejectLogins(t, "28P01")) + + _, err := errorMappingUnaryInterceptor(ctx, nil, &grpc.UnaryServerInfo{FullMethod: createMethod}, func(context.Context, any) (any, error) { + return nil, serviceerrors.NewInternal("Failed to reserve Session", rejected) + }) + + if got := status.Code(err); got != codes.Internal { + t.Fatalf("code = %v, want Internal", got) + } + if got := status.Convert(err).Message(); got != "internal server error" { + t.Fatalf("message = %q", got) + } + if !strings.Contains(logs.String(), "SQLSTATE 28P01") { + t.Fatalf("cause missing from logs: %s", logs.String()) + } + }) + + t.Run("Given a client error, when a handler returns it, assert that nothing is logged", func(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + + _, err := errorMappingUnaryInterceptor(ctx, nil, &grpc.UnaryServerInfo{FullMethod: createMethod}, func(context.Context, any) (any, error) { + return nil, serviceerrors.NewNotFound("missing", nil) + }) + + if got := status.Code(err); got != codes.NotFound { + t.Fatalf("code = %v, want NotFound", got) + } + if logs.Len() != 0 { + t.Fatalf("unexpected log: %s", logs.String()) + } + }) +} + +func contextWithCapturedLogs(t *testing.T) (context.Context, *bytes.Buffer) { + t.Helper() + var logs bytes.Buffer + return logging.IntoContext(t.Context(), slog.New(slog.NewTextHandler(&logs, nil))), &logs +} + +func refusedOrTimedOutCause(err error) string { + if errors.Is(err, context.DeadlineExceeded) { + return "context deadline exceeded" + } + return "connection refused" +} + +func connectErrorFromRefusedDial(t *testing.T) error { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + address := listener.Addr().String() + if err := listener.Close(); err != nil { + t.Fatal(err) + } + return connectError(t, t.Context(), address) +} + +func connectErrorFromUnansweredLogin(t *testing.T) error { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + t.Cleanup(func() { + close(done) + _ = listener.Close() + }) + go func() { + conn, err := listener.Accept() + if err != nil { + return + } + <-done + _ = conn.Close() + }() + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + err = connectError(t, ctx, listener.Addr().String()) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("connect error does not wrap DeadlineExceeded: %v", err) + } + return err +} + +func connectError(t *testing.T, ctx context.Context, address string) error { + t.Helper() + conn, err := pgconn.Connect(ctx, "postgres://kagent@"+address+"/kagent?sslmode=disable") + if err == nil { + _ = conn.Close(ctx) + t.Fatal("connect succeeded") + } + var connectErr *pgconn.ConnectError + if !errors.As(err, &connectErr) { + t.Fatalf("got %T, want *pgconn.ConnectError", err) + } + return err +} From 55fbb98e58d80c4c62ba32e863ab54a14e98e970 Mon Sep 17 00:00:00 2001 From: Danielle Mustillo Date: Thu, 5 Nov 2026 12:15:04 -0500 Subject: [PATCH 2/3] fix(database): retry failed PostgreSQL logins in store calls The store's pool makes one login attempt per acquire. When PostgreSQL is briefly unreachable, for example during a restart or failover, a store call fails at once even though the next login would succeed. - Client.db is a retryingPool with Exec, Query and BeginTx. It retries a failed login (database.IsUnavailable) with backoff, 0.5 s doubling to 5 s. A login fails before any SQL is sent, so no statement repeats. - database.pool.connectRetryTimeout sets how long new attempts may start, 90s by default, 0s to turn the retry off. The caller's context still ends the retry. - Each login's TCP dial stops after 10 s when the URL sets no connect_timeout, so one hung dial cannot outlast the retry window. - helm/README.md explains the retry window and how client timeouts, such as the web UI's 30 s, cut it short. Signed-off-by: Danielle Mustillo --- go/api/gen/kagent/config/v1/config.pb.go | 73 ++-- go/core/internal/database/client_postgres.go | 23 +- .../internal/database/connect_retry_test.go | 329 ++++++++++++++++++ .../database/push_notifications_test.go | 22 +- .../database/session_deletion_test.go | 2 +- go/core/internal/database/sql.go | 58 ++- go/core/pkg/app/app.go | 3 +- go/core/pkg/config/defaults.go | 2 +- go/core/pkg/config/helm_test.go | 4 + go/core/pkg/config/load_test.go | 27 ++ go/core/pkg/postgres/pool.go | 9 + helm/README.md | 6 + helm/kagent/values.yaml | 4 + proto/kagent/config/v1/config.proto | 6 + 14 files changed, 519 insertions(+), 49 deletions(-) create mode 100644 go/core/internal/database/connect_retry_test.go diff --git a/go/api/gen/kagent/config/v1/config.pb.go b/go/api/gen/kagent/config/v1/config.pb.go index 58b7630f7e..12c4aacddf 100644 --- a/go/api/gen/kagent/config/v1/config.pb.go +++ b/go/api/gen/kagent/config/v1/config.pb.go @@ -1109,8 +1109,11 @@ type PostgresPoolConfig struct { MaxConns *int32 `protobuf:"varint,2,opt,name=max_conns,json=maxConns,proto3,oneof" json:"max_conns,omitempty"` MaxConnIdleTime *durationpb.Duration `protobuf:"bytes,3,opt,name=max_conn_idle_time,json=maxConnIdleTime,proto3" json:"max_conn_idle_time,omitempty"` MaxConnLifetime *durationpb.Duration `protobuf:"bytes,4,opt,name=max_conn_lifetime,json=maxConnLifetime,proto3" json:"max_conn_lifetime,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // connect_retry_timeout is how long a database call retries a failed login, + // such as while a paused database resumes. 0s turns the retry off. Default 90s. + ConnectRetryTimeout *durationpb.Duration `protobuf:"bytes,5,opt,name=connect_retry_timeout,json=connectRetryTimeout,proto3" json:"connect_retry_timeout,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *PostgresPoolConfig) Reset() { @@ -1171,6 +1174,13 @@ func (x *PostgresPoolConfig) GetMaxConnLifetime() *durationpb.Duration { return nil } +func (x *PostgresPoolConfig) GetConnectRetryTimeout() *durationpb.Duration { + if x != nil { + return x.ConnectRetryTimeout + } + return nil +} + type SubstrateConfig struct { state protoimpl.MessageState `protogen:"open.v1"` ApiEndpoint *string `protobuf:"bytes,1,opt,name=api_endpoint,json=apiEndpoint,proto3,oneof" json:"api_endpoint,omitempty"` @@ -2223,14 +2233,16 @@ const file_kagent_config_v1_config_proto_rawDesc = "" + "\x0e_vector_schema\"D\n" + "\x13DatabaseLoginConfig\x12\x19\n" + "\burl_file\x18\x01 \x01(\tR\aurlFile\x12\x12\n" + - "\x04role\x18\x02 \x01(\tR\x04role\"\xa2\x03\n" + + "\x04role\x18\x02 \x01(\tR\x04role\"\x80\x04\n" + "\x12PostgresPoolConfig\x12$\n" + "\tmin_conns\x18\x01 \x01(\x05B\a\xbaH\x04\x1a\x02(\x00R\bminConns\x12)\n" + "\tmax_conns\x18\x02 \x01(\x05B\a\xbaH\x04\x1a\x02 \x00H\x00R\bmaxConns\x88\x01\x01\x12X\n" + "\x12max_conn_idle_time\x18\x03 \x01(\v2\x19.google.protobuf.DurationB\x10\xbaH\r\xaa\x01\n" + "\"\x06\b\x84\xfa\x85\xae\"*\x00R\x0fmaxConnIdleTime\x12W\n" + "\x11max_conn_lifetime\x18\x04 \x01(\v2\x19.google.protobuf.DurationB\x10\xbaH\r\xaa\x01\n" + - "\"\x06\b\x84\xfa\x85\xae\"*\x00R\x0fmaxConnLifetime:z\xbaHw\x1au\n" + + "\"\x06\b\x84\xfa\x85\xae\"*\x00R\x0fmaxConnLifetime\x12\\\n" + + "\x15connect_retry_timeout\x18\x05 \x01(\v2\x19.google.protobuf.DurationB\r\xbaH\n" + + "\xaa\x01\a\"\x03\b\xd8\x042\x00R\x13connectRetryTimeout:z\xbaHw\x1au\n" + "\x14postgres_pool.bounds\x12#min_conns must not exceed max_conns\x1a8!has(this.max_conns) || this.min_conns <= this.max_connsB\f\n" + "\n" + "_max_conns\"\xc6\x02\n" + @@ -2426,32 +2438,33 @@ var file_kagent_config_v1_config_proto_depIdxs = []int32{ 15, // 30: kagent.config.v1.DatabaseConfig.pool:type_name -> kagent.config.v1.PostgresPoolConfig 30, // 31: kagent.config.v1.PostgresPoolConfig.max_conn_idle_time:type_name -> google.protobuf.Duration 30, // 32: kagent.config.v1.PostgresPoolConfig.max_conn_lifetime:type_name -> google.protobuf.Duration - 30, // 33: kagent.config.v1.SessionsConfig.idle_ttl:type_name -> google.protobuf.Duration - 30, // 34: kagent.config.v1.SessionsConfig.share_max_ttl:type_name -> google.protobuf.Duration - 30, // 35: kagent.config.v1.SessionsConfig.expiration_interval:type_name -> google.protobuf.Duration - 30, // 36: kagent.config.v1.SessionsConfig.quiescence_recovery_interval:type_name -> google.protobuf.Duration - 30, // 37: kagent.config.v1.PushConfig.recovery_interval:type_name -> google.protobuf.Duration - 30, // 38: kagent.config.v1.PushConfig.send_timeout:type_name -> google.protobuf.Duration - 30, // 39: kagent.config.v1.SandboxConfig.default_ttl:type_name -> google.protobuf.Duration - 30, // 40: kagent.config.v1.SandboxConfig.max_ttl:type_name -> google.protobuf.Duration - 30, // 41: kagent.config.v1.SandboxConfig.expiration_interval:type_name -> google.protobuf.Duration - 30, // 42: kagent.config.v1.ScheduledRunsConfig.poll_interval:type_name -> google.protobuf.Duration - 30, // 43: kagent.config.v1.ScheduledRunsConfig.execution_interval:type_name -> google.protobuf.Duration - 30, // 44: kagent.config.v1.RuntimeRevisionsConfig.gc_interval:type_name -> google.protobuf.Duration - 25, // 45: kagent.config.v1.TelemetryConfig.traces:type_name -> kagent.config.v1.OTLPConfig - 25, // 46: kagent.config.v1.TelemetryConfig.metrics:type_name -> kagent.config.v1.OTLPConfig - 25, // 47: kagent.config.v1.TelemetryConfig.logs:type_name -> kagent.config.v1.OTLPConfig - 29, // 48: kagent.config.v1.TelemetryConfig.resource_attributes:type_name -> kagent.config.v1.TelemetryConfig.ResourceAttributesEntry - 30, // 49: kagent.config.v1.OTLPConfig.timeout:type_name -> google.protobuf.Duration - 4, // 50: kagent.config.v1.RuntimeDefaultsConfig.logging:type_name -> kagent.config.v1.LoggingConfig - 24, // 51: kagent.config.v1.RuntimeDefaultsConfig.telemetry:type_name -> kagent.config.v1.TelemetryConfig - 28, // 52: kagent.config.v1.RuntimeDefaultsConfig.token_propagation:type_name -> kagent.config.v1.TokenPropagationConfig - 27, // 53: kagent.config.v1.RuntimeDefaultsConfig.builtin_tools:type_name -> kagent.config.v1.BuiltinToolsConfig - 54, // [54:54] is the sub-list for method output_type - 54, // [54:54] is the sub-list for method input_type - 54, // [54:54] is the sub-list for extension type_name - 54, // [54:54] is the sub-list for extension extendee - 0, // [0:54] is the sub-list for field type_name + 30, // 33: kagent.config.v1.PostgresPoolConfig.connect_retry_timeout:type_name -> google.protobuf.Duration + 30, // 34: kagent.config.v1.SessionsConfig.idle_ttl:type_name -> google.protobuf.Duration + 30, // 35: kagent.config.v1.SessionsConfig.share_max_ttl:type_name -> google.protobuf.Duration + 30, // 36: kagent.config.v1.SessionsConfig.expiration_interval:type_name -> google.protobuf.Duration + 30, // 37: kagent.config.v1.SessionsConfig.quiescence_recovery_interval:type_name -> google.protobuf.Duration + 30, // 38: kagent.config.v1.PushConfig.recovery_interval:type_name -> google.protobuf.Duration + 30, // 39: kagent.config.v1.PushConfig.send_timeout:type_name -> google.protobuf.Duration + 30, // 40: kagent.config.v1.SandboxConfig.default_ttl:type_name -> google.protobuf.Duration + 30, // 41: kagent.config.v1.SandboxConfig.max_ttl:type_name -> google.protobuf.Duration + 30, // 42: kagent.config.v1.SandboxConfig.expiration_interval:type_name -> google.protobuf.Duration + 30, // 43: kagent.config.v1.ScheduledRunsConfig.poll_interval:type_name -> google.protobuf.Duration + 30, // 44: kagent.config.v1.ScheduledRunsConfig.execution_interval:type_name -> google.protobuf.Duration + 30, // 45: kagent.config.v1.RuntimeRevisionsConfig.gc_interval:type_name -> google.protobuf.Duration + 25, // 46: kagent.config.v1.TelemetryConfig.traces:type_name -> kagent.config.v1.OTLPConfig + 25, // 47: kagent.config.v1.TelemetryConfig.metrics:type_name -> kagent.config.v1.OTLPConfig + 25, // 48: kagent.config.v1.TelemetryConfig.logs:type_name -> kagent.config.v1.OTLPConfig + 29, // 49: kagent.config.v1.TelemetryConfig.resource_attributes:type_name -> kagent.config.v1.TelemetryConfig.ResourceAttributesEntry + 30, // 50: kagent.config.v1.OTLPConfig.timeout:type_name -> google.protobuf.Duration + 4, // 51: kagent.config.v1.RuntimeDefaultsConfig.logging:type_name -> kagent.config.v1.LoggingConfig + 24, // 52: kagent.config.v1.RuntimeDefaultsConfig.telemetry:type_name -> kagent.config.v1.TelemetryConfig + 28, // 53: kagent.config.v1.RuntimeDefaultsConfig.token_propagation:type_name -> kagent.config.v1.TokenPropagationConfig + 27, // 54: kagent.config.v1.RuntimeDefaultsConfig.builtin_tools:type_name -> kagent.config.v1.BuiltinToolsConfig + 55, // [55:55] is the sub-list for method output_type + 55, // [55:55] is the sub-list for method input_type + 55, // [55:55] is the sub-list for extension type_name + 55, // [55:55] is the sub-list for extension extendee + 0, // [0:55] is the sub-list for field type_name } func init() { file_kagent_config_v1_config_proto_init() } diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 5a4103db4e..e9df8eac61 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -12,21 +13,35 @@ import ( // Client persists control-plane state in PostgreSQL. Callers define the narrow // interfaces they need; SQL rows and protobuf encoding stay inside the store. type Client struct { - db *pgxpool.Pool + db retryingPool vectorCosineOperator string } +// ClientOption configures a Client built by NewClient. +type ClientOption func(*Client) + +// WithConnectRetry makes each store call retry a failed PostgreSQL login for up +// to timeout. A call can take up to timeout plus one attempt. Rejected +// credentials are not retried. Zero makes one attempt. +func WithConnectRetry(timeout time.Duration) ClientOption { + return func(c *Client) { c.db.retryFor = timeout } +} + // NewClient wraps an existing PostgreSQL pool without connecting or migrating. The caller // owns the pool and must close it. vectorSchema is the schema pgvector is installed in. -func NewClient(db *pgxpool.Pool, vectorSchema string) *Client { - return &Client{db: db, vectorCosineOperator: "OPERATOR(" + pgx.Identifier{vectorSchema}.Sanitize() + ".<=>)"} +func NewClient(db *pgxpool.Pool, vectorSchema string, opts ...ClientOption) *Client { + c := &Client{db: retryingPool{pool: db}, vectorCosineOperator: "OPERATOR(" + pgx.Identifier{vectorSchema}.Sanitize() + ".<=>)"} + for _, opt := range opts { + opt(c) + } + return c } // withTx commits all callback writes together on success and rolls them back on failure, // returning callback or transaction errors. The callback must keep external network work // outside the transaction. func (c *Client) withTx(ctx context.Context, fn func(pgx.Tx) error) error { - tx, err := c.db.Begin(ctx) + tx, err := c.db.BeginTx(ctx, pgx.TxOptions{}) if err != nil { return fmt.Errorf("failed to begin transaction: %w", err) } diff --git a/go/core/internal/database/connect_retry_test.go b/go/core/internal/database/connect_retry_test.go new file mode 100644 index 0000000000..0e8426dc39 --- /dev/null +++ b/go/core/internal/database/connect_retry_test.go @@ -0,0 +1,329 @@ +package database + +import ( + "context" + "encoding/binary" + "io" + "net" + "strconv" + "sync/atomic" + "syscall" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" + + "github.com/kagent-dev/kagent/go/core/internal/dbtest" +) + +func TestStoreRetriesFailedLogins(t *testing.T) { + t.Run("Given PostgreSQL drops the first two logins, when a caller stores a memory, assert that the memory is stored once", func(t *testing.T) { + inspect := NewClient(setupTestDB(t), "public") + proxy := newLoginProxy(t, dropFirst(2)) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Minute)) + + require.NoError(t, client.StoreAgentMemories(t.Context(), testMemory())) + + memories, err := inspect.ListAgentMemories(t.Context(), "team-a/agent", "owner") + require.NoError(t, err) + require.Len(t, memories, 1) + require.Equal(t, int32(3), proxy.logins.Load()) + }) + + t.Run("Given PostgreSQL drops the first two logins, when a caller stores memories in one transaction, assert that each memory is stored once", func(t *testing.T) { + inspect := NewClient(setupTestDB(t), "public") + proxy := newLoginProxy(t, dropFirst(2)) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Minute)) + + require.NoError(t, client.StoreAgentMemories(t.Context(), testMemory(), testMemory())) + + memories, err := inspect.ListAgentMemories(t.Context(), "team-a/agent", "owner") + require.NoError(t, err) + require.Len(t, memories, 2) + require.Equal(t, int32(3), proxy.logins.Load()) + }) + + t.Run("Given PostgreSQL drops the first login, when a caller deletes memories, assert that the memories are deleted", func(t *testing.T) { + inspect := NewClient(setupTestDB(t), "public") + require.NoError(t, inspect.StoreAgentMemories(t.Context(), testMemory())) + proxy := newLoginProxy(t, dropFirst(1)) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Minute)) + + require.NoError(t, client.DeleteAgentMemory(t.Context(), "team-a/agent", "owner")) + + memories, err := inspect.ListAgentMemories(t.Context(), "team-a/agent", "owner") + require.NoError(t, err) + require.Empty(t, memories) + require.Equal(t, int32(2), proxy.logins.Load()) + }) + + t.Run("Given PostgreSQL holds the first login past its connect timeout, when a caller reads, assert that the read succeeds on the next login", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, holdFirst(1)) + config := dbtest.PoolConfig(sharedConnStr) + config.ConnConfig.ConnectTimeout = 500 * time.Millisecond + client := proxy.client(t, config, WithConnectRetry(time.Minute)) + + _, err := client.ListAgentMemories(t.Context(), "team-a/agent", "owner") + + require.NoError(t, err) + require.Equal(t, int32(2), proxy.logins.Load()) + }) +} + +func TestStoreConnectRetryLimits(t *testing.T) { + t.Run("Given PostgreSQL drops every login, when the retry budget runs out, assert that the call fails with the connect error", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, dropAll) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Second)) + + start := time.Now() + _, err := client.ListAgentMemories(t.Context(), "team-a/agent", "owner") + elapsed := time.Since(start) + + require.True(t, IsUnavailable(err), "want a connect error, got %v", err) + require.GreaterOrEqual(t, elapsed, time.Second) + require.Less(t, elapsed, 2*time.Second) + require.Greater(t, proxy.logins.Load(), int32(1)) + }) + + t.Run("Given PostgreSQL drops every login, when the caller's deadline expires first, assert that the call returns at the deadline", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, dropAll) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Minute)) + ctx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := client.ListAgentMemories(ctx, "team-a/agent", "owner") + + require.Less(t, time.Since(start), time.Second) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.True(t, IsUnavailable(err), "want the connect error kept, got %v", err) + }) + + t.Run("Given PostgreSQL never answers the TCP handshake, when the retry budget runs out, assert that the call fails within the budget plus one dial", func(t *testing.T) { + config, err := poolConfig(&PostgresConfig{Pool: dbtest.PoolConfig("postgres://kagent@" + unansweredAddress(t) + "/kagent?sslmode=disable")}) + require.NoError(t, err) + pool, err := pgxpool.NewWithConfig(t.Context(), config) + require.NoError(t, err) + t.Cleanup(pool.Close) + client := NewClient(pool, "public", WithConnectRetry(15*time.Second)) + ctx, cancel := context.WithTimeout(t.Context(), 40*time.Second) + defer cancel() + + start := time.Now() + _, err = client.ListAgentMemories(ctx, "team-a/agent", "owner") + elapsed := time.Since(start) + + require.True(t, IsUnavailable(err), "want a connect error, got %v", err) + require.GreaterOrEqual(t, elapsed, 15*time.Second) + require.Less(t, elapsed, 25*time.Second, "a dial outlived its bound: %v", err) + }) + + t.Run("Given a wrong password, when a caller reads, assert that the login is tried once", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, forwardAll) + config := dbtest.PoolConfig(sharedConnStr) + config.ConnConfig.Password = "wrong" + client := proxy.client(t, config, WithConnectRetry(time.Minute)) + + start := time.Now() + _, err := client.ListAgentMemories(t.Context(), "team-a/agent", "owner") + + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Equal(t, "28P01", pgErr.Code) + require.Less(t, time.Since(start), defaultInitialDelay) + require.Equal(t, int32(1), proxy.logins.Load()) + }) + + scenarios := []struct { + name string + opts []ClientOption + }{ + {"Given no connect retry option, when PostgreSQL drops the first login, assert that the call fails after one login", nil}, + {"Given a zero connect retry timeout, when PostgreSQL drops the first login, assert that the call fails after one login", []ClientOption{WithConnectRetry(0)}}, + } + for _, scenario := range scenarios { + t.Run(scenario.name, func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, dropFirst(1)) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), scenario.opts...) + + _, err := client.ListAgentMemories(t.Context(), "team-a/agent", "owner") + + require.True(t, IsUnavailable(err), "want a connect error, got %v", err) + require.Equal(t, int32(1), proxy.logins.Load()) + }) + } + + t.Run("Given a connected database, when a statement fails, assert that the call fails without waiting", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, forwardAll) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(time.Minute)) + + start := time.Now() + err := client.DeleteSessionShare(t.Context(), "not-a-uuid", "owner") + + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + require.Less(t, time.Since(start), defaultInitialDelay) + }) + + t.Run("Given a retry budget shorter than a query, when the query runs past the budget, assert that it returns its rows", func(t *testing.T) { + setupTestDB(t) + proxy := newLoginProxy(t, forwardAll) + client := proxy.client(t, dbtest.PoolConfig(sharedConnStr), WithConnectRetry(100*time.Millisecond)) + + got, err := queryOne(t.Context(), client.db, `SELECT 1 FROM pg_sleep(0.3)`, pgx.RowTo[int32]) + + require.NoError(t, err) + require.Equal(t, int32(1), got) + }) +} + +func testMemory() *Memory { + return &Memory{AgentName: "team-a/agent", UserID: "owner", Embedding: makeEmbedding(0.5)} +} + +type loginAction int + +const ( + forwardLogin loginAction = iota + // dropLogin closes the connection after the startup message, like a proxy cutting a login. + dropLogin + // holdLogin never answers, like a resuming database. + holdLogin +) + +func dropFirst(n int32) func(int32) loginAction { + return func(login int32) loginAction { + if login <= n { + return dropLogin + } + return forwardLogin + } +} + +func holdFirst(n int32) func(int32) loginAction { + return func(login int32) loginAction { + if login <= n { + return holdLogin + } + return forwardLogin + } +} + +func dropAll(int32) loginAction { return dropLogin } + +func forwardAll(int32) loginAction { return forwardLogin } + +// unansweredAddress returns a loopback address with a full listen queue, so a +// dial hangs as it does against a host that is down. +func unansweredAddress(t *testing.T) string { + t.Helper() + fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, 0) + require.NoError(t, err) + t.Cleanup(func() { _ = syscall.Close(fd) }) + require.NoError(t, syscall.Bind(fd, &syscall.SockaddrInet4{Addr: [4]byte{127, 0, 0, 1}})) + require.NoError(t, syscall.Listen(fd, 1)) + bound, err := syscall.Getsockname(fd) + require.NoError(t, err) + address := net.JoinHostPort("127.0.0.1", strconv.Itoa(bound.(*syscall.SockaddrInet4).Port)) + for range 16 { + conn, err := net.DialTimeout("tcp", address, 200*time.Millisecond) + if err != nil { + return address + } + t.Cleanup(func() { _ = conn.Close() }) + } + t.Fatal("listen queue never filled") + return "" +} + +// loginProxy sits in front of the test database and handles each login by its +// 1-based position. Cancel requests are not counted. +type loginProxy struct { + address string + logins atomic.Int32 +} + +func newLoginProxy(t *testing.T, plan func(login int32) loginAction) *loginProxy { + t.Helper() + upstream, err := pgconn.ParseConfig(sharedConnStr) + require.NoError(t, err) + upstreamAddress := net.JoinHostPort(upstream.Host, strconv.Itoa(int(upstream.Port))) + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + done := make(chan struct{}) + t.Cleanup(func() { + close(done) + _ = listener.Close() + }) + proxy := &loginProxy{address: listener.Addr().String()} + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go proxy.serve(conn, upstreamAddress, plan, done) + } + }() + return proxy +} + +func (p *loginProxy) serve(conn net.Conn, upstreamAddress string, plan func(int32) loginAction, done <-chan struct{}) { + defer conn.Close() + header := make([]byte, 8) + if _, err := io.ReadFull(conn, header); err != nil { + return + } + const cancelRequestCode = 80877102 + if binary.BigEndian.Uint32(header[4:]) == cancelRequestCode { + return + } + switch plan(p.logins.Add(1)) { + case dropLogin: + _, _ = io.CopyN(io.Discard, conn, int64(binary.BigEndian.Uint32(header[:4]))-int64(len(header))) + return + case holdLogin: + <-done + return + } + upstream, err := net.Dial("tcp", upstreamAddress) + if err != nil { + return + } + defer upstream.Close() + if _, err := upstream.Write(header); err != nil { + return + } + copied := make(chan struct{}, 2) + go func() { _, _ = io.Copy(upstream, conn); copied <- struct{}{} }() + go func() { _, _ = io.Copy(conn, upstream); copied <- struct{}{} }() + select { + case <-copied: + case <-done: + } +} + +// client returns a Client that logs in through the proxy on first use. +func (p *loginProxy) client(t *testing.T, config *pgxpool.Config, opts ...ClientOption) *Client { + t.Helper() + config, err := poolConfig(&PostgresConfig{Pool: config, VectorEnabled: true, Schema: "public", VectorSchema: "public"}) + require.NoError(t, err) + host, port, err := net.SplitHostPort(p.address) + require.NoError(t, err) + portNumber, err := strconv.ParseUint(port, 10, 16) + require.NoError(t, err) + config.ConnConfig.Host, config.ConnConfig.Port = host, uint16(portNumber) + pool, err := pgxpool.NewWithConfig(t.Context(), config) + require.NoError(t, err) + t.Cleanup(pool.Close) + return NewClient(pool, "public", opts...) +} diff --git a/go/core/internal/database/push_notifications_test.go b/go/core/internal/database/push_notifications_test.go index f99cb060cf..944decd6bc 100644 --- a/go/core/internal/database/push_notifications_test.go +++ b/go/core/internal/database/push_notifications_test.go @@ -137,7 +137,7 @@ func TestPushRegistrationExpiresUnacceptedInput(t *testing.T) { require.NoError(t, execSQL(ctx, client.db, `UPDATE session_push_registration SET created_at = clock_timestamp() - interval '11 minutes' WHERE initial_message_id = $1`, "unaccepted")) require.NoError(t, client.ExpireUnboundPushRegistrations(ctx)) var closed bool - require.NoError(t, client.db.QueryRow(ctx, `SELECT closed_at IS NOT NULL FROM session_push_registration WHERE initial_message_id = $1`, "unaccepted").Scan(&closed)) + require.NoError(t, sharedDB.QueryRow(ctx, `SELECT closed_at IS NOT NULL FROM session_push_registration WHERE initial_message_id = $1`, "unaccepted").Scan(&closed)) require.True(t, closed) require.NoError(t, client.RegisterSessionPushNotification(ctx, session.Id, "unaccepted", "", config)) require.ErrorIs(t, client.RegisterSessionPushNotification(ctx, session.Id, "unaccepted", "", &a2a.PushConfig{ID: "default", URL: "http://changed"}), ErrIdempotencyConflict) @@ -152,7 +152,7 @@ func TestPushExpirationDoesNotCloseRegistrationBoundWhileWaiting(t *testing.T) { require.NoError(t, client.RegisterSessionPushNotification(ctx, session.Id, "pending", "", config)) require.NoError(t, execSQL(ctx, client.db, `UPDATE session_push_registration SET created_at = clock_timestamp() - interval '11 minutes' WHERE initial_message_id = $1`, "pending")) - tx, err := client.db.Begin(ctx) + tx, err := sharedDB.Begin(ctx) require.NoError(t, err) defer tx.Rollback(ctx) _, err = tx.Exec(ctx, `UPDATE session_push_registration SET task_id = $2 WHERE initial_message_id = $1`, "pending", task.ID) @@ -162,7 +162,7 @@ func TestPushExpirationDoesNotCloseRegistrationBoundWhileWaiting(t *testing.T) { go func() { expired <- client.ExpireUnboundPushRegistrations(ctx) }() require.Eventually(t, func() bool { var waiting bool - err := client.db.QueryRow(ctx, ` + err := sharedDB.QueryRow(ctx, ` SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname = current_database() AND query LIKE '%UPDATE session_push_registration SET closed_at%' @@ -173,7 +173,7 @@ func TestPushExpirationDoesNotCloseRegistrationBoundWhileWaiting(t *testing.T) { require.NoError(t, tx.Commit(ctx)) require.NoError(t, <-expired) var bound, closed bool - require.NoError(t, client.db.QueryRow(ctx, ` + require.NoError(t, sharedDB.QueryRow(ctx, ` SELECT task_id IS NOT NULL, closed_at IS NOT NULL FROM session_push_registration WHERE initial_message_id = $1 `, "pending").Scan(&bound, &closed)) @@ -430,7 +430,7 @@ func TestPushDeliveryBackoffAndAttemptLimit(t *testing.T) { var count int var delaySeconds float64 - require.NoError(t, client.db.QueryRow(ctx, ` + require.NoError(t, sharedDB.QueryRow(ctx, ` SELECT attempt_count, EXTRACT(EPOCH FROM (next_attempt_at - clock_timestamp()))::double precision FROM session_push_registration WHERE id = $1 @@ -465,7 +465,7 @@ func TestPushDeliveryBackoffAndAttemptLimit(t *testing.T) { require.NoError(t, err) require.True(t, current, "an exhausted generation cannot clear newer work") var count int - require.NoError(t, client.db.QueryRow(ctx, `SELECT attempt_count FROM session_push_registration WHERE id = $1`, next.ID).Scan(&count)) + require.NoError(t, sharedDB.QueryRow(ctx, `SELECT attempt_count FROM session_push_registration WHERE id = $1`, next.ID).Scan(&count)) require.Equal(t, 1, count) } @@ -858,7 +858,7 @@ func TestPushDeliveryClaimSkipsUncommittedPublication(t *testing.T) { task.Status.State = a2a.TaskStateCompleted version, err = client.UpdateSessionTask(ctx, session.Id, version, taskMutationHash("complete"), task, task, "") require.NoError(t, err) - tx, err := client.db.Begin(ctx) + tx, err := sharedDB.Begin(ctx) require.NoError(t, err) defer tx.Rollback(ctx) _, err = tx.Exec(ctx, `SELECT id FROM session_push_registration FOR UPDATE`) @@ -867,7 +867,7 @@ func TestPushDeliveryClaimSkipsUncommittedPublication(t *testing.T) { go func() { settled <- client.SettleSessionTask(ctx, session.Id, string(task.ID), version) }() require.Eventually(t, func() bool { var blocked bool - err := client.db.QueryRow(ctx, ` + err := sharedDB.QueryRow(ctx, ` SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname = current_database() AND query LIKE '%UPDATE session_push_registration SET pending_event_sequence%' @@ -966,7 +966,7 @@ func TestPushDeliveryInvalidTaskDoesNotBlockOtherRegistrations(t *testing.T) { stored, err := readSessionTask(ctx, client.db, session.HistoryID, string(badTask.ID)) require.NoError(t, err) var registrationID uuid.UUID - require.NoError(t, client.db.QueryRow(ctx, ` + require.NoError(t, sharedDB.QueryRow(ctx, ` SELECT id FROM session_push_registration WHERE history_id = $1 AND config_id = 'bad' `, session.HistoryID).Scan(®istrationID)) invalidData := []byte{0xff} @@ -986,7 +986,7 @@ func TestPushDeliveryInvalidTaskDoesNotBlockOtherRegistrations(t *testing.T) { require.Nil(t, delivery) var attempts int var scheduled, unleased bool - require.NoError(t, client.db.QueryRow(ctx, ` + require.NoError(t, sharedDB.QueryRow(ctx, ` SELECT attempt_count, next_attempt_at > clock_timestamp(), claim_token IS NULL AND lease_until IS NULL FROM session_push_registration WHERE config_id = 'bad' @@ -1039,7 +1039,7 @@ func TestPushDeliveryInvalidTaskDoesNotBlockOtherRegistrations(t *testing.T) { func assertPushDeliveryRetired(t *testing.T, client *Client, registrationID uuid.UUID) { t.Helper() var retired bool - require.NoError(t, client.db.QueryRow(t.Context(), ` + require.NoError(t, sharedDB.QueryRow(t.Context(), ` SELECT pending_event_sequence IS NULL AND attempt_count = 0 AND next_attempt_at IS NULL AND claim_token IS NULL AND lease_until IS NULL AND closed_at IS NULL FROM session_push_registration WHERE id = $1 diff --git a/go/core/internal/database/session_deletion_test.go b/go/core/internal/database/session_deletion_test.go index 0cb12b1588..bc04768a3e 100644 --- a/go/core/internal/database/session_deletion_test.go +++ b/go/core/internal/database/session_deletion_test.go @@ -52,7 +52,7 @@ func TestForceDeletionFencesTurnsAndSurvivesPreparationFailure(t *testing.T) { _, err = client.FinishSessionOperation(t.Context(), session.Id, operation.ID, uuid.Nil, "", "", "read failed") require.NoError(t, err) // A new controller/client can continue durable deletion without the force flag. - restarted := NewClient(client.db, "public") + restarted := NewClient(client.db.pool, "public") operation, err = restarted.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) require.NoError(t, err) finishAdmittedDeletion(t, restarted, operation) diff --git a/go/core/internal/database/sql.go b/go/core/internal/database/sql.go index b6bbd63445..35796c9a47 100644 --- a/go/core/internal/database/sql.go +++ b/go/core/internal/database/sql.go @@ -2,10 +2,15 @@ package database import ( "context" + "errors" + "fmt" + "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" + + "github.com/kagent-dev/kagent/go/pkg/logging" ) // dbExecutor lets the same store operation use a pool or its caller's transaction. @@ -15,10 +20,61 @@ type dbExecutor interface { } var ( - _ dbExecutor = (*pgxpool.Pool)(nil) + _ dbExecutor = retryingPool{} _ dbExecutor = (pgx.Tx)(nil) ) +// retryingPool retries failed logins on every store call. A login fails before +// any SQL is sent, so a retry never repeats a statement. +type retryingPool struct { + pool *pgxpool.Pool + // retryFor bounds when a new attempt may start. Zero makes one attempt. + retryFor time.Duration +} + +func (p retryingPool) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { + return retryLogin(ctx, p.retryFor, func() (pgconn.CommandTag, error) { return p.pool.Exec(ctx, sql, args...) }) +} + +func (p retryingPool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + return retryLogin(ctx, p.retryFor, func() (pgx.Rows, error) { return p.pool.Query(ctx, sql, args...) }) +} + +// BeginTx also lets pgx.BeginTxFunc take the pool. +func (p retryingPool) BeginTx(ctx context.Context, opts pgx.TxOptions) (pgx.Tx, error) { + return retryLogin(ctx, p.retryFor, func() (pgx.Tx, error) { return p.pool.BeginTx(ctx, opts) }) +} + +// retryLogin calls op until it succeeds, fails for another reason, ctx ends, or +// retryFor has passed. op runs on ctx, so a running statement is never cut. +func retryLogin[T any](ctx context.Context, retryFor time.Duration, op func() (T, error)) (T, error) { + if retryFor <= 0 { + return op() + } + start := time.Now() + delay := defaultInitialDelay + for attempt := 1; ; attempt++ { + result, err := op() + if err == nil || !IsUnavailable(err) { + return result, err + } + if ctx.Err() != nil { + return result, errors.Join(err, ctx.Err()) + } + remaining := retryFor - time.Since(start) + if remaining <= 0 { + return result, fmt.Errorf("database unavailable after %d login attempts in %s: %w", attempt, time.Since(start).Round(time.Millisecond), err) + } + logging.FromContext(ctx).WarnContext(ctx, "database login failed, retrying", "error", err, "attempt", attempt, "elapsed_ms", time.Since(start).Milliseconds()) + select { + case <-ctx.Done(): + return result, errors.Join(err, ctx.Err()) + case <-time.After(min(delay, remaining)): + } + delay = min(2*delay, defaultMaxDelay) + } +} + // queryOne runs parameterized SQL through the supplied pool or transaction and maps its // first row. It returns pgx.ErrNoRows for no results and propagates query/scan errors; it // closes rows and does not enforce uniqueness. diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 2b3458dbd5..dc998a9d41 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -242,7 +242,8 @@ func Run(ctx context.Context, input *configv1.ControllerConfig, opts Options) er return err } defer db.Close() - store := database.NewClient(db, cfg.GetDatabase().GetVectorSchema()) + store := database.NewClient(db, cfg.GetDatabase().GetVectorSchema(), + database.WithConnectRetry(cfg.GetDatabase().GetPool().GetConnectRetryTimeout().AsDuration())) kubeConfig, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( clientcmd.NewDefaultClientConfigLoadingRules(), &clientcmd.ConfigOverrides{}, diff --git a/go/core/pkg/config/defaults.go b/go/core/pkg/config/defaults.go index 2f5255990a..061144d0a1 100644 --- a/go/core/pkg/config/defaults.go +++ b/go/core/pkg/config/defaults.go @@ -61,7 +61,7 @@ func DefaultController(cfg *configv1.ControllerConfig) { } func defaultDatabase() *configv1.DatabaseConfig { - return &configv1.DatabaseConfig{Schema: new(consts.DefaultPostgresTableSchema), VectorSchema: new(consts.DefaultPgvectorSchema), Pool: &configv1.PostgresPoolConfig{MaxConns: new(int32(max(4, runtime.NumCPU()))), MaxConnIdleTime: durationpb.New(30 * time.Minute), MaxConnLifetime: durationpb.New(time.Hour)}} + return &configv1.DatabaseConfig{Schema: new(consts.DefaultPostgresTableSchema), VectorSchema: new(consts.DefaultPgvectorSchema), Pool: &configv1.PostgresPoolConfig{MaxConns: new(int32(max(4, runtime.NumCPU()))), MaxConnIdleTime: durationpb.New(30 * time.Minute), MaxConnLifetime: durationpb.New(time.Hour), ConnectRetryTimeout: durationpb.New(90 * time.Second)}} } func defaultTelemetry(cfg *configv1.TelemetryConfig) { diff --git a/go/core/pkg/config/helm_test.go b/go/core/pkg/config/helm_test.go index 8c49153290..9b4f975ab1 100644 --- a/go/core/pkg/config/helm_test.go +++ b/go/core/pkg/config/helm_test.go @@ -81,6 +81,10 @@ func TestHelmControllerConfig(t *testing.T) { require.Empty(t, base.Push.SigningPrivateKeyFile) require.EqualValues(t, 10, base.Push.GetMaxConcurrency()) require.Equal(t, 5*time.Second, base.Push.SendTimeout.AsDuration()) + require.Equal(t, 90*time.Second, base.GetDatabase().GetPool().GetConnectRetryTimeout().AsDuration()) + noRetry, _, err := render("config:\n database:\n pool: {connectRetryTimeout: 0s}\n") + require.NoError(t, err) + require.Zero(t, noRetry.GetDatabase().GetPool().GetConnectRetryTimeout().AsDuration()) verify, verifyChecksum, err := render("config:\n database: {skipMigrations: true}\n") require.NoError(t, err) require.True(t, verify.GetDatabase().GetSkipMigrations()) diff --git a/go/core/pkg/config/load_test.go b/go/core/pkg/config/load_test.go index 89d2dc2566..35e0e5d9f8 100644 --- a/go/core/pkg/config/load_test.go +++ b/go/core/pkg/config/load_test.go @@ -37,6 +37,10 @@ func TestControllerConfig(t *testing.T) { {"documents", "{}\n---\n{}", false}, {"field aliases", "kubernetes: {leaderElection: true, leader_election: false}", true}, {"pool bounds", "database: {pool: {minConns: 9, maxConns: 4}}", true}, + {"connect retry timeout", "database: {pool: {connectRetryTimeout: 90s}}", false}, + {"disabled connect retry", "database: {pool: {connectRetryTimeout: 0s}}", false}, + {"negative connect retry timeout", "database: {pool: {connectRetryTimeout: -1s}}", true}, + {"overlong connect retry timeout", "database: {pool: {connectRetryTimeout: 601s}}", true}, {"database schemas", "database: {schema: tenant, vectorSchema: public}", false}, {"invalid skip migrations type", "database: {skipMigrations: 'true'}", true}, {"empty database schema", "database: {schema: ''}", true}, @@ -100,6 +104,29 @@ func TestDatabaseSkipMigrations(t *testing.T) { } } +func TestDatabaseConnectRetryTimeout(t *testing.T) { + scenarios := []struct { + name, input string + want time.Duration + }{ + {"Given no connect retry timeout, when the controller config is defaulted, assert that logins are retried for 90s", "{}", 90 * time.Second}, + {"Given connectRetryTimeout 0s, when the controller config is defaulted, assert that the retry stays disabled", "database: {pool: {connectRetryTimeout: 0s}}", 0}, + {"Given connectRetryTimeout 90s, when the controller config is defaulted, assert that logins are retried for 90s", "database: {pool: {connectRetryTimeout: 90s}}", 90 * time.Second}, + } + for _, scenario := range scenarios { + t.Run(scenario.name, func(t *testing.T) { + cfg := &configv1.ControllerConfig{} + require.NoError(t, config.Decode([]byte(scenario.input), cfg)) + + config.DefaultController(cfg) + + setLogins(cfg) + require.NoError(t, config.Validate(cfg)) + require.Equal(t, scenario.want, cfg.GetDatabase().GetPool().GetConnectRetryTimeout().AsDuration()) + }) + } +} + func TestDefaultingPreservesPresence(t *testing.T) { cfg := &configv1.ControllerConfig{} require.NoError(t, config.Decode([]byte("kubernetes: {leaderElection: false}\nsessions: {idleTtl: 0s}\ndatabase: {pool: {minConns: 0}}"), cfg)) diff --git a/go/core/pkg/postgres/pool.go b/go/core/pkg/postgres/pool.go index 0c32a9c5dd..fa56dbaa45 100644 --- a/go/core/pkg/postgres/pool.go +++ b/go/core/pkg/postgres/pool.go @@ -5,7 +5,9 @@ import ( "context" "errors" "fmt" + "net" "slices" + "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" @@ -14,6 +16,10 @@ import ( var errInvalidURL = errors.New("invalid PostgreSQL connection string") +// dialTimeout bounds the TCP dial when the URL sets no connect_timeout. It does +// not bound the login, which a resuming database can hold for tens of seconds. +const dialTimeout = 10 * time.Second + // PoolConfig parses url into a pool configuration. Every connection assumes // role after authentication, so the login only needs membership in role. // Errors never include url, which can carry a password. @@ -25,6 +31,9 @@ func PoolConfig(url, role string) (*pgxpool.Config, error) { if err != nil { return nil, err } + if config.ConnConfig.ConnectTimeout == 0 { + config.ConnConfig.DialFunc = (&net.Dialer{Timeout: dialTimeout}).DialContext + } if usesTLS(config.ConnConfig) { // pgx reads TLS files once at parse time. Reparse before each new physical // connection so rotated client certificates take effect without a restart. diff --git a/helm/README.md b/helm/README.md index 51c82c55bd..e3148cff12 100644 --- a/helm/README.md +++ b/helm/README.md @@ -285,6 +285,12 @@ Kagent reads the connection strings from their Secrets once, at startup. After a When the connection string names TLS certificate files, Kagent rereads them before each new physical connection. Rotated client certificates, such as Pod Certificates, take effect without a restart. `config.database.pool.maxConnLifetime` limits how long a connection opened with an older certificate stays in use. +#### A database that pauses or restarts + +When the controller can't log in to PostgreSQL, for example while a serverless database resumes from a pause, a database call retries the login with backoff for `config.database.pool.connectRetryTimeout` (default `90s`, `0s` turns it off). No SQL has run before a login, so a retry repeats no work. The call returns `Unavailable` once the retry window ends. + +The retry also stops when the client stops waiting. The web UI waits 30 seconds for a request, so a resume longer than that fails the first request in the UI even though the controller would have waited longer. Set the retry window above your database's longest resume, and expect clients with shorter timeouts to retry. + #### OIDC authentication Set `config.server.auth.mode: trusted-proxy` together with diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 69c5b77417..838342a0b1 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -176,6 +176,10 @@ config: # Skip applying migrations at controller startup; verify they are already applied. # Apply migrations out-of-band before installing or upgrading when true. skipMigrations: false + # pool: + # # How long a database call retries a failed login, e.g. while a paused + # # database resumes. 0s turns the retry off. Default 90s. + # connectRetryTimeout: 90s push: allowHttp: false allowPrivateNetworks: false diff --git a/proto/kagent/config/v1/config.proto b/proto/kagent/config/v1/config.proto index fd97827559..c5370e695c 100644 --- a/proto/kagent/config/v1/config.proto +++ b/proto/kagent/config/v1/config.proto @@ -213,6 +213,12 @@ message PostgresPoolConfig { gt: {} lte: {seconds: 9223372036} }]; + // connect_retry_timeout is how long a database call retries a failed login, + // such as while a paused database resumes. 0s turns the retry off. Default 90s. + google.protobuf.Duration connect_retry_timeout = 5 [(buf.validate.field).duration = { + gte: {} + lte: {seconds: 600} + }]; } message SubstrateConfig { From f168e3e59735784fd07024064b5df67654c51fc8 Mon Sep 17 00:00:00 2001 From: Danielle Mustillo Date: Fri, 6 Nov 2026 14:23:34 -0500 Subject: [PATCH 3/3] fix(grpcserver): log rejected database logins during authentication A rejected controller login to PostgreSQL during actor authentication returned Unauthenticated "invalid credentials" and logged nothing. Every database login error now goes through the error mapper, so a rejected login returns Internal and logs its cause. Signed-off-by: Danielle Mustillo --- go/core/internal/grpcserver/interceptors.go | 4 +++- go/core/internal/grpcserver/interceptors_test.go | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/go/core/internal/grpcserver/interceptors.go b/go/core/internal/grpcserver/interceptors.go index 020f4ee1dd..5fb5edbcf1 100644 --- a/go/core/internal/grpcserver/interceptors.go +++ b/go/core/internal/grpcserver/interceptors.go @@ -12,6 +12,7 @@ import ( "time" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/jackc/pgx/v5/pgconn" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" @@ -86,7 +87,8 @@ func authenticate(ctx context.Context, fullMethod string, authenticator auth.Aut return ctx, status.Error(codes.Unauthenticated, "authentication is not configured") } session, err := authenticator.Authenticate(ctx, headers, url.Values{}) - if database.IsUnavailable(err) { + var connectErr *pgconn.ConnectError + if errors.As(err, &connectErr) { return ctx, mapAndLogError(ctx, err) } if err != nil || session == nil { diff --git a/go/core/internal/grpcserver/interceptors_test.go b/go/core/internal/grpcserver/interceptors_test.go index fdbd4890b8..d4714dc742 100644 --- a/go/core/internal/grpcserver/interceptors_test.go +++ b/go/core/internal/grpcserver/interceptors_test.go @@ -251,6 +251,21 @@ func TestRuntimeAuthenticationReportsDatabaseConnectFailure(t *testing.T) { } } +func TestRuntimeAuthenticationReportsRejectedDatabaseLogin(t *testing.T) { + ctx, logs := contextWithCapturedLogs(t) + rejected := connectError(t, t.Context(), dbtest.RejectLogins(t, "28P01")) + provider := &testAuthenticator{err: fmt.Errorf("resolve actor session: %w", rejected)} + + _, err := authenticate(ctx, runtimev1.TaskStoreService_GetTask_FullMethodName, provider, nil, DefaultMethodPolicies(), nil, nil) + + if got := status.Code(err); got != codes.Internal { + t.Fatalf("code = %v, want Internal: %v", got, err) + } + if !strings.Contains(logs.String(), "SQLSTATE 28P01") { + t.Fatalf("cause missing from logs: %s", logs.String()) + } +} + func TestRuntimeAuthenticationRejectsInvalidCredentials(t *testing.T) { provider := &testAuthenticator{err: errors.New("token signature is invalid")}