From 5c9fc61619af3afac2cdfcb95ca09d47858e5fc1 Mon Sep 17 00:00:00 2001 From: Meet Vasani Date: Mon, 21 Jul 2025 17:54:04 +0530 Subject: [PATCH] fix(agent): handle context cancellation (#5323) Signed-off-by: ivaltryek Co-authored-by: qwerty287 <80460567+qwerty287@users.noreply.github.com> --- cmd/agent/core/agent.go | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/cmd/agent/core/agent.go b/cmd/agent/core/agent.go index 93f057024..af15cac15 100644 --- a/cmd/agent/core/agent.go +++ b/cmd/agent/core/agent.go @@ -267,6 +267,11 @@ func run(ctx context.Context, c *cli.Command, backends []types.Backend) error { err := client.ReportHealth(grpcCtx) if err != nil { log.Err(err).Msg("failed to report health") + // Check if the error is due to context cancellation + if grpcCtx.Err() != nil || agentCtx.Err() != nil { + log.Debug().Msg("terminating health reporting due to context cancellation") + return nil + } } select { @@ -291,8 +296,18 @@ func run(ctx context.Context, c *cli.Command, backends []types.Backend) error { log.Debug().Msg("polling new steps") if err := runner.Run(agentCtx, shutdownCtx); err != nil { - log.Error().Err(err).Msg("runner done with error") - return err + log.Error().Err(err).Msg("runner error, retrying...") + // Check if context is canceled + if agentCtx.Err() != nil { + return nil + } + // Wait a bit before retrying to avoid hammering the server + select { + case <-agentCtx.Done(): + return nil + case <-time.After(time.Second * 5): + // Continue to next iteration + } } } })