diff --git a/agent_server.go b/agent_server.go
index 2f38092..fb583e0 100644
--- a/agent_server.go
+++ b/agent_server.go
@@ -209,6 +209,7 @@ func (s *AgentServer) Run(ctx context.Context) error {
 			return ctx.Err()
 		}
 	case <-ctx.Done():
+		server.Stop()
 		return ctx.Err()
 	}
 }