fix(agent): install connstats callback at statsReporter creation (#25819)

The stats reporter only installed the connstats callback on the TUN
device after the report loop negotiated an interval with the server.
Traffic that flowed before that point (e.g. an SSH handshake) was
silently dropped because the TUN wrapper's stats.Load() returned nil.

We now install the connstats callback immediately and we no longer
re-install the callback every interval unless the interval changed.

Fixes flaky TestAgent_Stats_SSH, TestAgent_Stats_ReconnectingPTY,
and TestAgent_Stats_Magic by ensuring the connstats callback is
always installed before network traffic can flow.

Closes coder/internal#505
Closes CODAGT-517
This commit is contained in:
Mathias Fredriksson
2026-06-04 21:16:26 +03:00
committed by GitHub
parent 6bd413163f
commit 20d678b886
5 changed files with 167 additions and 42 deletions
+37 -2
View File
@@ -32,7 +32,8 @@ import (
"github.com/coder/websocket"
)
const statsInterval = 500 * time.Millisecond
// StatsInterval is the report interval returned by FakeAgentAPI.UpdateStats.
const StatsInterval = 500 * time.Millisecond
func NewClient(t testing.TB,
logger slog.Logger,
@@ -128,6 +129,17 @@ func (c *Client) RefreshToken(context.Context) error {
return nil
}
// SetUpdateStatsOverride sets a function that wraps UpdateStats calls.
// The provided function receives a next callback for the default behavior.
func (c *Client) SetUpdateStatsOverride(fn func(
ctx context.Context,
req *agentproto.UpdateStatsRequest,
next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error),
) (*agentproto.UpdateStatsResponse, error),
) {
c.fakeAgentAPI.SetUpdateStatsOverride(fn)
}
func (c *Client) GetNumRefreshTokenCalls() int {
c.mu.Lock()
defer c.mu.Unlock()
@@ -246,6 +258,11 @@ type FakeAgentAPI struct {
subAgentDisplayApps map[uuid.UUID][]agentproto.CreateSubAgentRequest_DisplayApp
subAgentApps map[uuid.UUID][]*agentproto.CreateSubAgentRequest_App
updateStatsOverride func(
ctx context.Context,
req *agentproto.UpdateStatsRequest,
next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error),
) (*agentproto.UpdateStatsResponse, error)
getAnnouncementBannersFunc func() ([]codersdk.BannerConfig, error)
getResourcesMonitoringConfigurationFunc func() (*agentproto.GetResourcesMonitoringConfigurationResponse, error)
pushResourcesMonitoringUsageFunc func(*agentproto.PushResourcesMonitoringUsageRequest) (*agentproto.PushResourcesMonitoringUsageResponse, error)
@@ -320,8 +337,26 @@ func (f *FakeAgentAPI) PushResourcesMonitoringUsage(_ context.Context, req *agen
return f.pushResourcesMonitoringUsageFunc(req)
}
func (f *FakeAgentAPI) SetUpdateStatsOverride(fn func(
ctx context.Context,
req *agentproto.UpdateStatsRequest,
next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error),
) (*agentproto.UpdateStatsResponse, error),
) {
f.Lock()
defer f.Unlock()
f.updateStatsOverride = fn
}
func (f *FakeAgentAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) {
f.logger.Debug(ctx, "update stats called", slog.F("req", req))
if f.updateStatsOverride != nil {
return f.updateStatsOverride(ctx, req, f.updateStatsDefault)
}
return f.updateStatsDefault(ctx, req)
}
func (f *FakeAgentAPI) updateStatsDefault(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) {
// empty request is sent to get the interval; but our tests don't want empty stats requests
if req.Stats != nil {
select {
@@ -331,7 +366,7 @@ func (f *FakeAgentAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateSt
// OK!
}
}
return &agentproto.UpdateStatsResponse{ReportInterval: durationpb.New(statsInterval)}, nil
return &agentproto.UpdateStatsResponse{ReportInterval: durationpb.New(StatsInterval)}, nil
}
func (f *FakeAgentAPI) GetLifecycleStates() []codersdk.WorkspaceAgentLifecycle {